"""
Exp 29: Fine-tune wave4-trial-0009 on generated track.

wave4-trial-0009 is our best mini-monaco model (completed laps in eval).
It was trained on gentrack+mountain with continuous actions and LR≈0.00073.

What this experiment does:
  - Warm-start from wave4-trial-0009/model.zip (continuous Box action space)
  - Fine-tune on generated track only, very low LR to preserve driving skill
  - Same wrapper stack as wave4: ThrottleClampWrapper → StuckTermination → SpeedReward
  - NO DiscretizedActionWrapper — continuous actions throughout
  - 50K steps, checkpoint every 5K
  - Zero-shot eval on mini-monaco at the end

Goal: does additional generated-track exposure improve mini-monaco corner handling?
"""
import os
import sys
import time
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/exp29-wave4-finetune'
_PIDFILE    = os.path.join(_SAVE_DIR, 'current.pid')
_WARM_MODEL = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/wave4-trial-0009/model.zip'
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'[exp29] 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.utils import get_schedule_fn

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


HOST              = 'localhost'
PORT              = 9091
THROTTLE_MIN      = 0.2
LR                = 0.00005
TOTAL_STEPS       = 50_000
CHECKPOINT_EVERY  = 5_000
SCENE_RELOAD_WAIT = 5.0

TRAIN_TRACK = 'donkey-generated-track-v0'
EVAL_TRACK  = 'donkey-minimonaco-track-v0'

STUCK_STEPS           = 40
MIN_DISPLACEMENT      = 0.5
MAX_STUCK_SECONDS     = 12.0
MAX_EPISODE_SECONDS   = 30.0
LOW_SPEED_THRESHOLD   = 0.5
MAX_LOW_SPEED_SECONDS = 3.0
MAX_CTE               = 5.0
MAX_HIGH_CTE_SECONDS  = 1.0

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


def log(msg):
    print(f'[{datetime.now().strftime("%H:%M:%S")}] {msg}', flush=True)


def make_env(track_id, port):
    def _init():
        raw = gym.make(track_id, conf={'host': HOST, 'port': port})
        env = ThrottleClampWrapper(raw, throttle_min=THROTTLE_MIN)
        env = StuckTerminationWrapper(
            env,
            stuck_steps=STUCK_STEPS,
            min_displacement=MIN_DISPLACEMENT,
            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,
            max_high_cte_seconds=MAX_HIGH_CTE_SECONDS,
        )
        env = SpeedRewardWrapper(
            env,
            window_size=EFFICIENCY_WINDOW,
            min_efficiency=MIN_EFFICIENCY,
            max_cte=REWARD_MAX_CTE,
            min_lap_time=MIN_LAP_TIME,
            progress_patience=PROGRESS_PATIENCE,
        )
        return env
    return _init


def connect_env(track_id=TRAIN_TRACK):
    vec = DummyVecEnv([make_env(track_id, PORT)])
    return VecTransposeImage(vec)


def reconnect_env(old_env, track_id=TRAIN_TRACK):
    try:
        old_env.close()
    except Exception as e:
        log(f'  env.close() warning: {e}')
    time.sleep(SCENE_RELOAD_WAIT)
    return connect_env(track_id)


log('=' * 60)
log('Exp 29: wave4-trial-0009 fine-tune on generated track')
log(f'  Sim: {HOST}:{PORT} -> {TRAIN_TRACK}')
log(f'  Warm model: {_WARM_MODEL}')
log(f'  Action space: continuous Box (no discretization)')
log(f'  LR={LR}, total={TOTAL_STEPS:,}, checkpoint every {CHECKPOINT_EVERY:,}')
log(f'  After training: zero-shot eval on {EVAL_TRACK}')
log('=' * 60)

log('Connecting to sim...')
env = connect_env()
log(f'  obs={env.observation_space.shape}, action={env.action_space}')

