#!/usr/bin/env python3
"""Train Karman Cloak (2000x600) — GPU-native, no device="cpu" hack.

CUDA context management:
  - PPO model on GPU (PyTorch owns the default context)
  - CFD env wraps its calls in sim.ctx._ctx.push()/pop() (PyCUDA context)
  - This mirrors the exact pattern in legacy env: push before CFD, pop after.

Usage:
    conda run -n pycuda_3_10 python -u train_karman_2000x600.py --device-id 0 --seed 42 --total-episodes 50

    # Pre-record target to speed up:
    python -c "from env_karman_2000x600 import record_target; import numpy as np; np.save('target.npy', record_target(0))"
    conda run -n pycuda_3_10 python -u train_karman_2000x600.py --device-id 0 --seed 42 --target target.npy
"""
from __future__ import annotations

import argparse, json, os, sys, time
from pathlib import Path

import numpy as np
import pycuda.driver as cuda; cuda.init()

_REPO = Path(__file__).resolve().parents[3]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

import torch
from torch.nn import Module as TorchModule
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
from torch.utils.tensorboard import SummaryWriter

from env_karman_2000x600 import KarmanCloakEnv, record_target
from symmetry_wrapper import SymmetryAugmentWrapper


class Sin(TorchModule):
    def __init__(self): super().__init__()
    def forward(self, x): return torch.sin(x)


