#!/usr/bin/env python3
"""V5 Train model inference engine — skeleton injection with per-seed VecNormalize.

For each scene in TRAIN_SCENES:
  1. Create V5 CFD env (KarmanCloakEnv / IllusionCloakEnv)
  2. For every seed: fresh skeleton + VecNormalize.load(seed.pkl) + weight injection
  3. 360-step deterministic rollout, pick best by tail-180 reward
  4. Save signals.npz + vorticity PNGs (controlled, target, zero)

Uses skeleton injection per-seed (one skeleton per seed, one VecNormalize load per seed)
to avoid cloudpickle/numpy PPO.load() deserialization issues while still using
trained VecNormalize stats for proper observation normalization.

Usage:
    conda run -n pycuda_3_10 python -u infer_train.py --device-id 0
    conda run -n pycuda_3_10 python -u infer_train.py --device-id 0 --scene re100
"""
from __future__ import annotations

import argparse
import io
import json
import os
import sys
import time
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Tuple

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

_REPO = str(Path(__file__).resolve().parents[3])
_SRC = Path(_REPO) / "src"
for p in [_REPO, str(_SRC)]:
    if p not in sys.path:
        sys.path.insert(0, p)

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

from CelerisLab.common.render import compute_vorticity, render_vorticity_field

_TRAIN_DIR = _SRC / "drl_pinball" / "train"
sys.path.insert(0, str(_TRAIN_DIR))

from symmetry_wrapper import SymmetryAugmentWrapper

from drl_pinball.eval.scene_manifest import TRAIN_SCENES

_OUT_BASE = Path(__file__).resolve().parent / "output" / "train"


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


_device = None  # set at startup


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


# ---------------------------------------------------------------------------
# Env creation
# ---------------------------------------------------------------------------
def _create_env(scene: Dict[str, Any], device_id: int):
    scene_type = scene["scene_type"]
    config_path = Path(scene["config_path"]).resolve()
    cal_path = Path(scene["calibration_path"]).resolve()
    si = scene["si"]

    with open(cal_path) as f:
        cal = json.load(f)
    cal["SI"] = si

    cal_dir = cal_path.parent
    target_path = cal_dir / "target.npy"
    if not target_path.exists():
        raise FileNotFoundError(f"target.npy not found at {target_path}")
    target_states = np.load(str(target_path))

    if scene_type == "karman":
        from env_karman import KarmanCloakEnv
        log(f"  Creating KarmanCloakEnv (config={config_path.name}, SI={si})...")
        t0 = time.perf_counter()
        env = KarmanCloakEnv(
            device_id=device_id, seed=41,
            calibration=cal, config_path=str(config_path),
            target_states=target_states,
        )
        log(f"  Env ready in {time.perf_counter() - t0:.0f}s")
        return env
    elif scene_type == "illusion":
        from env_illusion import IllusionCloakEnv
        target_harmonics_path = cal_dir / "target_harmonics.json"
        if not target_harmonics_path.exists():
            raise FileNotFoundError(f"target_harmonics.json not found at {target_harmonics_path}")
        with open(target_harmonics_path) as f:
            target_harmonics = json.load(f)
        target_diam = float(scene.get("target_diam", 1.0))
        cal["target_diam"] = target_diam
        log(f"  Creating IllusionCloakEnv (config={config_path.name}, diam={target_diam}L, SI={si})...")
        t0 = time.perf_counter()
        env = IllusionCloakEnv(
            device_id=device_id, seed=41,
            calibration=cal, config_path=str(config_path),
            target_states=target_states,
            target_harmonics=target_harmonics,
            target_diam=target_diam,
        )
        log(f"  Env ready in {time.perf_counter() - t0:.0f}s")
        return env
    else:
        raise ValueError(f"Unknown scene_type: {scene_type}")


# ---------------------------------------------------------------------------
# Extract policy weights from SB3 .zip
# ---------------------------------------------------------------------------
def _extract_state_dict(zip_path: str) -> dict:
    with zipfile.ZipFile(zip_path, "r") as zf:
        with zf.open("policy.pth") as f:
            return torch.load(io.BytesIO(f.read()), map_location="cpu",
                              weights_only=False)


# ---------------------------------------------------------------------------
# Raw env extraction (walk wrapper chain)
# ---------------------------------------------------------------------------
def _get_raw_env(vec_env):
    dummy = vec_env.venv
    inner = dummy.envs[0]
    if hasattr(inner, 'prob'):
        return inner.env
    return inner


