"""Inference pipeline for Illusion (cylinder imitation) scenes.

Generates controlled data for a given target cylinder diameter using
LegacyCelerisLab + trained PPO model.

Usage:
    conda run -n pycuda_3_10 python scripts/infer_illusion.py \\
        --diameter 1.0 --device 0
    conda run -n pycuda_3_10 python scripts/infer_illusion.py \\
        --diameter 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 LegacyCelerisLab import utils as legacy_utils  # 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_illusion(
    scene_name: str,
    device_id: int,
    output_root: str,
    n_infer_steps: int = 200,
) -> dict:
    """Run full inference pipeline for one Illusion 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_total = cfg["n_objects_env"]
    sensor_x = cfg["sensor_x"]  # 30.0 for illusion
    front_x = cfg["pinball_front_x"]  # 19.0
    rear_x = cfg["pinball_rear_x"]  # 20.3
    target_diam = cfg["target_diameter"]

    os.makedirs(output_root, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"Scene: {scene_name}  Diam={target_diam}L  u0={u0}  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 and velocity
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))
    if u0 != 0.01:
        field_cfg = field_cfg._replace(velocity=float(u0))

    # -- Phase 1: Target recording (target cylinder + 3 sensors) ------------
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]

    # Add target cylinder at x=20*L0, with radius = target_diam * L0
    print(f"  Adding target cylinder: diam={target_diam}L, pos=({20*l0:.0f}, {ny/2:.0f})")
    ff.add_cylinder((20.0 * l0, (ny - 1) / 2, 0.0), target_diam * l0)

    # Add 3 sensors at x = sensor_x * L0
    for y_off in [2.0, 0.0, -2.0]:
        sc = (sensor_x * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)

    n_obj_phase1 = ff.obs.size // 2
    print(f"  Phase 1 objects: {n_obj_phase1}")

    # Stabilize
    stabilize_steps = int(4 * ff.FIELD_SHAPE[0] / u0)
    print(f"  Stabilising ({stabilize_steps} steps)...")
    ff.run(stabilize_steps, np.zeros(n_obj_phase1, dtype=DATA_TYPE))

    # Record target
    target_states = np.empty((0, 8), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, np.zeros(n_obj_phase1, dtype=DATA_TYPE))
        new_state = ff.obs.copy()[0:8]  # sensor[6] + cylinder force[2]
        target_states = np.vstack((target_states, new_state))
    print(f"  Target recorded: {target_states.shape}")

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

    # Clean up and create pinball env
    del ff
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    # -- Phase 2: Pinball env -----------------------------------------------
    # Add 3 sensors (same positions as target phase)
    for y_off in [2.0, 0.0, -2.0]:
        sc = (sensor_x * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)

    # Add 3 pinball cylinders (illusion positions)
    # Front at x=front_x*L0, rear at x=rear_x*L0
    ff.add_cylinder((front_x * l0, (ny - 1) / 2, 0.0), l0 / 2.0)
    ff.add_cylinder((rear_x * l0, (ny - 1) / 2 + 0.75 * l0, 0.0), l0 / 2.0)
    ff.add_cylinder((rear_x * 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 objects, got {n_obj}"

    # Stabilize with zero action
    print(f"  Stabilising pinball ({stabilize_steps} steps)...")
    ff.run(stabilize_steps, np.zeros(n_obj, dtype=DATA_TYPE))

    # Checkpoint
    ff.get_ddf()
    ff.save_ddf()

    # 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()[0:12])

    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-action rollout
    ff.apply_ddf()
    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)
    print(f"  bias action: {bias_arr}")

    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[0:12])
    save_states = np.array(list(fifo), dtype=DATA_TYPE)
    norm["save_states"] = save_states
    ff.apply_ddf()

    # Save norm
    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:
        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()

        # Re-bias FIFO
        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()[0:12])

        sens_list, forc_list, action_list = [], [], []
        obs = np.zeros(s_dim, 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())

            # Convert to legacy action array (6 objects: sensors[3] + pinball[3])
            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()[0:12]
            fifo.append(obs_slice)
            sens_list.append(obs_slice[0:6])
            forc_list.append(obs_slice[6:12])

            # Build normalized obs (just forces + sens for S_DIM=12)
            forces_norm = obs_slice[6:12] / force_norm_fact
            sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
            obs12 = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)

            if s_dim == 14:
                # Need target values -- for inference we zero-pad
                obs = np.zeros(14, dtype=np.float32)
                obs[:12] = obs12
            else:
                obs = obs12

        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))

        # Compute similarity (use the target cylinder's sensor-only signals)
        # For comparison, compute similarity between controlled sensors and target
        target_sensors = target_states[:, 0:6]
        sim = compute_similarity(target_sensors,
                                 np.array(sens_list, dtype=np.float32), CONV_LEN)
        print(f"  similarity (vs target cylinder) = {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="Illusion inference")
    ap.add_argument("--diameter", type=str, default="1.0",
                    help='Diameter: 0.75, 1.0, 1.5, or "all"')
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--steps", type=int, default=200)
    ap.add_argument("--out-root", type=str, default=None)
    args = ap.parse_args()

    # Determine scene names
    if args.diameter.lower() == "all":
        scene_names = get_scene_list("illusion")
    else:
        d = float(args.diameter)
        # Match by target_diameter field
        scene_names = []
        for sn in get_scene_list("illusion"):
            cfg = get_scene(sn)
            if abs(cfg["target_diameter"] - d) < 0.01:
                scene_names.append(sn)
        if not scene_names:
            print(f"ERROR: no illusion scene found for diameter={d}")
            return 1

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