"""Pareto analysis of SINDy threshold grid for any scene.

Loads sindy_results.json and prints Pareto-optimal tradeoffs.

Usage:
    python sindy/run_pareto.py --scene karman_re100
    python sindy/run_pareto.py --scene illusion_1L
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import List, 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)


def load_sindy(scene_name: str, sindy_dir: str) -> dict:
    """Load sindy_results.json for a scene."""
    path = os.path.join(sindy_dir, scene_name, "sindy_results.json")
    if os.path.isfile(path):
        return json.load(path)
    raise FileNotFoundError(f"Missing {path}")


def pareto(points: List[Tuple[float, float]]) -> List[Tuple[float, float]]:
    """Compute Pareto frontier: lower nz and lower error is better."""
    sp = sorted(points, key=lambda x: (x[0], x[1]))
    front, best = [], float("inf")
    for c, e in sp:
        if e < best:
            front.append((c, e))
            best = e
    return front


def fmt(fn: List[str], coef: List[float], threshold: float) -> str:
    """Format a control law string, showing terms above relative threshold."""
    ca = np.array(coef, dtype=np.float64)
    sc = np.max(np.abs(ca)) if np.max(np.abs(ca)) > 0 else 1.0
    mask = np.abs(ca) / sc >= threshold
    terms = [f"{ca[i]:+.4f}*{fn[i]}" for i in range(len(fn)) if mask[i]]
    return " ".join(terms) if terms else "0"


def analyze(name: str, feat_names: List[str], channel_data: dict) -> dict:
    """Print and return Pareto analysis for one channel."""
    grid = channel_data["results"]
    pts = [(g["nz"], 1.0 - g["r2"]) for g in grid]
    front = pareto(pts)
    best = channel_data["best"]
    coef = channel_data["best_coef"]

    print(f"\n  {name}:")
    for nz, err in front:
        r2 = 1.0 - err
        for g in grid:
            if g["nz"] == nz and abs(1.0 - g["r2"] - err) < 1e-10:
                th = g["threshold"]
                print(f"    nz={nz:2d} R2={r2:.6f} th={th:.4f}")
                if nz <= 8:
                    s = fmt(feat_names, coef, th)
                    print(f"      {s[:120]}")
    print(f"  Best: R2={best['r2']:.6f}")

    return {
        "channel": name,
        "best_r2": best["r2"],
        "best_nz": sum(1 for c in coef if abs(float(c)) > 1e-8),
        "pareto": [{"nz": nz, "r2": 1.0 - e} for nz, e in front],
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, required=True,
                    help="Scene name (e.g. karman_re100, illusion_1L)")
    ap.add_argument("--sindy-dir", type=str, default=None,
                    help="SINDy results directory")
    ap.add_argument("--out", type=str, default=None,
                    help="Output JSON path")
    args = ap.parse_args()

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

    # Map scene name to subdirectory
    # Extract series prefix: karman_* -> karman, illusion_* -> illusion, etc.
    first_part = args.scene.split("_")[0]
    known_series = {"karman": "karman", "illusion": "illusion", "vortex": "vortex", "steady": "steady"}
    series_dir = known_series.get(first_part, first_part)

    # Try sindy/{series}/sindy_results.json
    json_path = os.path.join(args.sindy_dir, series_dir, "sindy_results.json")

    if not os.path.isfile(json_path):
        # Fallback: try flat file
        json_path = os.path.join(args.sindy_dir, "sindy_results.json")
        if not os.path.isfile(json_path):
            print(f"ERROR: No sindy results found for {args.scene}")
            print(f"  Tried: {os.path.join(args.sindy_dir, series_dir, 'sindy_results.json')}")
            print(f"  Tried: {json_path}")
            return 1

    with open(json_path) as f:
        data = json.load(f)

    # Look up the scene in per_scene (multi-scene format)
    per = data.get("per_scene", {}).get(args.scene)
    if per is not None:
        fn_f = per["feature_names_front"]
        fn_r = per["feature_names_rear"]
        chs = [("front", fn_f, per["front"]),
               ("top", fn_r, per["top"]),
               ("bottom", fn_r, per["bottom"])]
    else:
        # Single-scene format
        fn_f = data.get("feature_names_front")
        fn_r = data.get("feature_names_rear")
        if fn_f is None:
            print(f"ERROR: No scene data found for {args.scene} in {json_path}")
            return 1
        chs = [("front", fn_f, data["front"]),
               ("top", fn_r, data["top"]),
               ("bottom", fn_r, data["bottom"])]

    print(f"Pareto SR: {args.scene}")
    results = {"scene": args.scene,
               "channels": [analyze(*c) for c in chs]}

    if args.out:
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        with open(args.out, "w") as f:
            json.dump(results, f, indent=2)
        print(f"Saved: {args.out}")


if __name__ == "__main__":
    main()
