#!/usr/bin/env python
"""Score every ablation on ONE common validation set.

The training runs each used their own valid-window set (requiring more history invalidates more
windows near episode starts), so their feat_mse values are not comparable. This evaluates every
checkpoint on the windows valid for the largest history, feeding each model only the context
frames it was trained with.

    policy_learning/lerobot/.venv/bin/python dynamics_model/eval_common.py
"""

import torch
import torch.nn.functional as F
from diffusers import AutoencoderKL
from transformers import AutoModel

from dino_dynamics import DINO_ID, DinoDynamics
from sd3_dynamics import build_datasets, dataset_features, per_camera_metrics

DATA = [
    "chomeed/mimicgen_coffee_d0_224x224_mtdit_flow_60k_success",
    "chomeed/mimicgen_coffee_d0_224x224_mtdit_flow_60k_failure",
    "chomeed/mimicgen_coffee_d0_224x224",
]
SD3 = "stabilityai/stable-diffusion-3.5-medium"
COMMON_HISTORY = 3  # most restrictive -> its windows are a subset of every arm's
BATCHES, BS = 8, 32

ARMS = [  # name, checkpoint, history it was trained with, use_action, use_state
    ("h1", "outputs/abl_h1/dino_step_10000.pt", 1, True, True),
    ("h2 (baseline)", "outputs/dino_step_10000.pt", 2, True, True),
    ("h3", "outputs/abl_h3/dino_step_10000.pt", 3, True, True),
    ("noaction", "outputs/abl_noaction/dino_step_10000.pt", 2, False, True),
    ("nostate", "outputs/abl_nostate/dino_step_10000.pt", 2, True, False),
]


def main():
    dev = "cuda"
    torch.manual_seed(0)
    _, val, cameras = build_datasets(DATA, (224, 224), "observation.state", 9, 20, COMMON_HISTORY, 8)
    print(f"common val windows: {len(val)}  cameras: {cameras}")
    action_dim = dataset_features(DATA[0])["action"]["shape"][0]

    dino = AutoModel.from_pretrained(DINO_ID, torch_dtype=torch.float32)
    vae = AutoencoderKL.from_pretrained(SD3, subfolder="vae", torch_dtype=torch.bfloat16)
    from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity

    lpips = LearnedPerceptualImagePatchSimilarity(net_type="alex", normalize=True).to(dev)

    # fixed batches, identical for every arm
    dl = torch.utils.data.DataLoader(val, batch_size=BS, shuffle=False, num_workers=8, drop_last=True)
    batches = []
    for i, b in enumerate(dl):
        if i >= BATCHES:
            break
        batches.append({k: v.to(dev) for k, v in b.items()})

    print(f"\n{'arm':<15}{'feat_mse':>10}{'copy':>9}{'ratio':>8}{'psnr':>8}{'ssim':>8}{'lpips':>8}")
    for name, ckpt, hist, use_act, use_st in ARMS:
        model = DinoDynamics(
            dino, vae, len(cameras), action_dim, 9, history=hist,
            use_action=use_act, use_state=use_st,
        ).to(dev)
        sd = torch.load(ckpt, map_location=dev)
        missing, unexpected = model.load_state_dict(sd, strict=False)
        assert not unexpected, f"{name}: unexpected keys {unexpected[:3]}"
        assert all(k.startswith(("dino.", "vae.")) for k in missing), f"{name}: missing {missing[:3]}"
        model.eval()

        sums = {}
        with torch.no_grad():
            for b in batches:
                b = dict(b, context=b["context"][:, -hist:])  # only the frames this arm was trained on
                pred = model(b)
                tgt = model.encode_views(b["future"][:, None])[:, 0]
                cur = model.encode_views(b["context"][:, -1:])[:, 0]
                sums["feat"] = sums.get("feat", 0.0) + F.mse_loss(pred, tgt).item()
                sums["copy"] = sums.get("copy", 0.0) + F.mse_loss(cur, tgt).item()
                for k, v in per_camera_metrics(
                    model.features_to_image(pred), b["future"], cameras, lpips
                ).items():
                    if "mean" in k:
                        sums[k] = sums.get(k, 0.0) + v
        n = len(batches)
        f, c = sums["feat"] / n, sums["copy"] / n
        print(f"{name:<15}{f:>10.4f}{c:>9.4f}{f / c:>8.3f}"
              f"{sums['mean/psnr'] / n:>8.2f}{sums['mean/ssim'] / n:>8.4f}{sums['mean/lpips'] / n:>8.4f}")
        del model
        torch.cuda.empty_cache()


if __name__ == "__main__":
    main()