"""Closed-loop validator for Illusion scenes.

Builds an illusion environment (no disturbance cylinder, 6 objects, obs[0:12]),
loads target data from existing target.npz, and validates SINDy control laws
in closed-loop CFD with v23 mode (front no-bias + rear shared-head).

Usage (from repo root):
    conda run -n pycuda_3_10 python src/SR_analysis/validate/run_closed_loop_illusion.py \\
        --scene illusion_1L --device 0 --steps 200

Default steps = max(320, 2*NX/U0/sample_interval) to ensure control propagates.
"""
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,
    vorticity_from_ddf, save_vorticity_png,
)
from SR_analysis.configs import get_scene, LEGACY_CFG_DIR, FIFO_LEN

from SR_analysis.validate.run_closed_loop import (
    predict_v23, predict_v23_deriv, load_sindy_coefs,
)

# Import target force harmonic reconstruction from inference pipeline
from SR_analysis.scripts.infer_illusion import gen_target_states_at

DATA_TYPE = np.float32
NX = 1280


def run_validation_illusion(
    scene_name: str,
    sindy_path: str,
    device_id: int,
    n_steps: int = 0,
    threshold: Optional[float] = None,
    out_dir: Optional[str] = None,
    mode_override: Optional[str] = None,
) -> dict:
    """Run closed-loop validation for an Illusion scene.

    Parameters
    ----------
    n_steps : int, default=0 (auto: >= 2*NX/U0 / sample_interval)
    mode_override : str, optional
        Override auto-detected mode: "v23", "absolute", or None for auto-detect.
    """
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    mu = cfg["mu"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    conv_len = cfg.get("conv_len", 36)
    action_scale = cfg["action_scale"]
    action_bias = cfg["action_bias"]
    n_obj_total = cfg["n_objects_env"]
    sensor_x = cfg["sensor_x"]
    front_x = cfg["pinball_front_x"]
    rear_x = cfg["pinball_rear_x"]

    # Auto-set steps: control must propagate at least 1*NX/U0 LBM steps
    min_steps = int(1 * NX / u0 / sample_interval)
    if n_steps == 0 or n_steps < min_steps:
        n_steps = max(min_steps, 200)
        print(f"  auto-set steps={n_steps} (min_steps={min_steps})")

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

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

    # Load SINDy coefficients
    coefs = load_sindy_coefs(sindy_path, scene_name, threshold=threshold)
    threshold_str = f"_th{threshold}" if threshold is not None else ""

    # Load target data (sensors for similarity, harmonics for target force reconstruction)
    data_dir = os.path.join(
        os.path.dirname(__file__), "..", "data", "illusion", scene_name,
    )
    target_npz = np.load(os.path.join(data_dir, "target.npz"))
    if "target_sensors" in target_npz:
        target_states = target_npz["target_sensors"]
    else:
        target_states = target_npz["target_states"][:, 2:8]
    print(f"  target loaded: {target_states.shape}")

    # Load target harmonics for real-time target force reconstruction
    target_harmonics = None
    harm_path = os.path.join(data_dir, "target_harmonics.json")
    if os.path.isfile(harm_path):
        with open(harm_path) as f:
            target_harmonics = json.load(f)
        print(f"  target harmonics loaded: {len(target_harmonics)} channels")

    # Build environment
    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, then 3 pinball cylinders (illusion positions)
    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)
    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
    assert n_obj == 6, f"Expected 6 objects, got {n_obj}"

    # Stabilize with zero action (4*NX/U0 = 512000 steps)
    stabilize_steps = int(4 * NX / u0)
    print(f"  stabilising pinball ({stabilize_steps} steps)...")
    ff.run(stabilize_steps, np.zeros(n_obj, dtype=DATA_TYPE))

    # --- save_ddf(1): post-stabilization checkpoint ---
    ff.get_ddf()
    ff.save_ddf()

    # Norm collection (zero action)
    print(f"  norm collection ({FIFO_LEN} steps)...")
    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])))
    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    # --- Bias FIFO rollout ---
    # Match legacy_env_imit.py line 142-143: [0, -U0, U0]
    ff.apply_ddf()  # restore save_ddf(1)
    bias_arr = np.zeros(n_obj, dtype=DATA_TYPE)
    bias_arr[3] = 0.0
    bias_arr[4] = -1.0 * u0
    bias_arr[5] = 1.0 * u0

    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[0:12])

    # --- save_ddf(2): post-bias checkpoint — consistent with FIFO ---
    ff.get_ddf()
    ff.save_ddf()
    ff.apply_ddf()  # restore to post-bias state for inference

    # Closed-loop inference
    sens_list, actions_list, force_list = [], [], []
    bias_norm = np.array([0.0, 0.125, -0.125], dtype=np.float64)
    a_prev = bias_norm.copy()
    a_prev2 = a_prev.copy()

    # Determine mode from SINDy results metadata (default: v23), unless overridden
    sindy_mode = mode_override if mode_override else coefs.get("mode", "v23")
    if mode_override:
        print(f"  mode overridden: {mode_override}")

    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

        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

        # Reconstruct target forces from harmonics (if available)
        target_forces_step = None
        if target_harmonics is not None:
            target_forces_step = gen_target_states_at(step, target_harmonics)

        if sindy_mode == "pysr":
            # PySR mode -- should not reach here (uses run_validation_pysr_illusion instead)
            raise ValueError("PySR mode should use run_validation_pysr_illusion()")
        elif sindy_mode == "absolute":
            omega_phys = predict_v23_deriv(
                obs, obs_prev, a_prev_phys, a_prev2_phys, mu, u0, sample_interval/2000.0,
                coefs["front_coef"], coefs["top_coef"],
                feat_keys_front=coefs["feat_keys_front"],
                feat_keys_rear=coefs["feat_keys_rear"],
                target_forces=target_forces_step,
                output_mode="absolute")
        else:
            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"],
                target_forces=target_forces_step)

        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.context.push()
        ff.run(sample_interval, action_arr)
        ff.context.pop()

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

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

    # Evaluate similarity over full run
    sens_arr = np.array(sens_list, dtype=np.float32)
    actions_arr = np.array(actions_list, dtype=np.float64)
    # Use last 2*conv_len steps for similarity (after control has propagated)
    tail_sens = sens_arr[-min(2*conv_len, n_steps):]
    tail_target = target_states[-min(2*conv_len, len(target_states)):]
    sim_tail = compute_similarity(tail_target, tail_sens, min(conv_len, len(tail_sens)//2))
    sim_full = compute_similarity(target_states, sens_arr, conv_len)
    action_range = float(np.max(np.abs(actions_arr)))

    print(f"  full-run similarity={sim_full:.4f}  tail similarity={sim_tail:.4f}  action_range={action_range:.4f}")

    # Save controlled vorticity
    try:
        omega_ctrl = vorticity_from_ddf(ff, u0)
        img_path = os.path.join(out_dir, f"{scene_name}_vorticity{threshold_str}.png")
        save_vorticity_png(img_path, omega_ctrl,
                           title=f"{scene_name} SINDy closed-loop (v23, sim={sim_tail:.3f})")
        print(f"  controlled vorticity saved: {img_path}")

        # Uncontrolled: restore from same post-bias DDF, run zero-action for same duration
        ff.restore_ddf()
        ff.apply_ddf()
        # Run bias steps then zero-action for same duration as controlled
        for _ in range(FIFO_LEN):
            ff.run(sample_interval, bias_arr)
        # Reset FIFO then run zero-action for n_steps
        for _ in range(n_steps):
            ff.run(sample_interval, np.zeros(n_obj, dtype=DATA_TYPE))
        omega_unc = vorticity_from_ddf(ff, u0)
        img_unc_path = os.path.join(out_dir, f"{scene_name}_uncontrolled{threshold_str}.png")
        save_vorticity_png(img_unc_path, omega_unc,
                           title=f"{scene_name} uncontrolled")
        print(f"  uncontrolled vorticity saved: {img_unc_path}")
    except Exception as e:
        print(f"  WARNING: controlled/uncontrolled vorticity export failed: {e}")
        import traceback; traceback.print_exc()

    del ff

    # Target cylinder reference (separate env, fully stabilized)
    try:
        ff_tgt = FlowField(field_cfg, cuda_cfg, device_id=device_id)
        tgt_diam = cfg["target_diameter"]
        ff_tgt.add_cylinder((20.0 * l0, (ny - 1) / 2, 0.0), tgt_diam * l0)
        for y_off in [2.0, 0.0, -2.0]:
            sc = (sensor_x * l0, (ny - 1) / 2 + y_off * l0, 0.0)
            ff_tgt.add_sensor(sc, l0 / 4.0)
        n_obj_tgt = ff_tgt.obs.size // 2
        ff_tgt.run(int(5 * NX / u0), np.zeros(n_obj_tgt, dtype=DATA_TYPE))
        omega_tgt = vorticity_from_ddf(ff_tgt, u0)
        img_tgt_path = os.path.join(out_dir, f"{scene_name}_target{threshold_str}.png")
        save_vorticity_png(img_tgt_path, omega_tgt,
                           title=f"{scene_name} target cylinder ({tgt_diam}L)")
        print(f"  target vorticity saved: {img_tgt_path}")
        del ff_tgt
    except Exception as e:
        print(f"  WARNING: target vorticity export failed: {e}")
        import traceback; traceback.print_exc()

    return {
        "scene": scene_name,
        "mode": "v23",
        "similarity_full": sim_full,
        "similarity_tail": sim_tail,
        "action_range": action_range,
        "n_steps": n_steps,
        "threshold": threshold if threshold is not None else "best",
    }


def run_validation_pysr_illusion(
    scene_name: str,
    device_id: int,
    front_func,
    top_func,
    feat_names_front: list,
    feat_names_rear: list,
    n_steps: int = 0,
    out_dir: Optional[str] = None,
) -> dict:
    """Run closed-loop validation using PySR formulas for an Illusion scene."""
    from SR_analysis.validate.predict_pysr import pysr_predict_v23

    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    mu = cfg["mu"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    conv_len = cfg.get("conv_len", 36)
    action_scale = cfg["action_scale"]
    action_bias = cfg["action_bias"]
    n_obj_total = cfg["n_objects_env"]
    sensor_x = cfg["sensor_x"]
    front_x = cfg["pinball_front_x"]
    rear_x = cfg["pinball_rear_x"]

    min_steps = int(1 * NX / u0 / sample_interval)
    if n_steps == 0 or n_steps < min_steps:
        n_steps = max(min_steps, 200)
        print(f"  auto-set steps={n_steps} (min_steps={min_steps})")

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

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

    # Load target data
    data_dir = os.path.join(
        os.path.dirname(__file__), "..", "data", "illusion", scene_name,
    )
    target_npz = np.load(os.path.join(data_dir, "target.npz"))
    target_states = target_npz.get("target_sensors", target_npz["target_states"][:, 2:8])
    print(f"  target loaded: {target_states.shape}")

    target_harmonics = None
    harm_path = os.path.join(data_dir, "target_harmonics.json")
    if os.path.isfile(harm_path):
        with open(harm_path) as f:
            target_harmonics = json.load(f)
        print(f"  target harmonics loaded: {len(target_harmonics)} channels")

    # Build environment
    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]

    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)
    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
    assert n_obj == 6, f"Expected 6 objects, got {n_obj}"

    stabilize_steps = int(4 * NX / u0)
    print(f"  stabilising pinball ({stabilize_steps} steps)...")
    ff.run(stabilize_steps, np.zeros(n_obj, dtype=DATA_TYPE))

    ff.get_ddf()
    ff.save_ddf()

    # Norm + bias FIFO (same as run_validation_illusion)
    print(f"  norm collection ({FIFO_LEN} steps)...")
    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])

    ff.apply_ddf()
    bias_arr = np.zeros(n_obj, dtype=DATA_TYPE)
    bias_arr[3] = 0.0
    bias_arr[4] = -1.0 * u0
    bias_arr[5] = 1.0 * u0
    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[0:12])

    ff.get_ddf()
    ff.save_ddf()
    ff.apply_ddf()

    # Closed-loop with PySR
    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"]
    sens_list, actions_list = [], []
    bias_norm = np.array([0.0, 0.125, -0.125], dtype=np.float64)
    a_prev = bias_norm.copy()
    a_prev2 = a_prev.copy()

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

        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

        target_forces_step = None
        if target_harmonics is not None:
            target_forces_step = gen_target_states_at(step, target_harmonics)

        omega_phys = pysr_predict_v23(
            obs, a_prev_phys, a_prev2_phys, mu, u0, sample_interval / 2000.0,
            front_func, top_func,
            feat_keys_front, feat_keys_rear,
            target_forces=target_forces_step,
        )

        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.context.push()
        ff.run(sample_interval, action_arr)
        ff.context.pop()

        obs_new = ff.obs.copy()[0:12]
        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)

    sens_arr = np.array(sens_list, dtype=np.float32)
    from SR_analysis.utils.cfd_interface import compute_similarity
    sim_full = compute_similarity(target_states, sens_arr, conv_len)
    sim_tail = compute_similarity(target_states[-conv_len:], sens_arr[-conv_len:], conv_len // 2)
    action_range = float(np.max(np.abs(actions_list)))

    print(f"  similarity (full)={sim_full:.4f}  (tail)={sim_tail:.4f}  action_range={action_range:.4f}")
    del ff

    return {
        "scene": scene_name,
        "mode": "pysr",
        "similarity_full": sim_full,
        "similarity_tail": sim_tail,
        "action_range": action_range,
        "n_steps": n_steps,
    }


def main():
    ap = argparse.ArgumentParser(description="Illusion closed-loop SINDy validation")
    ap.add_argument("--scene", type=str, required=True)
    ap.add_argument("--device", type=int, default=0, help="GPU device")
    ap.add_argument("--steps", type=int, default=0,
                    help="Steps (default: auto-set to cover 1*NX/U0)")
    ap.add_argument("--threshold", type=float, default=None)
    ap.add_argument("--sindy-results", type=str, default=None)
    ap.add_argument("--mode", type=str, default=None,
                    help='Override mode: "v23", "absolute", "pysr" (default: auto-detect)')
    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, pysr_predict_v23,
        )
        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_pysr_illusion(
            args.scene, args.device, front_func, top_func, fn_f, fn_r,
            n_steps=args.steps, out_dir=args.out,
        )
    else:
        if args.sindy_results is None:
            args.sindy_results = os.path.join(
                os.path.dirname(__file__), "..", "sindy", "illusion", "sindy_results_v2.json")
        result = run_validation_illusion(
            args.scene, args.sindy_results, args.device,
            n_steps=args.steps, threshold=args.threshold, out_dir=args.out,
            mode_override=args.mode,
        )

    if args.mode == "pysr":
        out_name = f"{args.scene}_pysr"
    else:
        th_str = f"_th{args.threshold}" if args.threshold is not None else ""
        out_name = f"{args.scene}_v23{th_str}"
    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()
