"""SINDy fitting for Illusion scenes.

Usage:
    conda run -n pycuda_3_10 python sindy/run_illusion.py
    conda run -n pycuda_3_10 python sindy/run_illusion.py --diameters 0.75,1.0
"""
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.configs import get_scene, get_scene_list

SINDY_DIR = os.path.join(os.path.dirname(__file__), "..", "sindy", "illusion")
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) -> tuple:
    data_dir = os.path.join(os.path.dirname(__file__), "..", "data", "illusion", 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):
    if scene_names is None:
        scene_names = get_scene_list("illusion")

    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)
        mu = cfg["mu"]
        print(f"  T={sensors.shape[0]}, mu={mu:.6f}")

        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: {Theta_f.shape}, Rear: {Theta_r.shape}")

        # Front channel
        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 (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)
        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],
            "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],
            },
        }

    os.makedirs(SINDY_DIR, exist_ok=True)
    out_path = os.path.join(SINDY_DIR, "sindy_results.json")
    result = {"thresholds": THRESHOLDS, "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("--diameters", type=str, default=None,
                    help="Comma-separated diameters (e.g. 0.75,1.0,1.5)")
    ap.add_argument("--scene-names", type=str, default=None)
    args = ap.parse_args()

    if args.scene_names:
        names = [s.strip() for s in args.scene_names.split(",")]
    elif args.diameters:
        names = [f"illusion_{d.strip()}L" for d in args.diameters.split(",")]
    else:
        names = None

    run(names)


if __name__ == "__main__":
    main()