"""Fair multi-seed + long-horizon benchmark for Spectral World Models."""
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=["dreamer","transformer","koopman","neural_operator","swm","no_spectral_state","no_spectral_transition","no_text_branch","no_stability_penalty"]
BASE_METRICS=[("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="Run fair multi-seed SWM benchmark with H=5/10/20/30 rollouts.")
    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])
    ap.add_argument("--device",default="auto"); ap.add_argument("--models",nargs="+",default=DEFAULT_MODELS)
    args=ap.parse_args(); root=Path(__file__).resolve().parents[1]
    data=Path(args.data) if args.data else root/"data"/"mini_moving_shapes.npz"
    out=Path(args.out) if args.out else root/"results"/"multiseed_horizons"
    rollout_data=Path(args.rollout_data) if args.rollout_data else root/"data"/"mini_moving_shapes_rollout_h31.npz"
    if not data.exists(): generate_benchmark_npz(data,DatasetConfig())
    need_len=max(args.horizons)+1
    regenerate=True
    if rollout_data.exists():
        try:
            import numpy as np
            with np.load(rollout_data,allow_pickle=True) as d: regenerate=d["test_images"].shape[1] < need_len
        except Exception: regenerate=True
    if regenerate:
        # Separate held-out long sequences: training data remain unchanged.
        generate_benchmark_npz(rollout_data,DatasetConfig(seq_len=need_len,train_sequences=1,val_sequences=1,test_sequences=64,seed=7007))
    device=("cuda" if torch.cuda.is_available() else "cpu") if args.device=="auto" else args.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 args.models}; raw=[]
    for seed in args.seeds:
        sd=out/f"seed_{seed}"
        for m in args.models:
            print(f"\n=== seed={seed} model={m} horizons={args.horizons} ===")
            r=train_one_model(m,data,sd,epochs=args.epochs,batch_size=args.batch_size,seed=seed,device=device,rollout_data_path=rollout_data,rollout_horizons=tuple(args.horizons))
            grouped[m].append(r); row={"seed":seed,"model":m,"params":r["params"]}
            for k,g in BASE_METRICS: row[k]=g(r)
            for h in args.horizons:
                rr=r["rollout_by_horizon"][str(h)]; row[f"rollout_mse_h{h}"]=rr["rollout_mse"]; row[f"time_per_rollout_step_ms_h{h}"]=rr["time_per_rollout_step_ms"]
            raw.append(row)
    with (out/"metrics_by_seed.csv").open("w",newline="") as f:
        w=csv.DictWriter(f,fieldnames=list(raw[0])); 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_METRICS:
            mu,sd=stats([float(g(r)) for r in runs]); row[k+"_mean"]=mu; row[k+"_std"]=sd
        for h in args.horizons:
            vals=[float(r["rollout_by_horizon"][str(h)]["rollout_mse"]) for r in runs]; mu,sd=stats(vals); row[f"rollout_mse_h{h}_mean"]=mu; row[f"rollout_mse_h{h}_std"]=sd
        summary.append(row)
    with (out/"metrics_summary.csv").open("w",newline="") as f:
        w=csv.DictWriter(f,fieldnames=list(summary[0])); w.writeheader(); w.writerows(summary)
    with (out/"run_config.json").open("w") as f: json.dump(vars(args)|{"device_resolved":device,"rollout_data_resolved":str(rollout_data)},f,indent=2)
    print("\n=== Long-horizon summary ===")
    for r in summary: print(r["model"]+": "+" | ".join(f'H{h} {r[f"rollout_mse_h{h}_mean"]:.6f} ± {r[f"rollout_mse_h{h}_std"]:.6f}' for h in args.horizons))
    print(f"\nWrote: {out/'metrics_summary.csv'}")
if __name__=="__main__": main()