"""Unified closed-loop validator for SINDy control laws.

Validates a SINDy-derived control law by running it in a closed-loop CFD
environment and measuring similarity to the target flow.

Two modes:
  - v23 (default): front no-bias + rear shared-head [bottom=-top(Gx)]
  - unstructured: front with bias + rear independent

Usage:
    conda run -n pycuda_3_10 python validate/run_closed_loop.py \\
        --scene karman_re70 --device 2 \\
        --sindy-results sindy/karman/sindy_results.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from collections import deque
from typing import Any, Dict, List, Optional, Tuple

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 (
    nu_from_re, load_legacy_configs, build_karman_cloak_env, add_pinball,
    scale_action, action_to_physical, compute_similarity,
    load_ppo_model,
)
from SR_analysis.utils.sindy_fitter import get_feature_matrix_from_data
from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    apply_G_x, ALL_FEAT_KEYS, ALL_FEAT_KEYS_V2, CORE_FEAT_KEYS_V2, U0,
)
from SR_analysis.utils.g_operator import apply_G_raw
from SR_analysis.configs import (
    get_scene, SCENES, LEGACY_CFG_DIR, FIFO_LEN, CONV_LEN,
)

DATA_TYPE = np.float32


def build_feature_vector(
    obs_slice: np.ndarray,
    a_prev: np.ndarray,
    a_prev2: np.ndarray,
    mu: float,
    u0: float,
    add_bias: bool,
    feat_keys: Optional[List[str]] = None,
    use_v2: bool = False,
    target_forces: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Build a single-row feature vector from raw obs and action state.

    Matches the feature_builder logic but for a single time step.

    Args:
        feat_keys: custom feature keys. If None, uses ALL_FEAT_KEYS_V2 or ALL_FEAT_KEYS.
        use_v2: if True, use CORE_FEAT_KEYS_V2 (no sin/cos). Ignored if feat_keys given.
        target_forces: (2,) raw lattice target forces [fx,fy], for Illusion scenes.
    """
    sensors = obs_slice[0:6].astype(np.float64).reshape(1, 6)
    forces = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    ap = a_prev.astype(np.float64).reshape(1, 3)
    ap2 = a_prev2.astype(np.float64).reshape(1, 3)

    dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
    # Detect whether to include mu terms based on feat_keys
    has_mu = feat_keys is not None and "mu" in feat_keys
    has_cos_sin = feat_keys is not None and "sin_ua" in feat_keys
    tf = target_forces.reshape(1, 2) if target_forces is not None else None
    sym = compute_features(dim, ap, ap2, mu, alpha_mode=False,
                           include_mu=has_mu,
                           include_cos_sin=has_cos_sin,
                           u0=u0, target_forces=tf)

    if feat_keys is None:
        if use_v2:
            feat_keys = CORE_FEAT_KEYS_V2
        else:
            feat_keys = ALL_FEAT_KEYS

    feat = build_feature_matrix(sym, feat_keys, add_bias=add_bias)
    return feat[0]  # single row


