"""Open-loop evaluation of a (text-conditioned) FACT checkpoint over many episodes.

Same per-anchor protocol as scripts/eval.py (K=1 chunk, pure-noise start,
seed 500+i, normalized 62-D delta-base action MSE, arm/hand split, first-11-step
horizon-matched MSE, future-state MSE), applied to ``--anchors-per-episode``
evenly spaced anchors of every requested episode.  Latents come from the
per-episode cache directory, text from a umT5 text cache.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from fact import FactPolicy
from fact.data import (
    FactEpisode,
    TextCache,
    latent_cache_path,
    load_episode_tasks,
)
from fact.runtime_stats import TrexRuntimeStatistics

ARM = list(range(0, 9)) + list(range(31, 40))
HAND = list(range(9, 31)) + list(range(40, 62))


def parse_episodes(spec: str) -> list[int]:
    out: list[int] = []
    for part in spec.split(","):
        part = part.strip()
        if "-" in part:
            a, b = part.split("-")
            out.extend(range(int(a), int(b) + 1))
        elif part:
            out.append(int(part))
    return out


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--dataset-root", required=True)
    p.add_argument("--cache-dir", required=True, help="dir with latents/episode_XXXXXX.npy")
    p.add_argument("--text-cache", default=None, help="text_cache.npz (default <cache-dir>/text_cache.npz)")
    p.add_argument("--stats-root", default=None, help="dataset root whose q01/q99 stats normalize inputs/outputs (default dataset-root)")
    p.add_argument("--episodes", required=True, help="e.g. 0-9,20")
    p.add_argument("--anchors-per-episode", type=int, default=8)
    p.add_argument("--out", required=True)
    p.add_argument("--bf16", action="store_true")
    p.add_argument("--dump", default=None, help="npz path: raw (denormalized delta-base) predictions/GT per anchor")
    p.add_argument("--tau-losses", action="store_true", help="also compute fixed-tau flow losses on K=4 anchors")
    p.add_argument("--tau-grid", default="0.8,0.6,0.4,0.2,0.05")
    args = p.parse_args()

    device = torch.device("cuda")
    model = FactPolicy.load(args.checkpoint).to(device).eval()
    cfg = model.cfg
    root = Path(args.dataset_root)
    cache = Path(args.cache_dir)
    text_cache = None
    if cfg.use_text:
        text_cache = TextCache(args.text_cache or cache / "text_cache.npz")
    stats = TrexRuntimeStatistics.from_dataset(args.stats_root or root)
    tasks = load_episode_tasks(root)

    per_episode: dict[int, dict[str, float]] = {}
    all_metrics: dict[str, list[float]] = {}
    dump: dict[str, list] = {"gt": [], "pred_cascade": [], "pred_coarse_to_zero": [], "episode": [], "anchor": []}
    tau_losses: dict[str, dict[str, list[float]]] = {}
    denorm = lambda x: (np.asarray(x, np.float32) + 1.0) * 0.5 * (stats.action_q99 - stats.action_q01) + stats.action_q01
    for ep in parse_episodes(args.episodes):
        episode = FactEpisode(root, ep, latent_cache_path(cache, ep), stats=stats, text_cache=text_cache, task=tasks.get(ep))
        anchors = episode.valid_anchor_times(1)
        picks = [anchors[int(i)] for i in np.linspace(0, len(anchors) - 1, min(args.anchors_per_episode, len(anchors)))]
        ep_metrics: dict[str, list[float]] = {}
        if args.tau_losses:
            k4 = episode.valid_anchor_times(4)
            for i, anchor in enumerate([k4[int(j)] for j in np.linspace(0, len(k4) - 1, min(3, len(k4)))] if k4 else []):
                sample4 = episode.build(anchor, blocks=4)
                batch4 = {k: torch.as_tensor(v)[None].to(device) for k, v in sample4.items()
                          if k not in ("action_delta_raw", "reference_state", "obs_frame_index")}
                for tau_value in (float(t) for t in args.tau_grid.split(",")):
                    for rep in range(2):
                        torch.manual_seed(100000 + rep * 1000 + i * 10)
                        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16):
                            o = model.forward_core(batch4, tau=torch.full((1, 4), tau_value, device=device))
                        for key in ("action_loss", "state_loss", "force_loss"):
                            tau_losses.setdefault(f"{tau_value:g}", {}).setdefault(key, []).append(float(o[key]))
        for i, anchor in enumerate(picks):
            sample = episode.build(anchor, blocks=1)
            gt = torch.as_tensor(sample["action"][0, :, :62])
            gt_fs = torch.as_tensor(sample["future_state"][0, :, :62])
            batch = {k: torch.as_tensor(v)[None].to(device) for k, v in sample.items()
                     if k not in ("action_delta_raw", "reference_state", "obs_frame_index")}
            for mode in ("cascade", "coarse_only", "coarse_to_zero"):
                with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.bf16):
                    r = model.sample(batch, seed=500 + i, mode=mode)
                pred = r["action_pred"][0, :, :62].float().cpu()
                ep_metrics.setdefault(mode, []).append(float(((pred - gt) ** 2).mean()))
                ep_metrics.setdefault(f"{mode}_first11", []).append(float(((pred[:11] - gt[:11]) ** 2).mean()))
                ep_metrics.setdefault(f"{mode}_arm", []).append(float(((pred[:, ARM] - gt[:, ARM]) ** 2).mean()))
                ep_metrics.setdefault(f"{mode}_hand", []).append(float(((pred[:, HAND] - gt[:, HAND]) ** 2).mean()))
                if mode == "cascade" and "future_state_pred" in r:
                    fp = r["future_state_pred"][0, :, :62].float().cpu()
                    ep_metrics.setdefault("future_state", []).append(float(((fp - gt_fs) ** 2).mean()))
                if mode in ("cascade", "coarse_to_zero"):
                    dump[f"pred_{mode}"].append(denorm(pred.numpy()))
            dump["gt"].append(sample["action_delta_raw"][0])
            dump["episode"].append(ep)
            dump["anchor"].append(anchor)
        per_episode[ep] = {k: float(np.mean(v)) for k, v in ep_metrics.items()}
        for k, v in ep_metrics.items():
            all_metrics.setdefault(k, []).extend(v)
        print(f"episode {ep:4d} ({episode.length} frames, {len(picks)} anchors): "
              + " ".join(f"{k} {per_episode[ep][k]:.4f}" for k in ("cascade", "coarse_to_zero", "future_state") if k in per_episode[ep]), flush=True)

    summary = {k: [float(np.mean(v)), float(np.std(v)), len(v)] for k, v in all_metrics.items()}
    if args.dump:
        np.savez_compressed(args.dump, **{k: np.asarray(v) for k, v in dump.items()})
        print("dumped", args.dump)
    record = {"checkpoint": args.checkpoint, "dataset_root": str(root), "episodes": sorted(per_episode),
              "anchors_per_episode": args.anchors_per_episode, "config": json.loads(cfg.to_json()),
              "params_total": sum(p.numel() for p in model.parameters()),
              "openloop_action_mse": summary, "per_episode": per_episode,
              "forward_tau_losses": {t: {k: [float(np.mean(v)), float(np.std(v))] for k, v in d.items()} for t, d in tau_losses.items()}}
    Path(args.out).write_text(json.dumps(record, indent=1))
    print("summary:", {k: round(v[0], 4) for k, v in summary.items() if not k.endswith(("_arm", "_hand"))})
    print("wrote", args.out)


if __name__ == "__main__":
    main()