log('Loading warm-start model from wave4-trial-0009...')
model = PPO.load(_WARM_MODEL, env=env, device='cpu')
# Must update lr_schedule — PPO.load restores the optimizer and schedule from
# the checkpoint. model.learning_rate = LR alone doesn't update the optimizer.
model.learning_rate = LR
model.lr_schedule = get_schedule_fn(LR)
for pg in model.policy.optimizer.param_groups:
    pg['lr'] = LR
log(f'  Warm model loaded. action={model.action_space}  LR={LR}')

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

best_total_steps  = float('-inf')
best_total_reward = float('-inf')
steps_done        = 0
run_tag           = datetime.now().strftime('%Y-%m-%d_%H%M%S') + '_wave4_finetune'
log_path          = os.path.join(_SAVE_DIR, f'run_{run_tag}.log')
best_model_path   = os.path.join(_SAVE_DIR, 'best_model.zip')

import logging
_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('exp29')
file_log.setLevel(logging.INFO)
file_log.propagate = False
file_log.addHandler(_fh)
file_log.addHandler(_sh)


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


flog('=' * 60)
flog(f'Exp 29 started — PID {os.getpid()}')
flog(f'Log: {log_path}')
flog(f'Warm start: wave4-trial-0009  |  LR={LR}')
flog(f'Track: {TRAIN_TRACK}  |  continuous actions')
flog('=' * 60)

# ── Training loop ─────────────────────────────────────────────────────────────
while steps_done < TOTAL_STEPS:
    seg_steps = min(CHECKPOINT_EVERY, TOTAL_STEPS - steps_done)
    model.learn(total_timesteps=seg_steps, reset_num_timesteps=False)
    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: {ckpt}.zip')

    flog('  Reconnecting for fresh track layout...')
    env = reconnect_env(env)
    model.set_env(env)

    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.sum())
        total_reward_eval = float(ep_rewards.sum())
        status = '✅' if ep_steps[0] >= 2000 else f'❌@{int(ep_steps[0])}'
        flog(f'  Eval: gentrack={total_reward_eval:.1f}r/{int(ep_steps[0])}s {status}')

        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('Training complete.')

# ── Zero-shot eval on mini-monaco ─────────────────────────────────────────────
flog('')
flog('=' * 60)
flog(f'ZERO-SHOT EVAL: best_model on {EVAL_TRACK}')
flog('=' * 60)

MINI_EPISODES  = 5
MINI_MAX_STEPS = 3000

time.sleep(SCENE_RELOAD_WAIT)
eval_env = connect_env(track_id=EVAL_TRACK)

try:
    eval_model = PPO.load(best_model_path, env=eval_env, device='cpu')
    rewards_mini, steps_mini = [], []

    for ep in range(1, MINI_EPISODES + 1):
        obs = eval_env.reset()
        total_r, steps, done = 0.0, 0, False
        while not done and steps < MINI_MAX_STEPS:
            action, _ = eval_model.predict(obs, deterministic=True)
            obs, r, d, info = eval_env.step(action)
            total_r += float(r[0])
            steps += 1
            done = bool(d[0])

        raw_info = info[0] if isinstance(info, (list, tuple)) else info
        hit = raw_info.get('hit', '?') if isinstance(raw_info, dict) else '?'
        status = '✅ timeout' if steps >= MINI_MAX_STEPS else f'❌ hit={hit}@{steps}'
        flog(f'  ep{ep}: {total_r:.1f}r / {steps}s  {status}')
        rewards_mini.append(total_r)
        steps_mini.append(steps)
        time.sleep(0.3)

    flog(f'  Mean: {np.mean(steps_mini):.0f} steps / {np.mean(rewards_mini):.1f} reward')
    flog(f'  {"✅ GENERALIZES" if np.mean(steps_mini) > 500 else "❌ DOES NOT GENERALIZE"}')

except Exception as e:
    flog(f'  Mini-monaco eval error: {e}')
finally:
    eval_env.close()

flog('')
flog('Exp 29 complete.')
flog(f'Log: {log_path}')