# ---------------------------------------------------------------------------
# Saving
# ---------------------------------------------------------------------------
def _save_vorticity(env, out_path: Path, scene: Dict[str, Any]) -> None:
    sim = env.sim
    macro = sim.get_macroscopic()
    vort = compute_vorticity(macro["ux"], macro["uy"])
    nx = int(sim.lbm_cfg.nx)
    ny = int(sim.lbm_cfg.ny)
    cylinders = []
    for body_id in range(sim.bodies.count):
        body = sim.bodies.get(body_id)
        if hasattr(body, 'center'):
            c = body.center
            r = body.radius if hasattr(body, 'radius') else 10.0
            cylinders.append(((c[0], c[1]), r))
    render_vorticity_field(vort, nx=nx, ny=ny, out_path=str(out_path),
                           cylinders=cylinders, vmin=-0.03, vmax=0.03)


# ---------------------------------------------------------------------------
# Per-seed eval — correct pattern: PPO.load(no env) + VecNormalize(env, training=False)
# ---------------------------------------------------------------------------
def _eval_one_seed(env, seed_label: str, seed_dir_str: str,
                   num_steps: int) -> Dict[str, Any]:
    """Load model WITHOUT env, wrap env with VecNormalize.load (single layer, frozen)."""
    seed_dir = Path(seed_dir_str)
    model_path = seed_dir / "best_model.zip"
    norm_path = seed_dir.parent / "vec_normalize.pkl"

    if not model_path.exists():
        log(f"    [seed {seed_label}] SKIP: no best_model.zip")
        return {"seed": seed_label, "avg_reward": -float("inf"), "skip": True}
    if not norm_path.exists():
        log(f"    [seed {seed_label}] SKIP: no vec_normalize.pkl")
        return {"seed": seed_label, "avg_reward": -float("inf"), "skip": True}

    log(f"    [seed {seed_label}] loading...")

    # All seeds: PPO.load fails (numpy._core.numeric cloudpickle issue).
    # Use skeleton injection directly — avoids GPU pollution from failed load.
    wrapped = SymmetryAugmentWrapper(env, prob=0.0, seed=41, rollout_len=2048)
    vec_env = DummyVecEnv([lambda: wrapped])
    vec_env = VecNormalize.load(str(norm_path), vec_env)
    vec_env.norm_reward = False

    skeleton = PPO(
        "MlpPolicy",
        policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]},
        env=vec_env, device=_device,
        n_steps=2048, batch_size=64, n_epochs=10,
        learning_rate=3e-4, gamma=0.995, verbose=0,
    )
    sd = _extract_state_dict(str(model_path))
    skeleton.policy.load_state_dict(sd, strict=False)

    raw_env = _get_raw_env(vec_env)
    sens_cc = float(raw_env._cal.get("SENSOR_CC", 78.0))
    return _rollout(skeleton, vec_env, raw_env, sens_cc, seed_label, num_steps)


def _rollout(model, vec_env, raw_env, sens_cc, seed_label, num_steps):
    """Deterministic rollout, model.predict receives normalized obs from vec_env."""
    ep_rew, ep_r_cd, ep_r_cl, ep_r_sim, ep_sim_raw = [], [], [], [], []
    sig_s, sig_f, sig_a = [], [], []
    obs = vec_env.reset()

    for step in range(num_steps):
        action, _ = model.predict(obs, deterministic=True)
        obs, reward, done, info = vec_env.step(action)

        inf = info[0] if isinstance(info, list) else info
        ep_rew.append(float(reward[0]))
        ep_r_cd.append(float(inf.get("r_cd", 0)))
        ep_r_cl.append(float(inf.get("r_cl", 0)))
        ep_r_sim.append(float(inf.get("r_sim", 0)))
        ep_sim_raw.append(float(inf.get("sim", 0)))

        try:
            raw = raw_env._read_obs()
            if "Karman" in type(raw_env).__name__:
                sl = raw[2:14]
            else:
                sl = raw[:]
            sig_s.append(sl[0:6] * sens_cc)
            sig_f.append(sl[6:12])
        except Exception:
            sig_s.append(np.zeros(6, dtype=np.float32))
            sig_f.append(np.zeros(6, dtype=np.float32))
        sig_a.append(np.asarray(action, dtype=np.float32).flatten())
        if done[0]:
            break

    tail = 180
    sl = slice(-tail, None) if len(ep_rew) >= tail else slice(None)
    avg_r = float(np.mean(ep_rew[sl]))

    log(f"      reward={avg_r:.4f}  r_cd={float(np.mean(ep_r_cd[sl])):.3f}  "
        f"r_cl={float(np.mean(ep_r_cl[sl])):.3f}  "
        f"r_sim={float(np.mean(ep_r_sim[sl])):.3f}  "
        f"sim_raw={float(np.mean(ep_sim_raw[sl])):.3f}")

    return {
        "seed": seed_label,
        "avg_reward": avg_r,
        "r_cd": float(np.mean(ep_r_cd[sl])), "r_cl": float(np.mean(ep_r_cl[sl])),
        "r_sim": float(np.mean(ep_r_sim[sl])),
        "sim_raw": float(np.mean(ep_sim_raw[sl])),
        "sensors": np.array(sig_s, dtype=np.float32),
        "forces": np.array(sig_f, dtype=np.float32),
        "actions": np.array(sig_a, dtype=np.float32),
        "rewards": np.array(ep_rew, dtype=np.float32),
        "skip": False,
    }