def predict_v23(
    obs_slice: np.ndarray,
    a_prev: np.ndarray,
    a_prev2: np.ndarray,
    mu: float,
    u0: float,
    front_coef: np.ndarray,
    top_coef: np.ndarray,
    feat_names_front: List[str],
    feat_names_rear: List[str],
    feat_keys_front: Optional[List[str]] = None,
    feat_keys_rear: Optional[List[str]] = None,
    target_forces: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Predict actions using v23: front no-bias + rear shared-head.

    Returns (3,) physical omega array: [front, bottom, top].
    """
    # Front channel: no bias
    fv_front = build_feature_vector(
        obs_slice, a_prev, a_prev2, mu, u0, add_bias=False,
        feat_keys=feat_keys_front,
        target_forces=target_forces,
    )
    front = float(np.dot(fv_front, front_coef))

    # Top channel: with bias
    fv_top = build_feature_vector(
        obs_slice, a_prev, a_prev2, mu, u0, add_bias=True,
        feat_keys=feat_keys_rear,
        target_forces=target_forces,
    )
    top = float(np.dot(fv_top, top_coef))

    # Bottom = -top(Gx) using shared-head
    G_obs, G_a_prev, G_a_prev2 = apply_G_raw(obs_slice, a_prev, a_prev2)
    # G transforms forces: swap bottom<->top, negate fy. Target force G transform:
    # target_Cd is a drag-like quantity that is even under G (stays),
    # target_Cl is a lift-like quantity that is odd under G (negates).
    G_target = np.array([target_forces[0], -target_forces[1]]) if target_forces is not None else None
    fv_bot = build_feature_vector(
        G_obs, G_a_prev, G_a_prev2, mu, u0, add_bias=True,
        feat_keys=feat_keys_rear,
        target_forces=G_target,
    )
    bottom = -float(np.dot(fv_bot, top_coef))

    return np.array([front, bottom, top], dtype=np.float64)


def predict_unstructured(
    obs_slice: np.ndarray,
    a_prev: np.ndarray,
    a_prev2: np.ndarray,
    mu: float,
    u0: float,
    front_coef: np.ndarray,
    bottom_coef: np.ndarray,
    top_coef: np.ndarray,
    feat_names: List[str],
    feat_keys: Optional[List[str]] = None,
) -> np.ndarray:
    """Predict actions using unstructured: each channel independent with bias."""
    feat = build_feature_vector(
        obs_slice, a_prev, a_prev2, mu, u0, add_bias=True,
        feat_keys=feat_keys,
    )
    front = float(np.dot(feat, front_coef))
    bottom = float(np.dot(feat, bottom_coef))
    bottom = float(np.dot(feat, bottom_coef))
    top = float(np.dot(feat, top_coef))
    return np.array([front, bottom, top], dtype=np.float64)


def predict_v23_deriv(
    obs_slice: np.ndarray,
    obs_prev: np.ndarray,
    a_prev_phys: np.ndarray,
    a_prev2_phys: np.ndarray,
    mu: float,
    u0: float,
    dt_c: float,
    front_coef: np.ndarray,
    top_coef: np.ndarray,
    feat_keys_front: Optional[List[str]] = None,
    feat_keys_rear: Optional[List[str]] = None,
    target_forces: Optional[np.ndarray] = None,
    output_mode: str = "deriv",
) -> np.ndarray:
    """Predict action using phase-state features, v23 structure.
    
    Two output modes:
        "deriv": Predict d(alpha)/dt, then integrate:
                 omega_new = (alpha_prev + dt_c * d_alpha/dt_pred) * u0
        "absolute": Predict alpha directly:
                 omega_new = alpha_pred * u0
    
    Returns (3,) physical omega array: [front, bottom, top].
    """
    from SR_analysis.utils.feature_builder import PHYSICS_FEAT_KEYS
    
    if feat_keys_front is None:
        feat_keys_front = PHYSICS_FEAT_KEYS
    if feat_keys_rear is None:
        feat_keys_rear = PHYSICS_FEAT_KEYS
    
    # Check if augmented features (derivatives or lags) are needed
    needs_aug = any(k.endswith("_dt") or k.endswith("_lag1") for k in feat_keys_front)
    
    # Build sensor/force arrays: if augmented features needed, stack prev+current
    if needs_aug:
        sensors_raw = np.vstack([obs_prev[0:6], obs_slice[0:6]]).astype(np.float64)
        forces_raw = np.vstack([obs_prev[6:12], obs_slice[6:12]]).astype(np.float64)
        # Current dim uses only the last row
        sensors_curr = obs_slice[0:6].astype(np.float64).reshape(1, 6)
        forces_curr = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    else:
        sensors_raw = sensors_curr = obs_slice[0:6].astype(np.float64).reshape(1, 6)
        forces_raw = forces_curr = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    
    ap = a_prev_phys.astype(np.float64).reshape(1, 3)
    ap2 = a_prev2_phys.astype(np.float64).reshape(1, 3)
    
    dim = compute_dimensionless(sensors_curr, forces_curr, u0=u0, d=20.0)
    has_mu = any("mu" in k for k in feat_keys_front)
    tf = target_forces.reshape(1, 2) if target_forces is not None else None
    sym = compute_features(dim, ap, ap2, mu, alpha_mode=False,
                           include_mu=has_mu,
                           include_cos_sin=False,
                           u0=u0, target_forces=tf,
                           sensors_raw=sensors_raw, forces_raw=forces_raw)
    
    # Front: no bias, predict d(alpha_F)/dt
    fv_front = build_feature_matrix(sym, feat_keys_front, add_bias=False)[0]
    pred_F = float(np.dot(fv_front, front_coef))
    
    # Top: with bias, predict d(alpha_T)/dt
    fv_top = build_feature_matrix(sym, feat_keys_rear, add_bias=True)[0]
    pred_T = float(np.dot(fv_top, top_coef))
    
    # Bottom = -top(Gx) using shared-head
    G_obs, G_a_prev, G_a_prev2 = apply_G_raw(obs_slice, a_prev_phys, a_prev2_phys)
    G_target = np.array([target_forces[0], -target_forces[1]]) if target_forces is not None else None
    G_sensors = G_obs[0:6].astype(np.float64).reshape(1, 6)
    G_forces = G_obs[6:12].astype(np.float64).reshape(1, 6)
    G_dim = compute_dimensionless(G_sensors, G_forces, u0=u0, d=20.0)
    G_sym = compute_features(G_dim, G_a_prev.reshape(1, 3), G_a_prev2.reshape(1, 3),
                              mu, alpha_mode=False, include_mu=has_mu,
                              include_cos_sin=False, u0=u0, target_forces=G_target)
    fv_bot = build_feature_matrix(G_sym, feat_keys_rear, add_bias=True)[0]
    pred_B = -float(np.dot(fv_bot, top_coef))
    
    # Convert to omega based on output_mode
    if output_mode == "absolute":
        # Direct prediction: pred = alpha, omega = alpha * u0
        alpha_new = np.array([pred_F, pred_B, pred_T])
    else:
        # Derivative + integrate: pred = d(alpha)/dt
        # omega = (alpha_prev + dt_c * d_alpha/dt) * u0
        alpha_prev = a_prev_phys / u0
        alpha_new = alpha_prev + dt_c * np.array([pred_F, pred_B, pred_T])
    omega_new = alpha_new * u0
    
    return omega_new.astype(np.float64)


def load_sindy_coefs(sindy_path: str, scene_name: str,
                     threshold: Optional[float] = None) -> Dict[str, Any]:
    """Load SINDy coefficients for a scene from results JSON.

    If threshold is given (e.g. 0.007), uses coefficients at that threshold
    instead of the best_R2 ones.

    Auto-detects v1 (21 features, sin_ua/cos_ua, mu) vs v2 (14 features, no sin/cos)
    feature sets based on the length of feature_names_front.

    Returns dict with keys: front_coef, top_coef, bottom_coef,
    feat_names_front, feat_names_rear, front_bias_mode,
    feat_keys_front, feat_keys_rear, is_v2.
    """
    with open(sindy_path) as f:
        data = json.load(f)

    per = data["per_scene"].get(scene_name)
    if per is None:
        raise KeyError(f"Scene {scene_name} not found in {sindy_path}")

    fn_f = per["feature_names_front"]
    fn_r = per["feature_names_rear"]

    def _coef_at_threshold(ch_data, th):
        """Get coef at given threshold, or best if th=None."""
        if th is None:
            return np.array(ch_data["best_coef"], dtype=np.float64)
        for r in ch_data.get("results", []):
            if abs(r["threshold"] - th) < 1e-6:
                return np.array(r.get("coef", ch_data["best_coef"]), dtype=np.float64)
        # Fallback: best
        return np.array(ch_data["best_coef"], dtype=np.float64)

    front_coef = _coef_at_threshold(per["front"], threshold)
    top_coef = _coef_at_threshold(per["top"], threshold)
    bottom_coef = _coef_at_threshold(per["bottom"], threshold)

    # Trim coefficients to match feature names
    front_coef = front_coef[:len(fn_f)]
    top_coef = top_coef[:len(fn_r)]
    bottom_coef = bottom_coef[:len(fn_r)]

    # Detect v1 vs v2: v2 has 14 features (no sin_ua/cos_ua), v1 has 16
    # v2 joint has 19 (14 core + 5 mu: mu, mu_u_a, mu_v_a, mu_Cd_tot, mu_Cl_diff)
    # v1 joint has 21 (16 core + 5 mu)
    # Illusion v2 has 16 (14 core + 2 target: target_Cd, target_Cl)
    # Illusion v2 + mu has 21 (14 core + 2 target + 5 mu)
    # Deriv mode: PHYSICS_FEAT_KEYS has 8, ILLUSION_ERR_KEYS has 10
    is_v2 = len(fn_f) in [8, 10, 14, 16, 19, 21]

    # Strip "bias" from feature keys (build_feature_matrix adds it via add_bias flag)
    feat_keys_front = [k for k in fn_f if k != "bias"]
    feat_keys_rear = [k for k in fn_r if k != "bias"]

    # Detect if front was fitted with bias or without
    front_has_bias = fn_f[0] == "bias" if len(fn_f) > 0 else False

    # Extract mode from per_scene metadata (default: "v23" for old-style results)
    sindy_mode = per.get("mode", "v23")

    return {
        "front_coef": front_coef,
        "top_coef": top_coef,
        "bottom_coef": bottom_coef,
        "feat_names_front": fn_f,
        "feat_names_rear": fn_r,
        "feat_keys_front": feat_keys_front,
        "feat_keys_rear": feat_keys_rear,
        "front_has_bias": front_has_bias,
        "is_v2": is_v2,
        "mode": sindy_mode,
        "threshold_used": threshold if threshold is not None else "best",
    }


def run_validation(
    scene_name: str,
    coefs: Dict[str, Any],
    device_id: int,
    n_steps: int = 0,
    mode: str = "v23",
) -> dict:
    """Run closed-loop validation using a SINDy control law.

    Parameters
    ----------
    n_steps : int, default=0 (auto: >= NX/U0 / sample_interval)
    """
    cfg = get_scene(scene_name)
    re_code = cfg["re_code"]
    nu = cfg["nu"]
    u0 = cfg["u0"]
    mu = cfg["mu"]
    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"]
    t0_steps = 2000
    dt_c = sample_interval / t0_steps

    # Auto-set steps: control must propagate at least 1*NX/U0 LBM steps
    NX = 1280
    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={mode}, device={device_id}, steps={n_steps}) ===")

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

    # Record target, add pinball, compute norm
    target_states, _ = build_karman_cloak_env(
        ff, u0=u0, l0=l0, sample_interval=sample_interval,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
    )
    norm = add_pinball(
        ff, l0=l0, u0=u0, sample_interval=sample_interval,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
        action_bias=action_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],
    )

    # Now ff is at save_ddf(2) = post-bias state (add_pinball saves after bias rollout).
    # Norm dict contains save_states from the bias FIFO.
    fifo = deque(maxlen=FIFO_LEN)
    for row in norm["save_states"]:
        fifo.append(row)

    # Closed-loop with SINDy law
    sens_list = []
    actions_list = []
    # For Karman: bias action = DRL action_bias = [0, -4, 4].
    # Normalized zero action [0,0,0] decodes to the same bias.
    a_prev = np.zeros(3, dtype=np.float64)
    a_prev2 = a_prev.copy()

    for _ 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  # for augmented features

        # Convert a_prev from normalized to physical omega for feature builder
        # (SINDy was trained on physical omega Y = actions_phys)
        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

        # Predict using SINDy law -- returns physical omega
        if 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"])
        elif mode == "deriv":
            omega_phys = predict_v23_deriv(
                obs, obs_prev, a_prev_phys, a_prev2_phys, mu, u0, dt_c,
                coefs["front_coef"], coefs["top_coef"],
                feat_keys_front=coefs["feat_keys_front"],
                feat_keys_rear=coefs["feat_keys_rear"],
                output_mode="deriv")
        elif mode == "abs":
            omega_phys = predict_v23_deriv(
                obs, obs_prev, a_prev_phys, a_prev2_phys, mu, u0, dt_c,
                coefs["front_coef"], coefs["top_coef"],
                feat_keys_front=coefs["feat_keys_front"],
                feat_keys_rear=coefs["feat_keys_rear"],
                output_mode="absolute")
        elif mode == "pysr":
            from SR_analysis.validate.predict_pysr import pysr_predict_v23
            omega_phys = pysr_predict_v23(
                obs, a_prev_phys, a_prev2_phys, mu, u0, dt_c,
                coefs["front_func"], coefs["top_func"],
                feat_keys_front=coefs["feat_keys_front"],
                feat_keys_rear=coefs["feat_keys_rear"])
        elif mode == "unstructured":
            omega_phys = predict_unstructured(
                obs, a_prev_phys, a_prev2_phys, mu, u0,
                coefs["front_coef"], coefs["bottom_coef"], coefs["top_coef"],
                coefs["feat_names_rear"],
                feat_keys=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()[2:14]
        fifo.append(obs_new)
        sens_list.append(obs_new[0:6])
        actions_list.append(omega_phys.copy())

        # Keep a_prev in normalized form for next iteration's memory features
        a_prev2 = a_prev.copy()
        a_prev = norm_a.copy().astype(np.float64)

    # Evaluate
    sens_arr = np.array(sens_list, dtype=np.float32)
    actions_arr = np.array(actions_list, dtype=np.float64)
    sim = compute_similarity(target_states, sens_arr, CONV_LEN)
    action_range = float(np.max(np.abs(actions_arr)))

    print(f"  similarity={sim:.4f}  action_range={action_range:.4f}")

    del ff

    return {
        "scene": scene_name,
        "mode": mode,
        "similarity": sim,
        "action_range": action_range,
        "n_steps": n_steps,
    }


def main():
    ap = argparse.ArgumentParser(description="Closed-loop SINDy validation")
    ap.add_argument("--scene", type=str, required=True, help="Scene name")
    ap.add_argument("--device", type=int, default=2, help="GPU device")
    ap.add_argument("--steps", type=int, default=100)
    ap.add_argument("--mode", type=str, default="v23",
                    choices=["v23", "unstructured", "deriv", "abs", "pysr"],
                    help="Control law mode: v23 (default), unstructured, deriv, abs, or pysr")
    ap.add_argument("--sindy-results", type=str, default=None,
                    help="Path to sindy_results.json (not used in pysr mode)")
    ap.add_argument("--threshold", type=float, default=None,
                    help="SINDy threshold for sparsity (default: best_R2)")
    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,
                    help="Output directory for result JSON")
    args = ap.parse_args()

    if args.mode == "pysr":
        # Load PySR formulas directly
        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)
        # Package as pseudo-coefs dict for run_validation
        feat_keys_front = [k for k in fn_f if k != "bias"]
        feat_keys_rear = [k for k in fn_r if k != "bias"]
        coefs = {
            "mode": "pysr",
            "front_func": front_func,
            "top_func": top_func,
            "feat_keys_front": feat_keys_front,
            "feat_keys_rear": feat_keys_rear,
            "feat_names_front": fn_f,
            "feat_names_rear": fn_r,
        }
    else:
        if args.sindy_results is None:
            args.sindy_results = os.path.join(
                os.path.dirname(__file__), "..", "sindy", "karman", "sindy_results_v2.json")
        coefs = load_sindy_coefs(args.sindy_results, args.scene, threshold=args.threshold)

    result = run_validation(args.scene, coefs, args.device,
                            n_steps=args.steps, mode=args.mode)

    th_str = f"_th{args.threshold}" if args.threshold is not None else ""
    if args.out is None:
        out_dir = os.path.join(os.path.dirname(__file__), "..",
                               "validate", "results")
    else:
        out_dir = args.out
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{args.scene}_{args.mode}{th_str}.json")
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Saved: {out_path}")


if __name__ == "__main__":
    main()
