"""Karman re400 closed-loop validation with configurable sample_interval.

Tests whether the existing joint deep PySR formula performs better
at shorter control intervals (SI=400, SI=200) than the default SI=800.

Usage:
    conda run -n pycuda_3_10 python validate/run_closed_loop_re400_si.py \\
        --si 400 --device 2

    conda run -n pycuda_3_10 python validate/run_closed_loop_re400_si.py \\
        --si 200 --device 2
"""
from __future__ import annotations

import argparse
import json
import os
import sys

_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 SR_analysis.validate.run_closed_loop import run_validation, load_sindy_coefs
from SR_analysis.validate.predict_pysr import load_pysr_formulas, build_pysr_functions
from SR_analysis.configs import SCENES


def main():
    ap = argparse.ArgumentParser(description="Karman re400 short-SI validation")
    ap.add_argument("--si", type=int, default=400, help="Sample interval (default: 400)")
    ap.add_argument("--device", type=int, default=2, help="GPU device")
    ap.add_argument("--steps", type=int, default=0,
                    help="Steps (0=auto: max(NX/u0/SI, 200))")
    args = ap.parse_args()

    scene = "karman_re400"
    si = args.si
    device = args.device
    n_steps = args.steps

    results_dir = os.path.join(os.path.dirname(__file__), "results")
    pysr_front = os.path.join(results_dir, "karman_joint_deep_front.json")
    pysr_top = os.path.join(results_dir, "karman_joint_deep_top.json")

    # Override sample_interval in scene config
    original_si = SCENES[scene]["sample_interval"]
    SCENES[scene]["sample_interval"] = si
    print(f"  Override: sample_interval {original_si} -> {si}")

    # Load PySR formulas
    front_expr, top_expr, fn_f, fn_r = load_pysr_formulas(pysr_front, pysr_top)
    front_func, top_func = build_pysr_functions(front_expr, top_expr, fn_f, fn_r)

    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,
    }

    result = run_validation(scene, coefs, device, n_steps=n_steps, mode="pysr")

    # Save with SI in name
    out_dir = os.path.join(os.path.dirname(__file__), "results")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"karman_re400_pysr_SI{si}.json")
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"Saved: {out_path}")

    # Restore original config
    SCENES[scene]["sample_interval"] = original_si


if __name__ == "__main__":
    main()
