"""
Exp 11c: Parallel DummyVecEnv, v6 reward, 250k steps (extended budget).

Exp 11b showed parallel envs + v6 reward works (no circles, stable training)
but plateaued at ~194 steps with only 90k total steps (45k per track).

This run extends to 250k steps (125k per track) to see if the model can
break through the plateau with more training budget.

Everything else identical to Exp 11b.
"""
import sys, os, time
sys.path.insert(0, '/home/paulh/projects/donkeycar-rl-autoresearch/agent')

from multitrack_runner import log, StuckTerminationWrapper
from donkeycar_sb3_runner import ThrottleClampWrapper
from reward_wrapper import SpeedRewardWrapper
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage
import gymnasium as gym
import numpy as np

HOST         = '10.0.0.55'
THROTTLE_MIN = 0.2
LR           = 0.000725
TOTAL_STEPS  = 250000
SAVE_DIR     = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp11c-parallel-v6-250k'
os.makedirs(SAVE_DIR, exist_ok=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=40, min_displacement=0.5)
        env = SpeedRewardWrapper(env)  # v6
        return env
    return _init

log('='*60)
log('Exp 11c: Parallel DummyVecEnv, v6 reward, 250k steps')
log(f'  Sim 1: {HOST}:9091 → generated_track')
log(f'  Sim 2: {HOST}:9093 → mountain_track')
log(f'  throttle_min={THROTTLE_MIN}, lr={LR}, total={TOTAL_STEPS:,}')
log(f'  Reward: v6 (speed × CTE_quality, gated by efficiency >= 0.15)')
log(f'  Budget: 250k steps = ~125k per track (2.8× more than Exp 11b)')
log('='*60)

env = DummyVecEnv([
    make_env('donkey-generated-track-v0', 9091),
    make_env('donkey-mountain-track-v0', 9093),
])
env = VecTransposeImage(env)
log(f'  VecEnv num_envs={env.num_envs}, obs={env.observation_space.shape}')

model = PPO('CnnPolicy', env, learning_rate=LR, verbose=1, device='cpu')
log('PPO created. Starting training...')

CHECKPOINT_EVERY = 10000
best_reward = float('-inf')
steps_done = 0

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'))
    log(f'[{steps_done:,}/{TOTAL_STEPS:,}] Checkpoint saved')

    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

        status0 = '✅' if ep_steps[0] >= 2000 else f'❌@{int(ep_steps[0])}'
        status1 = '✅' if ep_steps[1] >= 2000 else f'❌@{int(ep_steps[1])}'
        log(f'  Eval: gen_track={ep_rewards[0]:.1f}r/{int(ep_steps[0])}s {status0}  '
            f'mountain={ep_rewards[1]:.1f}r/{int(ep_steps[1])}s {status1}')

        total_reward = ep_rewards.sum()
        if total_reward > best_reward:
            best_reward = total_reward
            model.save(os.path.join(SAVE_DIR, 'best_model'))
            log(f'  ⭐ NEW BEST: {best_reward:.1f} (combined)')
    except Exception as e:
        log(f'  Eval error: {e}')

model.save(os.path.join(SAVE_DIR, 'model'))
log(f'\nTraining complete. Best combined reward: {best_reward:.1f}')

# Close training env
env.close()
time.sleep(5)

# --- Eval on all 4 tracks ---
log('\n' + '='*60)
log('EVALUATION: best_model on 4 tracks (3 sets each)')
log('='*60)

EVAL_TRACKS = [
    ('donkey-mountain-track-v0',   'mountain_track'),
    ('donkey-generated-track-v0',  'generated_track'),
    ('donkey-generated-roads-v0',  'generated_road'),
    ('donkey-minimonaco-track-v0', 'mini_monaco'),
]
EVAL_PORT = 9091
best_model_path = os.path.join(SAVE_DIR, 'best_model.zip')
results_by_track = {}

for track_id, track_name in EVAL_TRACKS:
    log(f'\n--- {track_name} ---')
    steps_list = []
    for s in range(1, 4):
        try:
            raw = gym.make(track_id, conf={'host': HOST, 'port': EVAL_PORT})
            eval_inner = ThrottleClampWrapper(raw, throttle_min=THROTTLE_MIN)
            eval_inner = StuckTerminationWrapper(eval_inner, stuck_steps=40, min_displacement=0.5)
            eval_inner = SpeedRewardWrapper(eval_inner)
            eval_env = VecTransposeImage(DummyVecEnv([lambda e=eval_inner: e]))
            eval_model = PPO.load(best_model_path, env=eval_env, device='cpu')
            
            obs = eval_env.reset()
            total_r, total_s, done = 0.0, 0, False
            while not done and total_s < 2000:
                action, _ = eval_model.predict(obs, deterministic=True)
                result = eval_env.step(action)
                if len(result) == 4: obs, r, d, _ = result; done = bool(d[0])
                else: obs, r, t, tr, _ = result; done = bool(t[0] or tr[0])
                total_r += float(r[0]); total_s += 1
            
            status = '✅' if total_s >= 2000 else f'❌@{total_s}'
            log(f'  Set{s}: {total_r:.1f}r / {total_s}s {status}')
            steps_list.append(total_s)
            eval_env.close(); time.sleep(3)
        except Exception as e:
            log(f'  Set{s}: ERROR — {e}')
            steps_list.append(0); time.sleep(3)
    
    results_by_track[track_name] = steps_list
    log(f'  Mean: {np.mean(steps_list):.0f} steps')

log('\n' + '='*60)
log('SUMMARY')
log('='*60)
for track_name, steps_list in results_by_track.items():
    steps_str = '/'.join(str(s) for s in steps_list)
    mean = np.mean(steps_list)
    verdict = '✅' if mean >= 1500 else '⚠️' if mean >= 500 else '❌'
    log(f'  {verdict} {track_name:20s}: {steps_str}  mean={mean:.0f}')

log(f'\n=== Exp 11c COMPLETE ===')