"""Inference pipeline for Karman cloak across Re.

Generates controlled/uncontrolled data for a given Re case using
LegacyCelerisLab + trained PPO model.

Usage:
    conda run -n pycuda_3_10 python scripts/infer_karman.py --re 100 --device 0
    conda run -n pycuda_3_10 python scripts/infer_karman.py --re all --device 2
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque

import numpy as np

# Add repo root for LegacyCelerisLab and src/ for SR_analysis
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from LegacyCelerisLab import FlowField  # noqa: E402

from SR_analysis.utils.cfd_interface import (
    nu_from_re, load_legacy_configs,
    build_karman_cloak_env, add_pinball, build_observation,
    scale_action, load_ppo_model, save_vorticity_png,
    vorticity_from_ddf, compute_similarity, ACTION_SMOOTH_WEIGHT,
)
from SR_analysis.configs import (
    SCENES, get_scene, get_scene_list, model_path_for_scene,
    LEGACY_CFG_DIR, FIFO_LEN, CONV_LEN,
)

DATA_TYPE = np.float32


def run_single_re(
    scene_name: str,
    device_id: int,
    output_root: str,
    n_infer_steps: int = 200,
) -> dict:
    """Run full inference pipeline for one Karman Re case."""
    cfg = get_scene(scene_name)
    re_code = cfg["re_code"]
    nu = cfg["nu"]
    u0 = cfg["u0"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    action_scale = cfg["action_scale"]
    action_bias = cfg["action_bias"]
    n_obj_total = cfg["n_objects_env"]

    os.makedirs(output_root, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"Scene: {scene_name}  Re_code={re_code}  nu={nu:.6f}  device={device_id}")
    print(f"{'='*60}")

    # Save config
    with open(os.path.join(output_root, "config.json"), "w") as f:
        json.dump({k: str(v) if not isinstance(v, (int, float, list, bool))
                    else v for k, v in cfg.items()}, f, indent=2)

    # Load legacy CFD configs with overridden viscosity
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))

    # Build env: dist cylinder + sensors, record target
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target_states, env_info = build_karman_cloak_env(
        ff, u0=u0, l0=l0, sample_interval=sample_interval,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
    )
    np.savez(os.path.join(output_root, "target.npz"), target_states=target_states)

    # Add pinball, compute norm
    norm = add_pinball(
        ff, l0=l0, u0=u0, sample_interval=sample_interval,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
        action_bias=action_bias, pinball_front_x=cfg["pinball_front_x"],
        pinball_rear_x=cfg["pinball_rear_x"],
        obs_slice_start=cfg["obs_slice"][0], obs_slice_end=cfg["obs_slice"][1],
    )

    norm_for_json = {k: v for k, v in norm.items()
                     if not isinstance(v, np.ndarray)}
    with open(os.path.join(output_root, "norm.json"), "w") as f:
        json.dump(norm_for_json, f, indent=2)

    # Uncontrolled rollout
    print("  uncontrolled rollout ...")
    ff.restore_ddf()
    ff.apply_ddf()
    sens_list, forc_list = [], []
    for _ in range(n_infer_steps):
        ff.run(sample_interval, np.zeros(n_obj_total, dtype=DATA_TYPE))
        obs_slice = ff.obs.copy()[2:14]
        sens_list.append(obs_slice[0:6])
        forc_list.append(obs_slice[6:12])

    np.savez(os.path.join(output_root, "uncontrolled.npz"),
             sensors=np.array(sens_list, dtype=np.float32),
             forces=np.array(forc_list, dtype=np.float32))

    omega_unc = vorticity_from_ddf(ff, u0=u0)
    save_vorticity_png(os.path.join(output_root, "vorticity_uncontrolled.png"),
                       omega_unc, title=f"{scene_name} uncontrolled")

    # Controlled rollout
    result = {"scene": scene_name, "controlled": False}
    model_path = model_path_for_scene(scene_name)

    if model_path is not None:
        s_dim = cfg.get("s_dim", 12)
        print(f"  loading model: {model_path}  (s_dim={s_dim})")
        model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=s_dim)
        model.set_random_seed(0)

        print(f"  controlled rollout ({n_infer_steps} steps) ...")
        ff.restore_ddf()
        ff.apply_ddf()

        # Bias FIFO init
        fifo = deque(maxlen=FIFO_LEN)
        bias_action = scale_action(
            np.zeros(3, dtype=np.float32),
            scale=action_scale, bias=action_bias, u0=u0,
            n_total_bodies=n_obj_total,
        )
        for _ in range(FIFO_LEN):
            ff.context.push()
            ff.run(sample_interval, bias_action)
            ff.context.pop()
            fifo.append(ff.obs.copy()[2:14])

        sens_list_c, forc_list_c, action_list_c = [], [], []
        reward_list_c = []
        obs = np.zeros(12, dtype=np.float32)

        for step in range(n_infer_steps):
            action, _states = model.predict(obs, deterministic=True)
            action = action.astype(np.float32).flatten()
            action_list_c.append(action.copy())

            action_arr = scale_action(
                action, scale=action_scale, bias=action_bias,
                u0=u0, n_total_bodies=n_obj_total,
            )

            ff.context.push()
            ff.run(sample_interval, action_arr)
            ff.context.pop()

            obs_slice = ff.obs.copy()[2:14]
            fifo.append(obs_slice)
            sens_list_c.append(obs_slice[0:6])
            forc_list_c.append(obs_slice[6:12])
            obs = build_observation(obs_slice, norm)

            # Compute reward
            states_arr = np.array(list(fifo), dtype=np.float32)
            if len(states_arr) >= CONV_LEN:
                forces = states_arr[-1, 6:12] / norm["force_norm_fact"]
                cd = float((forces[0] + forces[2] + forces[4]) / 3.0)
                cl = float((forces[1] + forces[3] + forces[5]) / 3.0)
                sim = compute_similarity(target_states, states_arr[:, 0:6], CONV_LEN)
                r_cd = np.exp(-abs(cd * 20.0))
                r_cl = np.exp(-abs(cl * 80.0))
                r_sim = np.exp(-10.0 * abs(sim - 1.0))
                reward = min(0.3 * r_cd + 0.4 * r_cl + 0.3 * r_sim, 1.0)
                reward_list_c.append(float(reward))

        np.savez(os.path.join(output_root, "controlled.npz"),
                 sensors=np.array(sens_list_c, dtype=np.float32),
                 forces=np.array(forc_list_c, dtype=np.float32),
                 actions=np.array(action_list_c, dtype=np.float32),
                 rewards=np.array(reward_list_c, dtype=np.float32))

        omega_con = vorticity_from_ddf(ff, u0=u0)
        save_vorticity_png(os.path.join(output_root, "vorticity_controlled.png"),
                           omega_con, title=f"{scene_name} controlled")

        avg_reward = (float(np.mean(reward_list_c[-100:]))
                      if len(reward_list_c) >= 100
                      else float(np.mean(reward_list_c)))
        sim_score = compute_similarity(
            target_states, np.array(sens_list_c, dtype=np.float32), CONV_LEN)
        result["controlled"] = True
        result["avg_reward_last100"] = avg_reward
        result["similarity"] = sim_score
        print(f"  avg_reward(last100)={avg_reward:.4f}  similarity={sim_score:.4f}")
    else:
        print(f"  no model for {scene_name}, skipping controlled rollout")

    del ff

    with open(os.path.join(output_root, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    return result


def main():
    ap = argparse.ArgumentParser(description="Karman cloak inference")
    ap.add_argument("--re", type=str, default="100",
                    help='Re case: 50,100,200,400, or "all", or "validation"')
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--steps", type=int, default=200,
                    help="Number of inference steps per rollout")
    ap.add_argument("--out-root", type=str, default=None,
                    help="Output root (default: SR_analysis/data/karman)")
    args = ap.parse_args()

    # Determine scene names
    selection = args.re.lower()
    if selection == "all":
        scene_names = get_scene_list("karman")
    elif selection == "validation":
        scene_names = [f"karman_re{rc}" for rc in [35, 70, 150]]
        # Check which are defined
        scene_names = [s for s in scene_names if s in SCENES]
    else:
        rc = int(selection)
        scene_names = [f"karman_re{rc}"]

    if args.out_root is None:
        out_root = os.path.join(os.path.dirname(__file__), "..", "data", "karman")
    else:
        out_root = args.out_root

    t_start = time.time()

    for sn in scene_names:
        case_dir = os.path.join(out_root, sn)
        result = run_single_re(sn, args.device, case_dir,
                               n_infer_steps=args.steps)
        print(f"  Done: {sn}  -> {case_dir}")

    elapsed = time.time() - t_start
    print(f"\nTotal time: {elapsed:.1f}s")


if __name__ == "__main__":
    main()
