"""Inference pipeline for Vortex cloak (Lamb dipole + Taylor monopole).

Generates controlled data for vortex cloak scenes using
LegacyCelerisLab + trained PPO model.

Vortex env characteristics:
- No disturbance cylinder
- Vortex is added AFTER DDF checkpoint
- Transient: MAX_STEPS=150
- Action scaling: action*4 + [0,-4,4]

Usage:
    conda run -n pycuda_3_10 python scripts/infer_vortex.py \\
        --type lamb --device 0
    conda run -n pycuda_3_10 python scripts/infer_vortex.py \\
        --type all --device 2
"""
from __future__ import annotations

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

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  # noqa: E402

from SR_analysis.utils.cfd_interface import (
    load_legacy_configs, build_observation,
    scale_action, load_ppo_model, compute_similarity,
)
from SR_analysis.configs import (
    get_scene, get_scene_list, model_path_for_scene,
    LEGACY_CFG_DIR, FIFO_LEN, CONV_LEN,
)

DATA_TYPE = np.float32


def run_single_vortex(
    scene_name: str,
    device_id: int,
    output_root: str,
    n_infer_steps: Optional[int] = None,
) -> dict:
    """Run full inference pipeline for one Vortex scene."""
    cfg = get_scene(scene_name)
    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_pinball = cfg["n_objects_env"]
    max_steps = cfg.get("max_steps", 150)
    vtype = cfg["vortex_type"]
    vstrength = cfg["vortex_strength"]

    if n_infer_steps is None:
        n_infer_steps = max_steps  # transient -- use full episode

    os.makedirs(output_root, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"Scene: {scene_name}  Vortex={vtype} strength={vstrength}  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
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]

    # -- Phase 1: Sensors only env, record target with vortex -----------------
    # Add 3 sensors at x=40*L0
    for y_off in [2.0, 0.0, -2.0]:
        sc = (40.0 * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)

    n_obj_sensors = ff.obs.size // 2
    print(f"  Sensor-only objects: {n_obj_sensors}")

    # Short stabilize (1*NX/U0 instead of 4* for vortex)
    stabilize_steps_short = int(1 * ff.FIELD_SHAPE[0] / u0)
    ff.run(stabilize_steps_short, np.zeros(n_obj_sensors, dtype=DATA_TYPE))

    # Save clean flow DDF (for later restore)
    ff.get_ddf()
    ff.save_ddf()

    # Add vortex
    print(f"  Adding vortex: type={vtype}, center=({10*l0:.0f}, {ny/2:.0f})")
    ff.add_vortex((10.0 * l0, (ny - 1) / 2, 0.0),
                  2.0 * l0, vstrength * u0, 0, vtype)

    # Record target (vortex evolving through sensor-only env)
    target_states = np.empty((0, 6), dtype=DATA_TYPE)
    for _ in range(max_steps):
        ff.run(sample_interval, np.zeros(n_obj_sensors, dtype=DATA_TYPE))
        target_states = np.vstack((target_states, ff.obs.copy()))
    print(f"  Target recorded: {target_states.shape}")

    np.savez(os.path.join(output_root, "target.npz"), target_states=target_states)

    # -- Phase 2: Restore clean flow, add pinball, add vortex, record norm ----
    ff.restore_ddf()
    ff.apply_ddf()

    # Add 3 pinball cylinders
    ff.add_cylinder((30.0 * l0, (ny - 1) / 2, 0.0), l0 / 2.0)
    ff.add_cylinder((31.3 * l0, (ny - 1) / 2 + 0.75 * l0, 0.0), l0 / 2.0)
    ff.add_cylinder((31.3 * l0, (ny - 1) / 2 - 0.75 * l0, 0.0), l0 / 2.0)

    n_obj = ff.obs.size // 2
    print(f"  Pinball env objects: {n_obj}")
    assert n_obj == 6, f"Expected 6, got {n_obj}"

    # Stabilize with zero action
    ff.run(stabilize_steps_short, np.zeros(n_obj, dtype=DATA_TYPE))

    # Stabilize with bias action (following vortex env code)
    bias_arr = np.zeros(n_obj, dtype=DATA_TYPE)
    bias_arr[3] = float(action_bias[0] * u0)
    bias_arr[4] = float(action_bias[1] * u0)
    bias_arr[5] = float(action_bias[2] * u0)
    ff.run(stabilize_steps_short, bias_arr)

    # Add vortex at x=15*L0 (pinball env) and save DDF
    print(f"  Adding vortex for pinball env: type={vtype}")
    ff.add_vortex((15.0 * l0, (ny - 1) / 2, 0.0),
                  2.0 * l0, vstrength * u0, 0, vtype)

    # SAVE DDF AFTER vortex (vortex env reset restores this mid-transient state)
    ff.get_ddf()
    ff.save_ddf()
    print("  DDF saved with vortex active")

    # Norm collection (zero action)
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, np.zeros(n_obj, dtype=DATA_TYPE))
        fifo.append(ff.obs.copy())

    temp_states = np.array(fifo, dtype=DATA_TYPE)
    force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
    sens_deviation = np.mean(temp_states[:, 0:6], axis=0).astype(DATA_TYPE)
    sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
    for i in range(6):
        sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))

    norm = {
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation.tolist(),
        "sens_norm_fact": sens_norm_fact.tolist(),
        "action_bias": list(action_bias),
    }
    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    # Bias rollout (restore before vortex was added, then add vortex again)
    ff.apply_ddf()  # restore to DDF with vortex active
    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy())
    save_states = np.array(list(fifo), dtype=DATA_TYPE)
    norm["save_states"] = save_states
    ff.apply_ddf()

    norm_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_json, f, indent=2)

    # -- Phase 3: Controlled inference ----------------------------------------
    result = {"scene": scene_name, "controlled": False}
    model_path = model_path_for_scene(scene_name)

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

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

        # Bias FIFO init (restore with vortex active)
        fifo = deque(maxlen=FIFO_LEN)
        for _ in range(FIFO_LEN):
            ff.context.push()
            ff.run(sample_interval, bias_arr)
            ff.context.pop()
            fifo.append(ff.obs.copy())

        sens_list, forc_list, action_list = [], [], []
        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.append(action.copy())

            # Action: action*4 + [0,-4,4]
            temp = np.zeros(n_obj, dtype=DATA_TYPE)
            temp[3:6] = np.array(
                (action * action_scale + list(action_bias)) * u0,
                dtype=DATA_TYPE)

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

            obs_slice = ff.obs.copy()
            fifo.append(obs_slice)
            sens_list.append(obs_slice[0:6])
            forc_list.append(obs_slice[6:12])

            forces_norm = obs_slice[6:12] / force_norm_fact
            sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
            obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)

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

        # Similarity: align by step index (no lag for transient)
        states_arr = np.array(sens_list, dtype=np.float32)
        target_arr = target_states[:n_infer_steps, :] if n_infer_steps <= max_steps else target_states

        n_align = min(states_arr.shape[0], target_arr.shape[0])
        if n_align >= CONV_LEN:
            sim = compute_similarity(target_arr, states_arr[:n_align], CONV_LEN)
        else:
            sim = 0.0
        print(f"  similarity (vs target) = {sim:.4f}")

        result["controlled"] = True
        result["similarity"] = sim
    else:
        print(f"  WARNING: no model for {scene_name}")

    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="Vortex cloak inference")
    ap.add_argument("--type", type=str, default="lamb",
                    help='Vortex type: lamb, taylor, or "all"')
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--steps", type=int, default=None,
                    help="Inference steps (default: max_steps for scene)")
    ap.add_argument("--out-root", type=str, default=None)
    args = ap.parse_args()

    if args.type.lower() == "all":
        scene_names = get_scene_list("vortex")
    else:
        scene_names = [f"vortex_{args.type.lower()}"]

    if args.out_root is None:
        out_root = os.path.join(os.path.dirname(__file__), "..", "data", "vortex")
    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_vortex(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()
