"""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, 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,
) -> 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.
    """
    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)
    sym = compute_features(dim, ap, ap2, mu, alpha_mode=False, include_mu=True, u0=u0)

    feat = build_feature_matrix(sym, ALL_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],
) -> np.ndarray:
    """Predict actions using v23: front no-bias + rear shared-head.

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

    # Top channel: with bias
    top = float(np.dot(
        build_feature_vector(obs_slice, a_prev, a_prev2, mu, u0, add_bias=True),
        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)
    bottom = -float(np.dot(
        build_feature_vector(G_obs, G_a_prev, G_a_prev2, mu, u0, add_bias=True),
        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],
) -> 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)
    front = float(np.dot(feat, front_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 load_sindy_coefs(sindy_path: str, scene_name: str) -> Dict[str, Any]:
    """Load SINDy coefficients for a scene from results JSON.

    Returns dict with keys: front_coef, top_coef, bottom_coef,
    feat_names_front, feat_names_rear, front_bias_mode.
    """
    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"]

    front_coef = np.array(per["front"]["best_coef"], dtype=np.float64)
    top_coef = np.array(per["top"]["best_coef"], dtype=np.float64)
    bottom_coef = np.array(per["bottom"]["best_coef"], dtype=np.float64)

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

    return {
        "front_coef": front_coef,
        "top_coef": top_coef,
        "bottom_coef": bottom_coef,
        "feat_names_front": fn_f,
        "feat_names_rear": fn_r,
        "front_has_bias": front_has_bias,
    }


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

    Parameters
    ----------
    scene_name : e.g. "karman_re70"
    coefs : dict from load_sindy_coefs()
    device_id : GPU device
    n_steps : number of closed-loop steps
    mode : "v23" (rear shared-head) or "unstructured"

    Returns dict with similarity, actions range, etc.
    """
    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"]

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

    # 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],
    )

    # Reset to checkpoint
    ff.restore_ddf()
    ff.apply_ddf()

    # Bias FIFO init
    fifo = deque(maxlen=FIFO_LEN)
    bias_arr = scale_action(np.zeros(3, dtype=np.float32), scale=action_scale,
                            bias=action_bias, u0=u0, n_total_bodies=n_obj_total)
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[2:14])

    # Closed-loop with SINDy law
    sens_list = []
    actions_list = []
    a_prev = action_to_physical(np.zeros((1, 3), dtype=np.float32),
                                scale=action_scale, bias=action_bias, u0=u0).flatten()
    a_prev2 = a_prev.copy()

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

        # Predict using SINDy law
        if mode == "v23":
            omega = predict_v23(
                obs, a_prev, a_prev2, mu, u0,
                coefs["front_coef"], coefs["top_coef"],
                coefs["feat_names_front"], coefs["feat_names_rear"])
        elif mode == "unstructured":
            omega = predict_unstructured(
                obs, a_prev, a_prev2, mu, u0,
                coefs["front_coef"], coefs["bottom_coef"], coefs["top_coef"],
                coefs["feat_names_rear"])
        else:
            raise ValueError(f"Unknown mode: {mode}")

        # Clip to valid action range
        norm_a = (omega / u0 - np.array(action_bias, dtype=np.float64)) / action_scale
        norm_a = np.clip(norm_a, -1.0, 1.0).astype(np.float32)

        # Apply to CFD
        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.copy())

        a_prev2 = a_prev.copy()
        a_prev = omega.copy()

    # 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"])
    ap.add_argument("--sindy-results", type=str, default=None,
                    help="Path to sindy_results.json")
    ap.add_argument("--out", type=str, default=None,
                    help="Output directory for result JSON")
    args = ap.parse_args()

    if args.sindy_results is None:
        args.sindy_results = os.path.join(
            os.path.dirname(__file__), "..", "sindy", "karman", "sindy_results.json")

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

    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}.json")
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Saved: {out_path}")


if __name__ == "__main__":
    main()
