"""Cross-scene support overlap analysis.

Compares SINDy support sets across scenes at a given relative threshold.

Usage:
    python compare/support_overlap.py --sindy-results sindy/karman/sindy_results.json \\
        --scenes karman_re100 karman_re200 illusion_1L vortex_lamb
"""
from __future__ import annotations

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

from SR_analysis.utils.sindy_fitter import get_active_support

RELATIVE_THRESHOLD = 0.02  # default: 2% of max coefficient


def load_sindy_scenes(sindy_path: str, scenes: List[str]) -> dict:
    """Load sindy results for the given scene names."""
    with open(sindy_path) as f:
        data = json.load(f)

    result = {}
    for sn in scenes:
        per = data["per_scene"].get(sn)
        if per is None:
            print(f"WARNING: {sn} not found in {sindy_path}")
            continue
        result[sn] = per
    return result


def feat_group(name: str) -> str:
    """Classify a feature into a group."""
    if name == "bias":
        return "bias"
    if name in ("u_m", "u_a", "u_c", "v_a", "sin_ua", "cos_ua"):
        return "sensor"
    if name.startswith("Cd") or name.startswith("Cl"):
        return "force"
    if "lag1" in name:
        return "memory_lag"
    if name.startswith("da"):
        return "memory_delta"
    if name == "mu" or name.startswith("mu_"):
        return "mu_mod"
    return "other"


def classify(
    a_active: Dict[str, float],
    b_active: Dict[str, float],
) -> Tuple[List[Tuple[str, float, float]], List[Tuple[str, float]], List[Tuple[str, float]]]:
    """Classify features as shared, A-only, B-only.

    Returns (shared, A_only, B_only) where shared has feature name + both coeffs.
    """
    a_keys = set(a_active.keys())
    b_keys = set(b_active.keys())

    shared = sorted(a_keys & b_keys)
    a_only = sorted(a_keys - b_keys)
    b_only = sorted(b_keys - a_keys)

    shared_out = [(k, a_active[k], b_active[k]) for k in shared]
    a_out = [(k, a_active[k]) for k in a_only]
    b_out = [(k, b_active[k]) for k in b_only]

    return shared_out, a_out, b_out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sindy-results", type=str, required=True)
    ap.add_argument("--scenes", type=str, nargs="+", required=True,
                    help="Scene names to compare")
    ap.add_argument("--threshold", type=float, default=RELATIVE_THRESHOLD)
    ap.add_argument("--channels", type=str, nargs="+",
                    default=["front", "top", "bottom"],
                    help="Which channels to compare")
    ap.add_argument("--out", type=str, default=None)
    args = ap.parse_args()

    data = load_sindy_scenes(args.sindy_results, args.scenes)

    if len(data) < 2:
        print("Need at least 2 scenes to compare")
        return 1

    scene_names = list(data.keys())
    scene_a, scene_b = scene_names[0], scene_names[1]

    print(f"Support Overlap: {scene_a} vs {scene_b} (th={args.threshold})")
    print("=" * 60)

    all_results = {}

    for ch_name in args.channels:
        # Map "front" -> "feature_names_front", etc
        fn_key = f"feature_names_{'front' if ch_name == 'front' else 'rear'}"
        fn_a = data[scene_a][fn_key]
        fn_b = data[scene_b][fn_key]

        fn_min = min(len(fn_a), len(fn_b))
        fn_a_trim = fn_a[:fn_min]
        fn_b_trim = fn_b[:fn_min]

        ch_a = get_active_support(np.array(data[scene_a][ch_name]["best_coef"])[:fn_min],
                                   fn_a_trim, args.threshold)
        ch_b = get_active_support(np.array(data[scene_b][ch_name]["best_coef"])[:fn_min],
                                   fn_b_trim, args.threshold)

        shared, a_only, b_only = classify(ch_a, ch_b)

        print(f"\n--- {ch_name} ---")
        print(f"  {scene_a} nz={len(ch_a)}  {scene_b} nz={len(ch_b)}  Shared={len(shared)}")

        for name, ca, cb in shared:
            print(f"    {name:20s}  A={ca:+9.6f}  B={cb:+9.6f}  [{feat_group(name)}]")
        for name, ca in a_only:
            print(f"    {scene_a[:10]:>10s} {name:20s}  A={ca:+9.6f}  [{feat_group(name)}]")
        for name, cb in b_only:
            print(f"    {scene_b[:10]:>10s} {name:20s}  B={cb:+9.6f}  [{feat_group(name)}]")

        all_results[ch_name] = {
            "scene_a_nz": len(ch_a),
            "scene_b_nz": len(ch_b),
            "shared_nz": len(shared),
            "shared": [{"name": n, "coef_a": ca, "coef_b": cb} for n, ca, cb in shared],
            f"{scene_a}_only": [{"name": n, "coef": ca} for n, ca in a_only],
            f"{scene_b}_only": [{"name": n, "coef": cb} for n, cb in b_only],
        }

    if args.out:
        output = {"scene_a": scene_a, "scene_b": scene_b,
                  "threshold": args.threshold,
                  "channels": all_results}
        os.makedirs(os.path.dirname(args.out), exist_ok=True)
        with open(args.out, "w") as f:
            json.dump(output, f, indent=2)
        print(f"\nSaved: {args.out}")


if __name__ == "__main__":
    main()
