#!/usr/bin/env python3
"""Bridge V5 eval outputs to SR_analysis/data/v5/ format.

For each effective V5 case:
  1. Reads signals.npz from eval/output/train/{case}/
  2. Finds best seed from metrics.json
  3. Finds target.npy and target_harmonics.json from train/output/{case}_seed{N}/
  4. Uses calibration.json from calibrations/
  5. Creates SR_analysis/data/v5/{case}/ with controlled.npz, target.npz, etc.

Effective cases (with best seed and DTW):
  - kar_re100 (seed45, dtw=0.918) — baseline
  - kar_d075 (seed44, dtw=0.911) — vardist
  - kar_d15 (seed45, dtw=0.892) — vardist
  - kar_d2 (seed45, dtw=0.816) — vardist, borderline
  - ill_075L (seed43, dtw=0.807)
  - ill_1L (seed43, dtw=0.900)
  - ill_15L (seed43, dtw=0.906)
  - ill_2L (seed43, dtw=0.800)

Usage:
    PYTHONPATH=src conda run -n pycuda_3_10 python -m SR_analysis.experiments.v5.bridge_v5
    PYTHONPATH=src conda run -n pycuda_3_10 python -m SR_analysis.experiments.v5.bridge_v5 --include-failed
    PYTHONPATH=src conda run -n pycuda_3_10 python -m SR_analysis.experiments.v5.bridge_v5 --cases kar_re100,ill_1L
"""
from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import numpy as np

from drl_pinball.case_registry import CASE_REGISTRY

_SR_ROOT = Path(__file__).resolve().parents[2]
_REPO = _SR_ROOT.parents[1]
_EVAL_OUT = _REPO / "src" / "drl_pinball" / "eval" / "output" / "train"
_TRAIN_OUT = _REPO / "src" / "drl_pinball" / "train" / "output"
_TRAIN_CAL = _REPO / "src" / "drl_pinball" / "train" / "calibrations"
_SR_DATA = _SR_ROOT / "data" / "v5"

L0 = 20.0
U0 = 0.01

# Evaluation attributes not supplied by the canonical registry.
_CASE_METRICS = {
    "kar_re100": (45, 0.918, 30, (2, 14), 7),
    "kar_d075": (44, 0.911, 30, (2, 14), 7),
    "kar_d15": (45, 0.892, 30, (2, 14), 7),
    "kar_d2": (45, 0.816, 30, (2, 14), 7),
    "ill_075L": (43, 0.807, 36, (0, 12), 6),
    "ill_1L": (43, 0.900, 36, (0, 12), 6),
    "ill_15L": (43, 0.906, 36, (0, 12), 6),
    "ill_2L": (43, 0.800, 36, (0, 12), 6),
    "kar_re60": (43, 0.364, 30, (2, 14), 7),
    "kar_re200": (43, 0.712, 30, (2, 14), 7),
    "kar_re400": (43, 0.565, 30, (2, 14), 7),
}


def _bridge_config(case_id: str) -> dict:
    case = CASE_REGISTRY[case_id]
    seed, dtw, conv_len, obs_slice, n_objects = _CASE_METRICS[case_id]
    cfg = {
        "seed": seed, "dtw": dtw, "scene_type": case.scene_type,
        "cal_name": case.calibration, "si": case.si,
        "action_scale": 12.0, "action_bias": (0.0, 0.0, 0.0),
        "s_dim": 14 if case.scene_type == "illusion" else 12,
        "conv_len": conv_len, "obs_slice": obs_slice, "n_objects": n_objects,
    }
    if case.target_diam is not None:
        cfg["target_diam"] = case.target_diam
    return cfg


EFFECTIVE_CASES = {case_id: _bridge_config(case_id) for case_id in (
    "kar_re100", "kar_d075", "kar_d15", "kar_d2",
    "ill_075L", "ill_1L", "ill_15L", "ill_2L",
)}
FAILED_CASES = {case_id: _bridge_config(case_id) for case_id in (
    "kar_re60", "kar_re200", "kar_re400",
)}


