"""
Exp 27: Fresh weights, truly random roads, variable throttle.

Changes from exp26:
  1. Fresh weights (no warm start) — exp26 peaked at 20k/300k then regressed.
  2. Random roads: regen_road TCP message with random seed each checkpoint.
     No close+reconnect (which was silently generating seed=2 road every time).
  3. Variable throttle: N_THROTTLE=3 → bins [0.2, 0.5, 1.0] after ThrottleClampWrapper.
  4. BrakeOnUpdateCallback: sends zero control before PPO gradient updates,
     preventing car from drifting into barriers during the ~5-15s CPU update pause.
  5. Tighter CTE termination: 2.0m / 0.5s (was 3.0m / 1.0s).
  6. Higher entropy: ent_coef=0.05 to prevent premature policy collapse.
  7. Smaller n_steps=1024: shorter rollout → shorter gradient update pause.
  8. set_ai_text: pushes training stats to sim overlay each checkpoint.
  9. 500k total steps — more budget for fresh weights to learn variable throttle.
"""
import os
import sys
import time
import random
from datetime import datetime

sys.path.insert(0, '/home/paulh/projects/donkeycar-rl-autoresearch/agent')

_SAVE_DIR = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp27-random-roads'
_PIDFILE  = os.path.join(_SAVE_DIR, 'current.pid')
os.makedirs(_SAVE_DIR, exist_ok=True)

if os.path.exists(_PIDFILE):
    try:
        _old = int(open(_PIDFILE).read().strip())
        if _old != os.getpid():
            import signal
            os.kill(_old, 0)
            print(f'[exp27] Another instance already running (PID {_old}). Exiting.', flush=True)
            sys.exit(1)
    except (OSError, ValueError):
        pass

import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage
from stable_baselines3.common.callbacks import BaseCallback

from discretize_action import DiscretizedActionWrapper
from donkeycar_sb3_runner import ThrottleClampWrapper
from multitrack_runner import StuckTerminationWrapper
from reward_wrapper import SpeedRewardWrapper


HOST              = 'localhost'
PORT              = 9091
TRACK_ID          = 'donkey-generated-roads-v0'
THROTTLE_MIN      = 0.2
LR                = 0.0003
ENT_COEF          = 0.05
N_STEPS           = 1024       # smaller rollout → shorter gradient-update pause
TOTAL_STEPS       = 500_000
CHECKPOINT_EVERY  = 10_000
REGEN_WAIT        = 3.0        # seconds after regen_road before reset

N_STEER    = 7
N_THROTTLE = 3   # throttle bins [0.0,0.5,1.0] → after ThrottleClampWrapper: [0.2,0.5,1.0]

MAX_STUCK_SECONDS     = 5.0
MAX_EPISODE_SECONDS   = 30.0
LOW_SPEED_THRESHOLD   = 1.0
MAX_LOW_SPEED_SECONDS = 1.5
MAX_CTE_TERMINATION   = 2.0    # tighter than exp26 (3.0m)
MAX_HIGH_CTE_SECONDS  = 0.5    # tighter than exp26 (1.0s)

EFFICIENCY_WINDOW   = 30
MIN_EFFICIENCY      = 0.15
MAX_CTE             = 8.0
MIN_LAP_TIME        = 12.0
PROGRESS_PATIENCE   = 100


import logging
_log_ts   = datetime.now().strftime('%Y-%m-%d_%H%M%S')
_log_path = os.path.join(_SAVE_DIR, f'run_{_log_ts}_random_roads.log')
_fh = logging.FileHandler(_log_path)
_fh.setFormatter(logging.Formatter('%(message)s'))
_sh = logging.StreamHandler(sys.stdout)
_sh.setFormatter(logging.Formatter('%(message)s'))
file_log = logging.getLogger('exp27')
file_log.setLevel(logging.INFO)
file_log.propagate = False
file_log.addHandler(_fh)
file_log.addHandler(_sh)


def flog(msg):
    ts = datetime.now().strftime('%H:%M:%S')
    file_log.info(f'[{ts}] {msg}')


def make_env():
    def _init():
        raw = gym.make(TRACK_ID, conf={'host': HOST, 'port': PORT})
        env = ThrottleClampWrapper(raw, throttle_min=THROTTLE_MIN)
        env = DiscretizedActionWrapper(env, n_steer=N_STEER, n_throttle=N_THROTTLE)
        env = StuckTerminationWrapper(
            env,
            stuck_steps=40,
            min_displacement=0.5,
            max_stuck_seconds=MAX_STUCK_SECONDS,
            max_episode_seconds=MAX_EPISODE_SECONDS,
            low_speed_threshold=LOW_SPEED_THRESHOLD,
            max_low_speed_seconds=MAX_LOW_SPEED_SECONDS,
            max_cte=MAX_CTE_TERMINATION,
            max_high_cte_seconds=MAX_HIGH_CTE_SECONDS,
        )
        env = SpeedRewardWrapper(
            env,
            window_size=EFFICIENCY_WINDOW,
            min_efficiency=MIN_EFFICIENCY,
            max_cte=MAX_CTE,
            min_lap_time=MIN_LAP_TIME,
            progress_patience=PROGRESS_PATIENCE,
        )
        return env
    return _init


def get_handler(vec_env):
    return vec_env.venv.envs[0].unwrapped.viewer.handler


def regen_road(vec_env, seed):
    msg = {
        'msg_type': 'regen_road',
        'road_style': '0',
        'rand_seed': str(seed),
        'turn_increment': '0.0',
    }
    get_handler(vec_env).queue_message(msg)
    time.sleep(REGEN_WAIT)


