"""Closed-loop validator for Vortex cloak scenes.

Adapted from run_closed_loop.py for vortex geometry:
- No disturbance cylinder (obs[0:12], n_objects=6)
- Action scaling: action*4 + [0,-4,4]
- Transient: n_steps=150, step-aligned similarity (no phase lag)
- Supports --mode pysr (Karman joint formula) and --mode v23 (vortex SINDy)

Usage:
    conda run -n pycuda_3_10 python src/SR_analysis/validate/run_closed_loop_vortex.py \
        --scene vortex_lamb --device 2 --steps 150 --mode pysr \
        --pysr-front validate/results/karman_joint_deep_front.json \
        --pysr-top validate/results/karman_joint_deep_top.json

    conda run -n pycuda_3_10 python src/SR_analysis/validate/run_closed_loop_vortex.py \
        --scene vortex_lamb --device 2 --mode v23 \
        --sindy-results sindy/vortex/sindy_results_v2.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from collections import deque
from typing import Any, Dict, 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, compute_similarity, scale_action,
)
from SR_analysis.configs import get_scene, LEGACY_CFG_DIR, FIFO_LEN, CONV_LEN


DATA_TYPE = np.float32
NX = 1280


def run_validation_vortex(
    scene_name: str,
    device_id: int,
    n_steps: int = 150,
    sindy_path: Optional[str] = None,
    threshold: Optional[float] = None,
    mode: str = "pysr",
    front_func=None,
    top_func=None,
    feat_names_front: Optional[list] = None,
    feat_names_rear: Optional[list] = None,
    out_dir: Optional[str] = None,
) -> dict:
    """Run closed-loop validation for a Vortex scene.

    Parameters
    ----------
    scene_name : str
        Scene name ("vortex_lamb" or "vortex_taylor")
    mode : str
        "pysr" for PySR formula evaluation, "v23" for SINDy coefficients
    front_func, top_func :
        PySR callable functions (required for mode="pysr")
    feat_names_front, feat_names_rear :
        Feature name lists (required for mode="pysr")
    """
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    mu = cfg["mu"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    action_scale = cfg["action_scale"]  # 4.0 for vortex
    action_bias = cfg["action_bias"]    # (0.0, -4.0, 4.0)
    n_obj_total = cfg["n_objects_env"]  # 6
    vtype = cfg["vortex_type"]
    vstrength = cfg["vortex_strength"]
    max_steps = cfg.get("max_steps", 150)

    if n_steps <= 0:
        n_steps = max_steps

    if out_dir is None:
        out_dir = os.path.join(os.path.dirname(__file__), "results")
    os.makedirs(out_dir, exist_ok=True)

    if mode == "pysr":
        from SR_analysis.validate.predict_pysr import pysr_predict_v23
    elif mode == "v23":
        from SR_analysis.validate.run_closed_loop import (
            load_sindy_coefs as _lsc, predict_v23,
        )
        if sindy_path is None:
            sindy_path = os.path.join(
                os.path.dirname(__file__), "..", "sindy", "vortex", "sindy_results_v2.json")
        coefs = _lsc(sindy_path, scene_name, threshold=threshold)
    else:
        raise ValueError(f"Unknown mode: {mode}")

    print(f"\n=== Validating {scene_name} (mode={mode}, device={device_id}, steps={n_steps}) ===")

    # -------------------------------------------------------------------
    # Phase 1: Sensor-only env, record target with vortex
    # -------------------------------------------------------------------
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]))

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

    # 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

    # Short stabilize (1*NX/U0)
    stabilize_steps_short = int(1 * NX / u0)
    print(f"  stabilising sensors ({stabilize_steps_short} steps)...")
    ff.run(stabilize_steps_short, np.zeros(n_obj_sensors, dtype=DATA_TYPE))

    # Save clean flow DDF
    ff.get_ddf()
    ff.save_ddf()

    # Add vortex at x=10*L0 (target position)
    ff.add_vortex((10.0 * l0, (ny - 1) / 2, 0.0),
                  2.0 * l0, vstrength * u0, 0, vtype)

    # Record target
    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}")

    # -------------------------------------------------------------------
    # Phase 2: Pinball env, norm collection
    # -------------------------------------------------------------------
    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
    assert n_obj == 6, f"Expected 6 objects, got {n_obj}"

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

    # Stabilize with bias action
    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 position)
    ff.add_vortex((15.0 * l0, (ny - 1) / 2, 0.0),
                  2.0 * l0, vstrength * u0, 0, vtype)

    # Save DDF after vortex (mid-transient reset 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])))

    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    # Bias rollout
    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)
    ff.apply_ddf()

    # Re-fill FIFO from save_states
    fifo = deque(maxlen=FIFO_LEN)
    for row in save_states:
        fifo.append(row)

    # -------------------------------------------------------------------
    # Phase 3: Closed-loop inference
    # -------------------------------------------------------------------
    sens_list, actions_list = [], []
    a_prev = np.zeros(3, dtype=np.float64)  # normalized action [0,0,0]
    a_prev2 = a_prev.copy()
    dt_c = sample_interval / 2000.0  # T0 = D/U0 = 2000 steps

    for step in range(n_steps):
        obs = fifo[-1] if fifo else np.zeros(12, dtype=np.float32)
        obs_prev = list(fifo)[-2] if len(fifo) >= 2 else obs

        # Convert a_prev from normalized to physical omega
        a_prev_phys = (a_prev * action_scale + np.array(action_bias, dtype=np.float64)) * u0
        a_prev2_phys = (a_prev2 * action_scale + np.array(action_bias, dtype=np.float64)) * u0

        if mode == "pysr":
            # Predict using PySR formula
            feat_keys_front = [k for k in feat_names_front if k != "bias"]
            feat_keys_rear = [k for k in feat_names_rear if k != "bias"]
            omega_phys = pysr_predict_v23(
                obs, a_prev_phys, a_prev2_phys, mu, u0, dt_c,
                front_func, top_func,
                feat_keys_front, feat_keys_rear,
                sensors_raw=np.vstack([obs_prev[0:6], obs[0:6]]),
                forces_raw=np.vstack([obs_prev[6:12], obs[6:12]]),
            )
        elif mode == "v23":
            omega_phys = predict_v23(
                obs, a_prev_phys, a_prev2_phys, mu, u0,
                coefs["front_coef"], coefs["top_coef"],
                coefs["feat_names_front"], coefs["feat_names_rear"],
                feat_keys_front=coefs["feat_keys_front"],
                feat_keys_rear=coefs["feat_keys_rear"],
            )
        else:
            raise ValueError(f"Unknown mode: {mode}")

        # Convert physical omega -> normalized action -> legacy array
        norm_a = (omega_phys / u0 - np.array(action_bias, dtype=np.float64)) / action_scale
        norm_a = np.clip(norm_a, -1.0, 1.0).astype(np.float32)
        action_arr = scale_action(norm_a, scale=action_scale, bias=action_bias,
                                  u0=u0, n_total_bodies=n_obj_total)

        ff.run(sample_interval, action_arr)

        obs_new = ff.obs.copy()
        fifo.append(obs_new)
        sens_list.append(obs_new[0:6])
        actions_list.append(omega_phys.copy())

        a_prev2 = a_prev.copy()
        a_prev = norm_a.copy().astype(np.float64)

    del ff

    # Similarity: step-aligned (no phase lag for transient)
    sens_arr = np.array(sens_list, dtype=np.float32)
    target_arr = target_states[:n_steps]
    n_align = min(sens_arr.shape[0], target_arr.shape[0])

    if n_align >= CONV_LEN:
        sim = compute_similarity(target_arr, sens_arr[:n_align], CONV_LEN)
    else:
        sim = 0.0

    action_range = float(np.max(np.abs(actions_list)))
    print(f"  similarity={sim:.4f}  action_range={action_range:.4f}")

    return {
        "scene": scene_name,
        "mode": mode,
        "similarity": sim,
        "action_range": action_range,
        "n_steps": n_steps,
        "threshold": threshold if threshold is not None else "best",
    }


def main():
    ap = argparse.ArgumentParser(description="Vortex closed-loop validation")
    ap.add_argument("--scene", type=str, required=True,
                    help="Scene name: vortex_lamb or vortex_taylor")
    ap.add_argument("--device", type=int, default=2, help="GPU device")
    ap.add_argument("--steps", type=int, default=150,
                    help="Validation steps (default: 150 for vortex)")
    ap.add_argument("--mode", type=str, default="pysr",
                    choices=["pysr", "v23"],
                    help='Control law mode: "pysr" or "v23"')
    ap.add_argument("--sindy-results", type=str, default=None,
                    help="Path to sindy_results.json (for --mode v23)")
    ap.add_argument("--threshold", type=float, default=None,
                    help="SINDy threshold (for --mode v23)")
    ap.add_argument("--pysr-front", type=str, default=None,
                    help="PySR front JSON (required for --mode pysr)")
    ap.add_argument("--pysr-top", type=str, default=None,
                    help="PySR top JSON (required for --mode pysr)")
    ap.add_argument("--out", type=str, default=None)
    args = ap.parse_args()

    if args.mode == "pysr":
        if args.pysr_front is None or args.pysr_top is None:
            raise ValueError("--pysr-front and --pysr-top required for --mode pysr")
        from SR_analysis.validate.predict_pysr import (
            load_pysr_formulas, build_pysr_functions,
        )
        front_expr, top_expr, fn_f, fn_r = load_pysr_formulas(
            args.pysr_front, args.pysr_top)
        front_func, top_func = build_pysr_functions(
            front_expr, top_expr, fn_f, fn_r)
        result = run_validation_vortex(
            args.scene, args.device, n_steps=args.steps, mode="pysr",
            front_func=front_func, top_func=top_func,
            feat_names_front=fn_f, feat_names_rear=fn_r,
            out_dir=args.out,
        )
    else:
        result = run_validation_vortex(
            args.scene, args.device, n_steps=args.steps, mode="v23",
            sindy_path=args.sindy_results, threshold=args.threshold,
            out_dir=args.out,
        )

    out_name = f"{args.scene}_{args.mode}"
    out_dir = args.out or os.path.join(os.path.dirname(__file__), "results")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{out_name}.json")
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Saved: {out_path}")


if __name__ == "__main__":
    main()
