"""Shared core detection across scenes.

Finds features that are active across ALL scenes in a group (e.g. all Karman Re,
all Illusion diameters) and identifies the cross-scene shared core.

Usage:
    python compare/shared_core.py --sindy-results sindy/karman/sindy_results.json \\
        --scenes karman_re50 karman_re100 karman_re200 karman_re400
    python compare/shared_core.py \\
        --sindy-results sindy/results.json \\
        --scenes karman_re100 illusion_1L vortex_lamb steady
"""
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


def feat_group(name: str) -> str:
    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 detect_core(scene_data: Dict[str, dict], channels: List[str],
                threshold: float) -> dict:
    """Find features active in ALL scenes for each channel."""
    scene_names = list(scene_data.keys())
    results = {}

    for ch_name in channels:
        fn_key = f"feature_names_{'front' if ch_name == 'front' else 'rear'}"

        # Collect active sets per scene
        active_per_scene = {}
        for sn in scene_names:
            fn = scene_data[sn][fn_key]
            coef = scene_data[sn][ch_name]["best_coef"]
            active = get_active_support(np.array(coef, dtype=np.float64)[:len(fn)],
                                        fn, threshold)
            active_per_scene[sn] = set(active.keys())

        # Intersection = shared core
        core_keys = set.intersection(*active_per_scene.values()) if active_per_scene else set()

        # Union for scene-specific detection
        all_keys = set.union(*active_per_scene.values()) if active_per_scene else set()
        scene_specific = {}
        for sn in scene_names:
            others = set.union(*[v for k, v in active_per_scene.items() if k != sn])
            diff = active_per_scene[sn] - others
            if diff:
                scene_specific[sn] = sorted(diff)

        # Coef means for core features
        core_coefs = {}
        for k in sorted(core_keys):
            vals = []
            for sn in scene_names:
                fn = scene_data[sn][fn_key]
                coef = scene_data[sn][ch_name]["best_coef"]
                if k in fn:
                    idx = fn.index(k)
                    vals.append(float(coef[idx]))
            core_coefs[k] = {
                "mean": float(np.mean(vals)),
                "std": float(np.std(vals)),
                "per_scene": {sn: vals[i] for i, sn in enumerate(scene_names)},
            }

        results[ch_name] = {
            "n_scenes": len(scene_names),
            "n_core": len(core_keys),
            "core_features": {k: {"group": feat_group(k), "coef": v}
                              for k, v in core_coefs.items()},
            "scene_specific": {sn: sorted(v) for sn, v in scene_specific.items()},
        }

    return results


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sindy-results", type=str, required=True)
    ap.add_argument("--scenes", type=str, nargs="+", required=True)
    ap.add_argument("--threshold", type=float, default=RELATIVE_THRESHOLD)
    ap.add_argument("--out", type=str, default=None)
    args = ap.parse_args()

    with open(args.sindy_results) as f:
        all_data = json.load(f)

    per = all_data.get("per_scene", {})
    scene_data = {sn: per[sn] for sn in args.scenes if sn in per}
    if len(scene_data) < 2:
        print(f"Need >=2 scenes. Found: {list(scene_data.keys())}")
        return 1

    print(f"Shared Core Detection: {len(scene_data)} scenes")
    for sn in scene_data:
        print(f"  {sn}")
    print(f"  threshold={args.threshold}")

    results = detect_core(scene_data, ["front", "top", "bottom"], args.threshold)

    for ch_name, ch_data in results.items():
        print(f"\n--- {ch_name} ---")
        print(f"  Core features: {ch_data['n_core']}")
        for k, v in ch_data["core_features"].items():
            c = v["coef"]
            print(f"    {k:20s}  mean={c['mean']:+.6f}  std={c['std']:.6f}  [{v['group']}]")
        for sn, keys in ch_data["scene_specific"].items():
            if keys:
                print(f"  {sn} specific: {', '.join(keys)}")

    if args.out:
        output = {"scenes": args.scenes, "threshold": args.threshold,
                  "channels": 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()