# ---------------------------------------------------------------------------
# Per-scene pipeline
# ---------------------------------------------------------------------------
def evaluate_train_scene(scene: Dict[str, Any], device_id: int) -> None:
    scene_id = scene["scene_id"]
    num_steps = scene["num_steps"]
    out_dir = _OUT_BASE / scene_id
    out_dir.mkdir(parents=True, exist_ok=True)

    log(f"\n{'='*70}")
    log(f"  Train eval: {scene_id}  (SI={scene['si']}, steps={num_steps})")
    log(f"{'='*70}")

    # Create CFD env once, reused across seeds
    env = _create_env(scene, device_id)

    # Evaluate each seed (independent skeleton + VecNormalize per seed)
    all_results = []
    for seed_label, seed_dir in scene["seeds"]:
        t0 = time.perf_counter()
        result = _eval_one_seed(env, seed_label, seed_dir, num_steps)
        result["_dt"] = time.perf_counter() - t0
        all_results.append(result)

    # Pick best seed
    valid = [r for r in all_results if not r.get("skip", False)]
    if not valid:
        log("  WARNING: No valid seeds found, skipping scene.")
        env.close()
        return
    best = max(valid, key=lambda r: r["avg_reward"])
    log(f"  Best: seed={best['seed']}, reward={best['avg_reward']:.4f}")

    # Save all-seeds summary
    summary = [{"seed": r["seed"], "reward": r["avg_reward"],
                "r_cd": r.get("r_cd", 0), "r_cl": r.get("r_cl", 0),
                "r_sim": r.get("r_sim", 0), "sim_raw": r.get("sim_raw", 0),
                "dt_sec": r.get("_dt", 0)} for r in all_results]
    with open(out_dir / "all_seeds.json", "w") as f:
        json.dump(summary, f, indent=2)

    # Save best seed's full signals
    np.savez_compressed(out_dir / "signals.npz",
                        sensors=best["sensors"], forces=best["forces"],
                        actions=best["actions"], rewards=best["rewards"])

    # Vorticity: re-create env, load best seed skeleton, run, capture
    log("  Generating vorticity for best seed...")
    env2 = _create_env(scene, device_id)
    seed_dir = Path(next(s[1] for s in scene["seeds"] if s[0] == best["seed"]))
    norm_path = seed_dir.parent / "vec_normalize.pkl"
    model_path = seed_dir / "best_model.zip"

    wrapped2 = SymmetryAugmentWrapper(env2, prob=0.0, seed=41, rollout_len=2048)
    vec_env2 = DummyVecEnv([lambda: wrapped2])
    vec_env2 = VecNormalize.load(str(norm_path), vec_env2)
    vec_env2.norm_reward = False
    skeleton2 = PPO(
        "MlpPolicy",
        policy_kwargs={"activation_fn": Sin, "net_arch": [64, 64]},
        env=vec_env2, device=_device,
        n_steps=2048, batch_size=64, n_epochs=10,
        learning_rate=3e-4, gamma=0.995, verbose=0,
    )
    sd = _extract_state_dict(str(model_path))
    skeleton2.policy.load_state_dict(sd, strict=False)

    obs = vec_env2.reset()
    for _ in range(num_steps):
        action, _ = skeleton2.predict(obs, deterministic=True)
        obs, reward, done, info = vec_env2.step(action)
        if done[0]:
            break
    _save_vorticity(env2, out_dir / "vorticity_controlled.png", scene)
    env2.close()

    # Target vorticity
    log("  Generating target vorticity...")
    _generate_target_vorticity(scene, device_id, out_dir)

    # Zero-action baseline
    log("  Generating zero-action baseline...")
    env3 = _create_env(scene, device_id)
    zero_omega = env3._action_to_omega(np.zeros(3, dtype=np.float32))
    for _ in range(num_steps):
        smoothed = env3.smoother(zero_omega)
        env3._set_omega(smoothed)
        env3._gpu_block(lambda: env3.sim.run(scene["si"], zero_obs=True))
    _save_vorticity(env3, out_dir / "vorticity_zero.png", scene)
    env3.close()

    # Metrics
    metrics = {
        "best_seed": best["seed"],
        "dtw_sim_v5": float(best.get("sim_raw", 0)),
        "reward_mean": float(best["avg_reward"]),
        "r_cd_mean": float(best.get("r_cd", 0)),
        "r_cl_mean": float(best.get("r_cl", 0)),
        "r_sim_mean": float(best.get("r_sim", 0)),
        "sim_raw_mean": float(best.get("sim_raw", 0)),
        "aF_mean": float(np.mean(best["actions"][:, 0])),
        "aB_mean": float(np.mean(best["actions"][:, 1])),
        "aT_mean": float(np.mean(best["actions"][:, 2])),
    }
    with open(out_dir / "metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)

    env.close()
    log(f"  {scene_id} complete.")


def _generate_target_vorticity(scene: Dict[str, Any], device_id: int,
                                out_dir: Path) -> None:
    from CelerisLab import Simulation

    scene_type = scene["scene_type"]
    config_path = Path(scene["config_path"]).resolve()
    si = scene["si"]
    warmup = int(4.0 * 2000 / 0.01)
    L0 = 20.0
    U0 = 0.01

    sim = Simulation(lbm_config_path=str(config_path), device_id=device_id)

    if scene_type == "karman":
        with open(scene["calibration_path"]) as f:
            cal = json.load(f)
        dist_radius = float(cal.get("dist_radius", 1.0)) * L0
        nx_cfg = int(cal.get("grid", {}).get("nx", 2000))
        ny_cfg = int(cal.get("grid", {}).get("ny", 600))
        CENTER_Y = float(ny_cfg - 1) / 2.0
        DIST_X = 600.0
        SENS_X = 1200.0

        sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=dist_radius)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y + 40.0, 0.0), radius=5.0)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y, 0.0), radius=5.0)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y - 40.0, 0.0), radius=5.0)
    elif scene_type == "illusion":
        with open(scene["calibration_path"]) as f:
            cal = json.load(f)
        target_diam = float(scene.get("target_diam", 1.0))
        nx_cfg = int(cal.get("grid", {}).get("nx", 2000))
        ny_cfg = int(cal.get("grid", {}).get("ny", 600))
        CENTER_Y = float(ny_cfg - 1) / 2.0
        TARGET_X = 400.0
        SENS_X = 600.0

        sim.add_body("circle", center=(TARGET_X, CENTER_Y, 0.0), radius=target_diam * L0)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y + 40.0, 0.0), radius=5.0)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y, 0.0), radius=5.0)
        sim.add_body("sensor", center=(SENS_X, CENTER_Y - 40.0, 0.0), radius=5.0)
    else:
        sim.close()
        return

    sim.initialize()
    sim.run(warmup, zero_obs=True)
    for _ in range(30):
        sim.run(si, zero_obs=True)

    macro = sim.get_macroscopic()
    vort = compute_vorticity(macro["ux"], macro["uy"])
    cylinders = [
        ((ctr[0], ctr[1]), rad)
        for body in [sim.bodies.get(i) for i in range(sim.bodies.count)]
        if hasattr(body, 'center')
        for ctr, rad in [((body.center[0], body.center[1]),
                          body.radius if hasattr(body, 'radius') else 10.0)]
    ]
    render_vorticity_field(vort, nx=nx_cfg, ny=ny_cfg,
                           out_path=str(out_dir / "vorticity_target.png"),
                           cylinders=cylinders, vmin=-0.03, vmax=0.03)
    sim.close()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    global _device
    parser = argparse.ArgumentParser(description="Evaluate V5 train models")
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--scene", type=str, default=None)
    args = parser.parse_args()

    _device = torch.device(f"cuda:{args.device_id}")
    log(f"GPU: {args.device_id}, torch device: {_device}")

    scenes = TRAIN_SCENES
    if args.scene:
        key = args.scene
        scenes = [s for s in TRAIN_SCENES if s["scene_id"] == key
                  or s["scene_id"].startswith(key)]
        if not scenes:
            log(f"ERROR: No scene matching '{args.scene}'")
            return 1

    for i, scene in enumerate(scenes):
        if i > 0:
            prev_cfg = scenes[i - 1]["config_path"]
            curr_cfg = scene["config_path"]
            if prev_cfg != curr_cfg:
                log(f"Waiting 120s before config switch...")
                time.sleep(120)

        evaluate_train_scene(scene, args.device_id)

    log("All train scenes complete.")
    return 0


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