def main() -> int:
    parser = argparse.ArgumentParser(description="Train Karman Cloak 2000x600")
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--total-episodes", type=int, default=500,
                        help="Total episodes (default 500)")
    parser.add_argument("--n-steps", type=int, default=2048,
                        help="PPO rollout buffer size (default 2048 = legacy)")
    parser.add_argument("--learn-timesteps", type=int, default=2048,
                        help="Timesteps per model.learn() call (should = n_steps)")
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--n-epochs", type=int, default=10,
                        help="PPO n_epochs (default 10 = legacy)")
    parser.add_argument("--target", type=str, default=None)
    parser.add_argument("--symmetry-prob", type=float, default=0.5,
                        help="Probability of G-symmetry augmentation (0=off, 0.5=half)")
    parser.add_argument("--no-bias", action="store_true",
                        help="Train without action bias (ACTION_BIAS=[0,0,0], scale=12)")
    args = parser.parse_args()

    bias_tag = "nobias" if args.no_bias else "bias"
    run_name = f"{bias_tag}_seed{args.seed}_s{args.n_steps}_e{args.n_epochs}_v4"
    out_dir = Path(__file__).resolve().parent / "output" / run_name
    (out_dir / "models").mkdir(parents=True, exist_ok=True)

    log_path = out_dir / "train.log"
    def log(msg):
        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
        print(line, flush=True)
        with open(log_path, "a") as f: f.write(line + "\n"); f.flush()

    writer = SummaryWriter(log_dir=str(out_dir / "tb"))
    log(f"=== {run_name} ===")
    log(f"  Device={args.device_id}, Seed={args.seed}, lr={args.lr}, episodes={args.total_episodes}")

    if args.target:
        target_states = np.load(args.target)
        log(f"  Loaded target from {args.target}")
    else:
        log("  Recording target...")
        t0 = time.perf_counter()
        target_states = record_target(args.device_id)
        log(f"  Target recorded in {time.perf_counter()-t0:.0f}s")

    log("  Creating env...")
    t0 = time.perf_counter()
    if args.no_bias:
        action_bias = np.array([0.0, 0.0, 0.0], dtype=np.float32)
        action_scale = 12.0  # wider range to cover Bias omega range [-12,12]×U0
        log(f"  NoBias mode: scale={action_scale}, bias={action_bias}")
    else:
        action_bias = None  # use default [0,-4,4]
        action_scale = None  # use default 8.0
    env = KarmanCloakEnv(device_id=args.device_id, seed=args.seed,
                          target_states=target_states,
                          action_bias=action_bias,
                          action_scale=action_scale)
    # Symmetry augmentation: 50% prob G-mirror obs/action to encourage
    # learning a symmetric (up/down) control law.
    env = SymmetryAugmentWrapper(env, prob=args.symmetry_prob, seed=args.seed,
                                  rollout_len=args.n_steps)
    vec_env = DummyVecEnv([lambda: env])
    # VecNormalize: automatic running mean/std normalization for obs.
    # norm_reward=False because reward is already in [0,1] and we want
    # to preserve its magnitude for eval. clip_obs=10 is generous (physical
    # norm already keeps obs in ~[-1,1] range).
    vec_env = VecNormalize(vec_env, norm_obs=True, norm_reward=False,
                           clip_obs=10.0, gamma=0.99)
    log(f"  Env ready in {time.perf_counter()-t0:.0f}s (symmetry_prob={args.symmetry_prob}, VecNormalize)")

    # PPO on GPU — env handles CUDA context via push/pop around CFD calls
    device = torch.device(f"cuda:{args.device_id}")
    log(f"  PPO device: {device}")

    model = PPO(
        "MlpPolicy",
        policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]},
        env=vec_env,
        device=device,
        n_steps=args.n_steps,
        batch_size=args.batch_size,
        n_epochs=args.n_epochs,
        learning_rate=args.lr,
        # Legacy-matching: use SB3 defaults (no ent_coef, no target_kl)
    # ent_coef=0.0 (default), target_kl=None (default)
        gamma=0.995,   # keep higher gamma for normalized DTW delay propagation
        verbose=0,
    )
    log("  Created from scratch.")

    best_reward = -float("inf")
    t_last = time.perf_counter()

    # Save VecNormalize statistics periodically for inference reuse
    norm_path = str(out_dir / "vec_normalize.pkl")

    for ep in range(1, args.total_episodes + 1):
        model.learn(total_timesteps=args.learn_timesteps, reset_num_timesteps=False)

        # --- Evaluation: disable symmetry augmentation for clean policy eval ---
        eval_env = model.get_env()
        # Access underlying env to turn off symmetry
        try:
            inner = eval_env.venv.envs[0]
            if hasattr(inner, 'prob'):
                inner.prob = 0.0
        except Exception:
            pass
        eval_obs = eval_env.reset()
        ep_rewards = []
        ep_r_cd, ep_r_cl, ep_r_sim = [], [], []
        for _ in range(360):
            action, _ = model.predict(eval_obs)  # stochastic eval (matches legacy, smoother curves)
            eval_obs, reward, done, info = eval_env.step(action)
            ep_rewards.append(float(reward[0]))
            inf = info[0] if isinstance(info, list) else info
            if "r_cd" in inf:
                ep_r_cd.append(float(inf["r_cd"]))
                ep_r_cl.append(float(inf["r_cl"]))
                ep_r_sim.append(float(inf["r_sim"]))
            if done[0]: break
        # Re-enable symmetry for next training round
        try:
            inner = eval_env.venv.envs[0]
            if hasattr(inner, 'prob'):
                inner.prob = args.symmetry_prob
        except Exception:
            pass
        avg_r = np.mean(ep_rewards[-180:]) if len(ep_rewards) >= 180 else np.mean(ep_rewards)
        dt = time.perf_counter() - t_last; t_last = time.perf_counter()

        writer.add_scalar("eval/avg_reward", avg_r, ep)
        if ep_r_cd:
            writer.add_scalar("eval/r_cd", float(np.mean(ep_r_cd[-180:])), ep)
            writer.add_scalar("eval/r_cl", float(np.mean(ep_r_cl[-180:])), ep)
            writer.add_scalar("eval/r_sim", float(np.mean(ep_r_sim[-180:])), ep)

        if avg_r > best_reward:
            best_reward = avg_r
            model.save(str(out_dir / "models" / "best_model.zip"))
            # Save normalizer alongside best model
            vec_env.save(norm_path)
            log(f"  Ep {ep:3d}: reward={avg_r:.4f} (BEST, {dt:.0f}s/ep) ** "
                f"r_cd={np.mean(ep_r_cd[-180:]):.3f} r_cl={np.mean(ep_r_cl[-180:]):.3f} "
                f"r_sim={np.mean(ep_r_sim[-180:]):.3f}")
        elif ep % 5 == 0:
            log(f"  Ep {ep:3d}: reward={avg_r:.4f} (best={best_reward:.4f}, {dt:.0f}s/ep)")

        if ep % 10 == 0:
            model.save(str(out_dir / "models" / f"chkpt_ep{ep}.zip"))
            vec_env.save(norm_path)

    model.save(str(out_dir / "models" / "final_model.zip"))
    vec_env.save(norm_path)

    meta = {"seed": args.seed, "total_episodes": args.total_episodes,
            "best_reward": float(best_reward), "n_steps": args.n_steps,
            "batch_size": args.batch_size, "n_epochs": args.n_epochs, "lr": args.lr}
    with (out_dir / "meta.json").open("w") as f:
        json.dump(meta, f, indent=2)

    env.close()
    log(f"Done. Best reward: {best_reward:.4f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
