#!/usr/bin/env python3
"""Train Hydrodynamic Illusion (V5) - parameterized, no_bias, GPU-native.

Usage:
    conda run -n pycuda_3_10 python -u train_illusion.py \
        --case-name illusion_1L --device-id 0 --seed 42 \
        --config configs/config_lbm_karman_2000x600.json \
        --calibration calibrations/illusion_1L/calibration.json \
        --total-episodes 500
"""
from __future__ import annotations

import argparse, json, re, 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_illusion import IllusionCloakEnv
from normalization import POLICY_SENSOR_UNITS_LEGACY
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 Illusion V5")
    parser.add_argument("--case-name", type=str, required=True)
    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)
    parser.add_argument("--calibration", type=str, required=True)
    parser.add_argument("--si", type=int, default=None)
    parser.add_argument("--symmetry-prob", type=float, default=0.0)
    parser.add_argument("--transfer-model", type=str, default=None)
    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()

    with open(args.calibration, "r") as f:
        cal = json.load(f)
    if args.si is not None:
        cal["SI"] = int(args.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"
        log_existed = log_path_tmp.exists()
        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 Illusion {run_name} (RESUME from ep {resume_ep}) ===")
        # Resume is checkpoint-based: model and normalization state are restored,
        # while CFD state and exact RNG streams intentionally restart.
        best_ep = None
        if log_existed:
            with open(log_path_tmp, "r") as f:
                for line in f:
                    m = re.search(r'Ep\s+(\d+):\s+reward=([-+\d.eE]+)\s+\(BEST', line)
                    if m and int(m.group(1)) <= resume_ep:
                        best_ep = int(m.group(1))
                        resume_best_reward = float(m.group(2))
            _log(f"  Recovered best through ep {resume_ep}: "
                 f"episode={best_ep}, reward={resume_best_reward:.4f}")
        else:
            _log("  WARNING: train.log not found; resumed best_reward starts at -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()

    log(f"=== V5 Illusion {run_name} ===")
    log(f"  Device={args.device_id}, Seed={args.seed}, Ep={args.total_episodes}")
    log(f"  Config={args.config}, SI={cal['SI']}")

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

    # Illusion training requires both products from the same calibration.
    cal_dir = Path(args.calibration).resolve().parent
    target_npy = cal_dir / "target.npy"
    harmonics_path = cal_dir / "target_harmonics.json"
    missing = [path for path in (target_npy, harmonics_path) if not path.exists()]
    if missing:
        missing_text = "\n".join(f"  - {path}" for path in missing)
        raise FileNotFoundError(
            f"Required illusion calibration artifact(s) missing:\n{missing_text}\n"
            "Run illusion calibration first, for example:\n"
            f"  conda run -n pycuda_3_10 python -u "
            f"{Path(__file__).resolve().parent / 'calibrate.py'} "
            f"--case {args.case_name} --device-id {args.device_id} "
            f"--config {args.config} --scene illusion\n"
            "Then pass the generated calibration.json via --calibration."
        )
    target_states = np.load(str(target_npy))
    with open(harmonics_path, "r") as f:
        target_harmonics = json.load(f)
    log(f"  Loaded target from {cal_dir}")

    config_path = args.config
    t0 = time.perf_counter()
    writer = SummaryWriter(log_dir=str(out_dir / "tb"))

    log("  Creating env...")
    env = IllusionCloakEnv(device_id=args.device_id, seed=args.seed,
                            calibration=cal, config_path=config_path,
                            target_states=target_states,
                            target_harmonics=target_harmonics,
                            policy_sensor_units=(POLICY_SENSOR_UNITS_LEGACY
                                                 if resume_ep or args.transfer_model else None))
    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():
        vec_env = VecNormalize.load(str(ckpt_norm), vec_env)
        vec_env.training = True
        vec_env.norm_reward = False
        log("  Loaded VecNormalize from checkpoint.")
    else:
        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 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 = best_ep if resume_ep else 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)

        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) r_cd={np.mean(ep_r_cd[-180:]):.3f} "
                f"r_cl={np.mean(ep_r_cl[-180:]):.3f} 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)

        # 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_episode": int(best_ep) if best_ep is not None else None,
            "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": args.config, "calibration_path": args.calibration}
    with (out_dir / "meta.json").open("w") as f:
        json.dump(meta, f, indent=2)

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


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