"""Targeted attribution benchmark for SWM transition stability."""
from __future__ import annotations
import argparse, csv, json
from pathlib import Path
from statistics import mean, stdev
import torch
from spectral_world_models.data import DatasetConfig, generate_benchmark_npz
from spectral_world_models.train import train_one_model

DEFAULT_MODELS=["koopman","neural_operator","swm","swm_selective","swm_spectral_norm","no_spectral_transition","no_stability_penalty"]
BASE=[("test_psnr",lambda r:r["test"]["psnr"]),("test_ssim",lambda r:r["test"]["ssim"]),("test_nll",lambda r:r["test"]["nll"]),("test_mse",lambda r:r["test"]["mse"]),("test_text_acc",lambda r:r["test"]["text_acc"])]
def stats(xs): return mean(xs), stdev(xs) if len(xs)>1 else 0.0

def main():
    ap=argparse.ArgumentParser(description="SWM v4 benchmark: selective spectral regulation and long-horizon multimodal information preservation.")
    ap.add_argument("--data",default=None); ap.add_argument("--rollout-data",default=None); ap.add_argument("--out",default=None)
    ap.add_argument("--epochs",type=int,default=3); ap.add_argument("--batch-size",type=int,default=64)
    ap.add_argument("--seeds",type=int,nargs="+",default=[0,1,2,3,4]); ap.add_argument("--horizons",type=int,nargs="+",default=[5,10,20,30,50,100])
    ap.add_argument("--device",default="auto"); ap.add_argument("--models",nargs="+",default=DEFAULT_MODELS); ap.add_argument("--lambda-stability",type=float,default=0.01)
    a=ap.parse_args(); root=Path(__file__).resolve().parents[1]
    data=Path(a.data) if a.data else root/"data"/"mini_moving_shapes.npz"
    out=Path(a.out) if a.out else root/"results"/"stability_attribution"
    rollout=Path(a.rollout_data) if a.rollout_data else root/f"data/mini_moving_shapes_rollout_h{max(a.horizons)+1}.npz"
    if not data.exists(): generate_benchmark_npz(data,DatasetConfig())
    need=max(a.horizons)+1; regen=True
    if rollout.exists():
        try:
            import numpy as np
            with np.load(rollout,allow_pickle=True) as d: regen=d["test_images"].shape[1] < need
        except Exception: pass
    if regen: generate_benchmark_npz(rollout,DatasetConfig(seq_len=need,train_sequences=1,val_sequences=1,test_sequences=64,seed=7007))
    device=("cuda" if torch.cuda.is_available() else "cpu") if a.device=="auto" else a.device
    if device.startswith("cuda") and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")
    out.mkdir(parents=True,exist_ok=True); grouped={m:[] for m in a.models}; raw=[]
    for seed in a.seeds:
        for m in a.models:
            print(f"\n=== seed={seed} model={m} horizons={a.horizons} ===")
            r=train_one_model(m,data,out/f"seed_{seed}",epochs=a.epochs,batch_size=a.batch_size,seed=seed,device=device,rollout_data_path=rollout,rollout_horizons=tuple(a.horizons),lambda_stability=a.lambda_stability)
            grouped[m].append(r); row={"seed":seed,"model":m,"params":r["params"]}
            for k,g in BASE: row[k]=g(r)
            for h in a.horizons:
                rr=r["rollout_by_horizon"][str(h)]
                for metric in ["rollout_mse","rollout_text_acc","rollout_latent_cosine","rollout_joint_score"]: row[f"{metric}_h{h}"]=rr[metric]
            for k,v in r.get("stability_diagnostics",{}).items(): row[k]=v
            raw.append(row)
    fields=[]
    for r in raw:
        for k in r:
            if k not in fields: fields.append(k)
    with (out/"metrics_by_seed.csv").open("w",newline="") as f:
        w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(raw)
    summary=[]
    for m,runs in grouped.items():
        row={"model":m,"params":runs[0]["params"],"n_seeds":len(runs)}
        for k,g in BASE:
            mu,sd=stats([float(g(r)) for r in runs]); row[k+"_mean"]=mu; row[k+"_std"]=sd
        for h in a.horizons:
            for metric in ["rollout_mse","rollout_text_acc","rollout_latent_cosine","rollout_joint_score"]:
                vals=[float(r["rollout_by_horizon"][str(h)][metric]) for r in runs]; mu,sd=stats(vals); row[f"{metric}_h{h}_mean"]=mu; row[f"{metric}_h{h}_std"]=sd
        diag_keys=set().union(*(r.get("stability_diagnostics",{}).keys() for r in runs))
        for k in sorted(diag_keys):
            vals=[float(r["stability_diagnostics"][k]) for r in runs if k in r.get("stability_diagnostics",{})]; mu,sd=stats(vals); row[k+"_mean"]=mu; row[k+"_std"]=sd
        summary.append(row)
    fields=[]
    for r in summary:
        for k in r:
            if k not in fields: fields.append(k)
    with (out/"metrics_summary.csv").open("w",newline="") as f:
        w=csv.DictWriter(f,fieldnames=fields); w.writeheader(); w.writerows(summary)
    with (out/"run_config.json").open("w") as f: json.dump(vars(a)|{"device_resolved":device,"rollout_data_resolved":str(rollout)},f,indent=2)
    print("\nWrote",out/"metrics_summary.csv")
if __name__=="__main__": main()