#!/usr/bin/env python3
"""Train Karman Cloak (V5) — parameterized, no-bias only, GPU-native.

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)

Usage:
    conda run -n pycuda_3_10 python -u train_karman.py \
        --case-name re100_karman --device-id 0 --seed 42 \
        --config configs/config_lbm_karman_2000x600.json \
        --calibration calibrations/re100/calibration.json \
        --total-episodes 500
"""
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 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 V5")
    parser.add_argument("--case-name", type=str, required=True,
                        help="Case name for output dir (e.g. re100_karman)")
    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)
    parser.add_argument("--n-steps", type=int, default=2048)
    parser.add_argument("--learn-timesteps", type=int, default=2048)
    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)
    parser.add_argument("--config", type=str, required=True,
                        help="Path to LBM config JSON")
    parser.add_argument("--calibration", type=str, required=True,
                        help="Path to calibration.json")
    parser.add_argument("--si", type=int, default=None,
                        help="Override sample interval (default: from calibration)")
    parser.add_argument("--symmetry-prob", type=float, default=0.0,
                        help="G-symmetry augmentation probability (0=off, 0.5=half)")
    parser.add_argument("--transfer-model", type=str, default=None,
                        help="Path to .zip model for transfer learning")
    parser.add_argument("--resume-from", type=int, default=0,
                        help="Resume from episode N (0=start fresh). Loads ep{N}_model.zip "
                             "and ep{N}_vecnormalize.pkl, parses train.log for best_reward, "
                             "continues from ep N+1.")
    args = parser.parse_args()

    # Load calibration
    with open(args.calibration, "r") as f:
        cal = json.load(f)

    # Override SI if specified
    if args.si is not None:
        cal["SI"] = int(args.si)

    config_path = args.config
    si_val = cal["SI"]

    run_name = f"{args.case_name}_seed{args.seed}"
    out_dir = Path(__file__).resolve().parent / "output" / run_name
    (out_dir / "models").mkdir(parents=True, exist_ok=True)

    # ── Resume mode: load checkpoint and restore best_reward ──────────────
    resume_ep = args.resume_from
    resume_best_reward = -float("inf")

    if resume_ep > 0:
        log_path_tmp = out_dir / "train.log"
        def _log(msg):
            line = f"[{time.strftime('%H:%M:%S')}] {msg}"
            print(line, flush=True)
            with open(log_path_tmp, "a") as f: f.write(line + "\n"); f.flush()
        _log(f"=== V5 {run_name} (RESUME from ep {resume_ep}) ===")
        if log_path_tmp.exists():
            with open(log_path_tmp, "r") as f:
                for line in f:
                    if "BEST" in line:
                        import re
                        m = re.search(r'reward=([\d.]+)\s+\(BEST', line)
                        if m:
                            resume_best_reward = float(m.group(1))
            _log(f"  Recovered best_reward={resume_best_reward:.4f} from log")
        else:
            _log("  WARNING: train.log not found, best_reward set to -inf")
        ckpt_model = out_dir / "models" / f"ep{resume_ep:04d}_model.zip"
        ckpt_norm  = out_dir / "models" / f"ep{resume_ep:04d}_vecnormalize.pkl"
        if not ckpt_model.exists():
            _log(f"  ERROR: Checkpoint not found: {ckpt_model}")
            return 1
        if not ckpt_norm.exists():
            _log(f"  ERROR: Checkpoint not found: {ckpt_norm}")
            return 1
        _log(f"  Resume checkpoint: {ckpt_model}")

    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"=== V5 {run_name} ===")
    log(f"  Device={args.device_id}, Seed={args.seed}, Episodes={args.total_episodes}")
    log(f"  Config={config_path}, SI={si_val}")
    log(f"  FORCE_SCALE={cal['FORCE_SCALE']:.6f}, SENS_SCALE={cal['SENS_SCALE']:.4f}")
    log(f"  dtw_norm_scale={cal['dtw_norm_scale']:.4f}")

    # Save calibration and config alongside training output
    cal_copy = dict(cal)
    cal_copy["config_path"] = config_path
    if args.transfer_model:
        cal_copy["transfer_model"] = args.transfer_model
    if not resume_ep:
        with open(out_dir / "calibration.json", "w") as f:
            json.dump(cal_copy, f, indent=2)

    # Load target
    cal_dir = Path(args.calibration).resolve().parent
    target_path = cal_dir / "target.npy"
    if target_path.exists():
        target_states = np.load(str(target_path))
        log(f"  Loaded target from {target_path}")
    else:
        log("  Recording target...")
        t0 = time.perf_counter()
        dist_radius = float(cal.get("dist_radius", 1.0))
        target_states = record_target(config_path, args.device_id, si_val, dist_radius=dist_radius)
        np.save(str(out_dir / "target.npy"), target_states)
        log(f"  Target recorded in {time.perf_counter()-t0:.0f}s")

    log("  Creating env...")
    t0 = time.perf_counter()
    env = KarmanCloakEnv(device_id=args.device_id, seed=args.seed,
                          calibration=cal, config_path=config_path,
                          target_states=target_states)
    env = SymmetryAugmentWrapper(env, prob=args.symmetry_prob, seed=args.seed,
                                  rollout_len=args.n_steps)
    vec_env = DummyVecEnv([lambda: env])
    if resume_ep and ckpt_norm.exists():
        # Resume: restore single-layer VecNormalize from checkpoint
        vec_env = VecNormalize.load(str(ckpt_norm), vec_env)
        vec_env.training = True
        vec_env.norm_reward = False
        log("  Loaded VecNormalize from checkpoint.")
    else:
        # Fresh training or transfer: create new VecNormalize
        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")

    device = torch.device(f"cuda:{args.device_id}")
    log(f"  PPO device: {device}")

    if args.transfer_model:
        log(f"  Transfer learning from: {args.transfer_model}")
        model = PPO.load(args.transfer_model, env=vec_env, device=device,
                          custom_objects={"activation_fn": Sin})
        log("  Loaded base model.")
    elif resume_ep:
        log(f"  Resuming from: {ckpt_model}")
        model = PPO.load(str(ckpt_model), env=vec_env, device=device,
                          custom_objects={"activation_fn": Sin})
        log("  Loaded checkpoint model.")
    else:
        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,
            gamma=0.995,
            verbose=0,
        )
        log("  Created from scratch.")

    best_reward = resume_best_reward if resume_ep else -float("inf")
    best_ep = None
    t_last = time.perf_counter()
    norm_path = str(out_dir / "vec_normalize.pkl")

    # Save initial vec_normalize state (important after checkpoint load)
    if resume_ep:
        vec_env.save(norm_path)

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

        # Evaluation (disable symmetry for clean policy eval)
        eval_env = model.get_env()
        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, deterministic=True)
            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
        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
            best_ep = ep
            model.save(str(out_dir / "models" / "best_model.zip"))
            vec_env.save(str(out_dir / "best_vecnormalize.pkl"))
            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)")

        # Save model + eval summary after EVERY episode
        model.save(str(out_dir / "models" / f"ep{ep:04d}_model.zip"))
        vec_env.save(str(out_dir / "models" / f"ep{ep:04d}_vecnormalize.pkl"))

    model.save(str(out_dir / "models" / "final_model.zip"))
    vec_env.save(str(out_dir / "final_vecnormalize.pkl"))
    # Keep the historical root name as a final-state alias for old consumers.
    vec_env.save(norm_path)

    meta = {"case_name": args.case_name, "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,
            "config_path": config_path, "calibration_path": args.calibration,
            "transfer_model": args.transfer_model,
            "best_episode": best_ep}
    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())