def bridge_case(case_name: str, cfg: dict, out_dir: Path) -> bool:
    seed = cfg["seed"]
    scene_type = cfg["scene_type"]
    cal_name = cfg["cal_name"]

    signals_path = _EVAL_OUT / case_name / "signals.npz"
    if not signals_path.exists():
        print(f"  SKIP {case_name}: signals.npz not found")
        return False

    target_dir = _TRAIN_OUT / "{}_seed{}".format(case_name, seed)
    target_npy = target_dir / "target.npy"
    target_harm = target_dir / "target_harmonics.json"

    if not target_npy.exists():
        print(f"  SKIP {case_name}: target.npy not found at {target_npy}")
        return False

    cal_path = _TRAIN_CAL / cal_name / "calibration.json"
    if not cal_path.exists():
        cal_path = target_dir / "calibration.json"
    if not cal_path.exists():
        print(f"  SKIP {case_name}: calibration.json not found")
        return False

    out_dir.mkdir(parents=True, exist_ok=True)
    print(f"  Bridging {case_name} (seed={seed}, dtw={cfg['dtw']:.3f})...")

    # 1. signals.npz -> controlled.npz
    data = dict(np.load(signals_path, allow_pickle=True))
    np.savez_compressed(out_dir / "controlled.npz",
                        sensors=data["sensors"],
                        forces=data["forces"],
                        actions=data["actions"])

    # 2. target
    target = np.load(target_npy)
    np.savez_compressed(out_dir / "target.npz", target_states=target)

    # 3. Harmonics (illusion only)
    if scene_type == "illusion" and target_harm.exists():
        shutil.copy(target_harm, out_dir / "target_harmonics.json")

    # 4. norm.json from calibration
    with open(cal_path) as f:
        cal = json.load(f)

    force_scale = float(cal.get("FORCE_SCALE", 1.0))
    sens_scale = float(cal.get("SENS_SCALE", 1.0))
    action_bias = cal.get("ACTION_BIAS", [0.0, 0.0, 0.0])

    norm = {
        "force_norm_fact": force_scale,
        "sens_deviation": [0.0] * 6,
        "sens_norm_fact": [sens_scale] * 6,
        "action_bias": action_bias,
    }
    with open(out_dir / "norm.json", "w") as f:
        json.dump(norm, f, indent=2)

    # 5. config.json
    config = {
        "case_name": case_name,
        "scene_type": scene_type,
        "seed": seed,
        "SI": cfg["si"],
        "conv_len": cfg["conv_len"],
        "action_scale": cfg["action_scale"],
        "action_bias": list(cfg["action_bias"]),
        "u0": U0,
        "l0": L0,
        "s_dim": cfg["s_dim"],
        "obs_slice": cfg["obs_slice"],
        "n_objects": cfg["n_objects"],
        "dtw_sim_v5": cfg["dtw"],
    }
    if "target_diam" in cfg:
        config["target_diam"] = cfg["target_diam"]
    with open(out_dir / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"    -> {out_dir}")
    return True


def main() -> int:
    parser = argparse.ArgumentParser(description="Bridge V5 eval data to SR_analysis format")
    parser.add_argument("--cases", type=str, default=None,
                        help="Comma-separated case names (default: all effective)")
    parser.add_argument("--include-failed", action="store_true",
                        help="Also bridge cross-Re failed cases")
    args = parser.parse_args()

    all_cases = dict(EFFECTIVE_CASES)
    if args.include_failed:
        all_cases.update(FAILED_CASES)

    if args.cases:
        requested = [c.strip() for c in args.cases.split(",")]
        cases = {k: v for k, v in all_cases.items() if k in requested}
    else:
        cases = all_cases

    print(f"Bridging {len(cases)} V5 cases to SR_analysis/data/v5/...")
    print(f"SR data root: {_SR_DATA}")

    success = 0
    for case_name, cfg in cases.items():
        out_dir = _SR_DATA / case_name
        if bridge_case(case_name, cfg, out_dir):
            success += 1

    print(f"\nDone: {success}/{len(cases)} cases bridged.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
