"""SINDy fitting for Karman cloak scenes.

Runs STLSQ threshold grid for Karman scenes (training Re or all Re).

Usage:
    conda run -n pycuda_3_10 python sindy/run_karman.py --re-codes 50,100,200
    conda run -n pycuda_3_10 python sindy/run_karman.py
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import List, 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 SR_analysis.utils.sindy_fitter import fit_sindy, get_feature_matrix_from_data
from SR_analysis.utils.feature_builder import ALL_FEAT_KEYS, U0
from SR_analysis.configs import get_scene, get_scene_list, SCENES


SINDY_DIR = os.path.join(os.path.dirname(__file__), "..", "sindy", "karman")
THRESHOLDS = [0.0, 0.001, 0.002, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1]


def load_data(scene_name: str, scene_subdir: str = "karman") -> tuple:
    """Load sensors/forces/actions from a scene's controlled.npz."""
    data_dir = os.path.join(os.path.dirname(__file__), "..", "data", scene_subdir, scene_name)
    npz = np.load(os.path.join(data_dir, "controlled.npz"))
    sensors = npz["sensors"].astype(np.float64)
    forces = npz["forces"].astype(np.float64)
    actions = npz["actions"].astype(np.float64)
    return sensors, forces, actions


def run(scene_names: Optional[List[str]] = None):
    """Run SINDy fitting for given scene names."""
    if scene_names is None:
        scene_names = get_scene_list("karman")

    per_scene = {}

    for sn in scene_names:
        print(f"\n{'='*60}")
        print(f"Scene: {sn}")
        print(f"{'='*60}")

        cfg = get_scene(sn)
        sensors, forces, actions_phys = load_data(sn, cfg["scene_id"])
        T = sensors.shape[0]
        mu = cfg["mu"]
        print(f"  T={T}, mu={mu:.6f}")

        # Build feature matrices
        Theta_f, Theta_r, Y, fn_f, fn_r = get_feature_matrix_from_data(
            sensors, forces, actions_phys, mu, u0=cfg["u0"],
            alpha_mode=False, include_mu=True, n_warmup=2,
        )
        print(f"  Front features: {Theta_f.shape}")
        print(f"  Rear features:  {Theta_r.shape}")
        print(f"  Y:              {Y.shape}")

        # Front channel (ci=0, no bias)
        print(f"\n  --- Front (no bias) ---")
        front_results = fit_sindy(Theta_f, Y[:, 0], THRESHOLDS)
        best_f = max(front_results, key=lambda r: r["r2"])
        print(f"    Best: th={best_f['threshold']:.4f} nz={best_f['nz']:2d} R2={best_f['r2']:.6f}")

        # Top channel (ci=2, rear shared-head)
        print(f"\n  --- Top (rear shared-head) ---")
        top_results = fit_sindy(Theta_r, Y[:, 2], THRESHOLDS)
        best_t = max(top_results, key=lambda r: r["r2"])
        print(f"    Best: th={best_t['threshold']:.4f} nz={best_t['nz']:2d} R2={best_t['r2']:.6f}")

        # Bottom (independent, for comparison)
        print(f"\n  --- Bottom (independent) ---")
        bot_results = fit_sindy(Theta_r, Y[:, 1], THRESHOLDS)
        best_b = max(bot_results, key=lambda r: r["r2"])
        print(f"    Best: th={best_b['threshold']:.4f} nz={best_b['nz']:2d} R2={best_b['r2']:.6f}")

        per_scene[sn] = {
            "scene": sn,
            "re_code": cfg["re_code"],
            "mu": mu,
            "n_samples": Theta_f.shape[0],
            "n_features_front": Theta_f.shape[1],
            "n_features_rear": Theta_r.shape[1],
            "feature_names_front": fn_f,
            "feature_names_rear": fn_r,
            "front": {
                "results": [{k: v for k, v in r.items() if k != "coef"}
                            for r in front_results],
                "best": {k: v for k, v in best_f.items() if k != "coef"},
                "best_coef": best_f["coef"],
                "sparsity_curve": [(r["threshold"], r["nz"], r["r2"])
                                   for r in front_results],
            },
            "top": {
                "results": [{k: v for k, v in r.items() if k != "coef"}
                            for r in top_results],
                "best": {k: v for k, v in best_t.items() if k != "coef"},
                "best_coef": best_t["coef"],
                "sparsity_curve": [(r["threshold"], r["nz"], r["r2"])
                                   for r in top_results],
            },
            "bottom": {
                "results": [{k: v for k, v in r.items() if k != "coef"}
                            for r in bot_results],
                "best": {k: v for k, v in best_b.items() if k != "coef"},
                "best_coef": best_b["coef"],
                "sparsity_curve": [(r["threshold"], r["nz"], r["r2"])
                                   for r in bot_results],
            },
        }

    # Save
    os.makedirs(SINDY_DIR, exist_ok=True)
    out_path = os.path.join(SINDY_DIR, "sindy_results.json")
    result = {
        "thresholds": THRESHOLDS,
        "all_feature_names_front": fn_f,
        "all_feature_names_rear": fn_r,
        "per_scene": per_scene,
    }
    with open(out_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"\nSaved: {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--re-codes", type=str, default=None,
                    help="Comma-separated Re codes (default: all karman)")
    ap.add_argument("--scene-names", type=str, default=None,
                    help="Comma-separated scene names (overrides --re-codes)")
    args = ap.parse_args()

    if args.scene_names:
        names = [s.strip() for s in args.scene_names.split(",")]
    elif args.re_codes:
        codes = [int(r) for r in args.re_codes.split(",")]
        names = [f"karman_re{rc}" for rc in codes]
    else:
        names = None  # all karman

    run(names)


if __name__ == "__main__":
    main()
