"""Karman cloak DRL inference (Re=100).

Uses utils/cfd_interface.py (copied and verified from SR_analysis).

Usage:
    conda run -n pycuda_3_10 python scripts/collect_karman.py --device 2 --steps 200

Output: data/karman/karman_re200/
"""
from __future__ import annotations

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

import numpy as np

_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

from CCD_analysis.configs import get_scene, get_scene_list, data_dir_for_scene, model_path_for_scene, LEGACY_CFG_DIR
from CCD_analysis.utils.dual_clock import DualClockCollector, field_steps_from_interval
from CCD_analysis.utils.cfd_interface import (
    load_legacy_configs,
    build_karman_cloak_env, add_pinball, build_observation,
    scale_action, load_ppo_model, save_vorticity_png,
    vorticity_from_ddf, compute_similarity,
)

DATA_TYPE = np.float32
L0 = 20.0


def run_single(
    scene_name: str,
    device_id: int,
    n_steps: int,
    *,
    field_interval: int | None = None,
    output_dir: str | None = None,
) -> dict:
    cfg = get_scene(scene_name)
    out_dir = output_dir or data_dir_for_scene(scene_name)
    Path(out_dir).mkdir(parents=True, exist_ok=True)
    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    ac_scale = cfg["action_scale"]
    ac_bias = cfg["action_bias"]
    n_obj = cfg["n_objects_env"]

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))

    # Save config
    with open(os.path.join(out_dir, "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)

    # --- Target recording ---
    print("=== Target recording ===")
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target_states, _ = build_karman_cloak_env(
        ff, u0=u0, l0=L0, sample_interval=si, fifo_len=150, data_type=DATA_TYPE)
    np.savez(os.path.join(out_dir, "target.npz"), target_states=target_states)

    # --- Add pinball + norm ---
    print("=== Norm ===")
    norm = add_pinball(
        ff, l0=L0, u0=u0, sample_interval=si, fifo_len=150, data_type=DATA_TYPE,
        action_bias=ac_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],
    )
    # Save norm (without save_states array)
    norm_json = {k: v for k, v in norm.items() if not isinstance(v, np.ndarray)}
    with open(os.path.join(out_dir, "norm.json"), "w") as f:
        json.dump(norm_json, f, indent=2)

    # --- Uncontrolled rollout ---
    print("=== Uncontrolled ===")
    ff.restore_ddf()
    ff.apply_ddf()
    sens_u, forc_u = [], []
    for _ in range(n_steps):
        ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
        obs_slice = ff.obs.copy()[2:14]
        sens_u.append(obs_slice[0:6])
        forc_u.append(obs_slice[6:12])
    np.savez(os.path.join(out_dir, "uncontrolled.npz"),
             sensors=np.array(sens_u, dtype=np.float32),
             forces=np.array(forc_u, dtype=np.float32))
    save_vorticity_png(os.path.join(out_dir, "vorticity_uncontrolled.png"),
                       vorticity_from_ddf(ff, u0=u0),
                       title=f"{scene_name} uncontrolled")

    # --- Controlled rollout ---
    print("=== Controlled ===")
    model_path = model_path_for_scene(scene_name)
    s_dim = cfg.get("s_dim", 12)
    model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=s_dim)
    model.set_random_seed(0)

    ff.restore_ddf()
    ff.apply_ddf()
    fifo = deque(maxlen=150)
    bias_action = scale_action(np.zeros(3, dtype=np.float32),
                               scale=ac_scale, bias=ac_bias, u0=u0, n_total_bodies=n_obj)
    for _ in range(150):
        ff.context.push()
        ff.run(si, bias_action)
        ff.context.pop()
        fifo.append(ff.obs.copy()[2:14])

    # Save DDF+FIFO checkpoint for replay (state right after warmup, before step 0)
    save_states_arr = np.array(fifo, dtype=DATA_TYPE)
    ff.get_ddf()
    np.save(os.path.join(out_dir, "ddf_checkpoint.npy"), ff.ddf)
    np.save(os.path.join(out_dir, "fifo_checkpoint.npy"), save_states_arr)

    sens_c, forc_c, act_c, rew_c = [], [], [], []
    obs = np.zeros(s_dim, dtype=np.float32)
    collector = None
    if field_interval is not None:
        collector = DualClockCollector(
            ff,
            control_interval=si,
            control_count=n_steps,
            u0=u0,
            field_steps=field_steps_from_interval(si * n_steps, field_interval),
        )

    for step in range(n_steps):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        act_c.append(action.copy())

        action_arr = scale_action(action, scale=ac_scale, bias=ac_bias,
                                  u0=u0, n_total_bodies=n_obj)
        ff.context.push()
        if collector is None:
            ff.run(si, action_arr)
        else:
            collector.run_interval(step, action_arr)
        ff.context.pop()

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

        # Reward
        sarr = np.array(fifo, dtype=np.float32)
        if len(sarr) >= 30:
            f = sarr[-1, 6:12] / norm["force_norm_fact"]
            cd = float((f[0] + f[2] + f[4]) / 3.0)
            cl = float((f[1] + f[3] + f[5]) / 3.0)
            sim = compute_similarity(target_states, sarr[:, 0:6], 30)
            r = min(0.3 * np.exp(-abs(cd * 20)) + 0.4 * np.exp(-abs(cl * 80))
                    + 0.3 * np.exp(-10 * abs(sim - 1)), 1.0)
            rew_c.append(float(r))

    sens_arr = np.array(sens_c, dtype=np.float32)
    forc_arr = np.array(forc_c, dtype=np.float32)
    act_arr = np.array(act_c, dtype=np.float32)
    rew_arr = np.array(rew_c, dtype=np.float32)

    np.savez(os.path.join(out_dir, "controlled.npz"),
             sensors=sens_arr, forces=forc_arr, actions=act_arr, rewards=rew_arr)
    if collector is not None:
        collector.save(Path(out_dir) / "fields.npz")

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

    avg_reward = float(np.mean(rew_arr[-100:])) if len(rew_arr) >= 100 else 0.0
    sim_score = compute_similarity(target_states, sens_arr, 30)
    print(f"  reward={avg_reward:.4f}  similarity={sim_score:.4f}")

    result = {
        "scene": scene_name,
        "similarity": sim_score,
        "avg_reward": avg_reward,
        "control_interval": si,
        "field_interval": field_interval,
    }
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    del ff, model
    return result


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, default="karman_re100",
                    help="Scene name (karman_re50, re100, re200, re400)")
    ap.add_argument("--re", type=int, default=None,
                    help="Re number shortcut (50, 100, 200, 400)")
    ap.add_argument("--device", type=int, default=2)
    ap.add_argument("--steps", type=int, default=200)
    ap.add_argument("--field-interval", type=int, default=None,
                    help="Save fields every N absolute lattice steps")
    ap.add_argument("--output-dir", type=str, default=None,
                    help="Optional isolated output directory")
    args = ap.parse_args()

    if args.re is not None:
        scene_name = f"karman_re{args.re}"
    else:
        scene_name = args.scene

    if scene_name not in get_scene_list("karman"):
        print(f"Unknown scene: {scene_name}. Available: {get_scene_list('karman')}")
        return 1

    t0 = time.time()
    r = run_single(
        scene_name, args.device, args.steps,
        field_interval=args.field_interval, output_dir=args.output_dir,
    )
    print(f"Done in {time.time()-t0:.1f}s: sim={r['similarity']:.4f}")


if __name__ == "__main__":
    sys.exit(main())