def set_ai_text(vec_env, text):
    try:
        get_handler(vec_env).queue_message({'msg_type': 'set_ai_text', 'text': text})
    except Exception:
        pass


class BrakeOnUpdateCallback(BaseCallback):
    """
    Sends zero-throttle control to sim before PPO gradient updates begin.

    on_rollout_end() fires after n_steps rollouts are collected, right before
    PPO starts gradient updates (which can take 5-15s on CPU). Without this,
    the sim holds the last action → car drifts into barriers during the pause.
    """
    def __init__(self, vec_env):
        super().__init__(verbose=0)
        self._vec_env = vec_env

    def _on_rollout_end(self):
        try:
            get_handler(self._vec_env).queue_message({
                'msg_type': 'control',
                'steering': '0.0',
                'throttle': '0.0',
                'brake': '0.0',
            })
        except Exception:
            pass

    def _on_step(self):
        return True


flog('=' * 60)
flog('Exp 27: fresh weights | truly random roads | variable throttle')
flog(f'  Sim: {HOST}:{PORT} → {TRACK_ID}')
flog(f'  Steering: {N_STEER} bins | Throttle: {N_THROTTLE} bins → [0.2, 0.5, 1.0]')
flog(f'  LR={LR}, ent_coef={ENT_COEF}, n_steps={N_STEPS}')
flog(f'  Total={TOTAL_STEPS:,} steps, checkpoint every {CHECKPOINT_EVERY:,}')
flog(f'  CTE term: >{MAX_CTE_TERMINATION}m for >{MAX_HIGH_CTE_SECONDS}s')
flog(f'  Speed term: <{LOW_SPEED_THRESHOLD} for >{MAX_LOW_SPEED_SECONDS}s')
flog(f'  Episode cap: {MAX_EPISODE_SECONDS}s | Road regen: random seed each checkpoint')
flog(f'  BrakeOnUpdateCallback: enabled')
flog('=' * 60)

flog('Connecting to sim...')
env = DummyVecEnv([make_env()])
env = VecTransposeImage(env)
flog(f'  Connected. obs={env.observation_space.shape}, action={env.action_space}')

first_seed = random.randint(0, 100000)
flog(f'  Initial road regen (seed={first_seed})...')
regen_road(env, first_seed)
flog('  Road ready.')

flog('Creating fresh PPO model (no warm start)...')
model = PPO(
    'CnnPolicy',
    env,
    learning_rate=LR,
    n_steps=N_STEPS,
    ent_coef=ENT_COEF,
    device='cpu',
    verbose=1,
)
flog(f'  Model created. Action space: {env.action_space.n} discrete actions')

with open(_PIDFILE, 'w') as f:
    f.write(str(os.getpid()))

flog(f'Exp 27 started — PID {os.getpid()}')
flog(f'Log: {_log_path}')

best_total_steps  = float('-inf')
best_total_reward = float('-inf')
steps_done        = 0
best_model_path   = os.path.join(_SAVE_DIR, 'best_model.zip')
brake_cb          = BrakeOnUpdateCallback(env)
current_seed      = first_seed

while steps_done < TOTAL_STEPS:
    seg_steps = min(CHECKPOINT_EVERY, TOTAL_STEPS - steps_done)
    model.learn(total_timesteps=seg_steps, reset_num_timesteps=False, callback=brake_cb)
    steps_done += seg_steps

    ckpt = os.path.join(_SAVE_DIR, f'checkpoint_{steps_done:07d}')
    model.save(ckpt)
    model.save(os.path.join(_SAVE_DIR, 'model'))
    flog(f'[{steps_done:,}/{TOTAL_STEPS:,}] Checkpoint saved')

    current_seed = random.randint(0, 100000)
    flog(f'  Regenerating road (seed={current_seed})...')
    regen_road(env, current_seed)
    flog('  Road ready.')

    try:
        obs = env.reset()
        ep_rewards = np.zeros(env.num_envs)
        ep_steps   = np.zeros(env.num_envs)
        done_mask  = np.zeros(env.num_envs, dtype=bool)

        for _ in range(2000):
            action, _ = model.predict(obs, deterministic=True)
            obs, rewards, dones, infos = env.step(action)
            for i in range(env.num_envs):
                if not done_mask[i]:
                    ep_rewards[i] += rewards[i]
                    ep_steps[i]   += 1
                    if dones[i]:
                        done_mask[i] = True
            if done_mask.all():
                break

        total_steps_eval  = int(ep_steps[0])
        total_reward_eval = float(ep_rewards[0])
        status = '✅' if total_steps_eval >= 2000 else f'❌@{total_steps_eval}'
        flog(f'  Eval (seed={current_seed}): {total_reward_eval:.1f}r/{total_steps_eval}s {status}')

        overlay = (f'Exp27 {steps_done//1000:3d}k/{TOTAL_STEPS//1000}k\n'
                   f'R:{total_reward_eval:6.1f} Seed:{current_seed} {status}')
        set_ai_text(env, overlay)

        if (total_steps_eval > best_total_steps
                or (total_steps_eval == best_total_steps
                    and total_reward_eval > best_total_reward)):
            best_total_steps  = total_steps_eval
            best_total_reward = total_reward_eval
            model.save(best_model_path)
            flog(f'  NEW BEST: steps={best_total_steps} reward={best_total_reward:.1f}')

    except Exception as e:
        flog(f'  Eval error: {e}')

env.close()

flog('=' * 60)
flog('Exp 27 complete.')
flog(f'Best model: {best_model_path}')
flog(f'Best eval: steps={best_total_steps} reward={best_total_reward:.1f}')
flog('=' * 60)