"""
Exp 11: Parallel DummyVecEnv — generated_track + mountain_track on two sim instances.

THIS IS THE KEY EXPERIMENT: Does parallel multi-track training via DummyVecEnv
produce reliable results, unlike the close-and-switch approach (Wave 4, Exp 10)?

Setup:
  - Sim 1: 10.0.0.55:9091 → generated_track
  - Sim 2: 10.0.0.55:9093 → mountain_track
  - DummyVecEnv wraps both → PPO sees both tracks in every rollout batch
  - NO env closing, NO set_env(), NO track switching
  - Same hyperparameters as Wave 4 Trial 9 (the lottery winner)

Hypothesis: Parallel envs eliminate catastrophic forgetting. Results should be
CONSISTENT across runs, not a 16% lottery like Wave 4.
"""
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  = 90000
SAVE_DIR     = '/home/paulh/projects/donkeycar-rl-autoresearch/agent/models/exp11-parallel-envs'
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)
        return env
    return _init

log('='*60)
log('Exp 11: Parallel DummyVecEnv — two sims, two tracks')
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'  Method: DummyVecEnv (both tracks in every PPO batch)')
log(f'  NO close_and_switch — single stable env for entire training')
log('='*60)

# Create parallel env
log('Creating DummyVecEnv with two tracks...')
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}')

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

# Train in segments for checkpointing
CHECKPOINT_EVERY = 6000
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

    # Save checkpoint
    ckpt = os.path.join(SAVE_DIR, f'checkpoint_{steps_done:07d}')
    model.save(ckpt)
    log(f'[{steps_done:,}/{TOTAL_STEPS:,}] Checkpoint saved: {ckpt}.zip')

    # Quick eval on the parallel env (both tracks)
    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
        log(f'  Eval: env0(gen_track)={ep_rewards[0]:.1f}r/{int(ep_steps[0])}steps  '
            f'env1(mountain)={ep_rewards[1]:.1f}r/{int(ep_steps[1])}steps')

        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'))
env.close()
time.sleep(3)

log(f'\nTraining complete. Best combined reward: {best_reward:.1f}')
log(f'Checkpoints in {SAVE_DIR}:')
for f in sorted(os.listdir(SAVE_DIR)):
    size = os.path.getsize(os.path.join(SAVE_DIR, f)) // (1024*1024)
    log(f'  {f}  ({size}MB)')

# Run standard eval
log('\nRunning standard 3-set eval on all tracks...')
import subprocess
subprocess.run([
    'python3',
    '/home/paulh/projects/donkeycar-rl-autoresearch/agent/run_eval.py',
    '--model', os.path.join(SAVE_DIR, 'best_model.zip'),
    '--sets', '3', '--steps', '2000'
], cwd='/home/paulh/projects/donkeycar-rl-autoresearch/agent')

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