#!/usr/bin/env python
"""Auto-eval for C-mega on held-out LeKiwi: text-prompt x numeric-descriptor ablation.

Compares (numeric descriptor correct/wrong) x (text prompt correct/none) to see
whether the natural-language robot description adds anything beyond the C-scheme
numeric descriptor, zero-shot, no fine-tuning. Also reports the zero-floor and the
prior C-diverse baseline (295mm, numeric-only) for context.
"""

from __future__ import annotations

import numpy as np
import torch
import yaml
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
from scipy.spatial.transform import Rotation
from transformers import AutoTokenizer

from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize
from tinyvla.data.eval_utils import StateAdapter
from tinyvla.modeling_tinyvla import TinyVLAPolicy
from tinyvla.modules.embodiment import MORPH_FIELDS

CKPT = "outputs/tv2_C_mega/final"
DS = "lekiwi_cleanup"
ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}"
_SC = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3,
       "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416,
       "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2}


def mvec(d):
    return torch.tensor([d.get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32)


@torch.no_grad()
def main():
    pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval()
    cfg = pol.config
    tok = AutoTokenizer.from_pretrained(cfg.lm_model_name)
    desc = yaml.safe_load(open("configs/morphology/descriptors.yaml"))
    prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml"))

    m = LeRobotDatasetMetadata(DS, root=ROOT)
    ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]},
                        video_backend="torchcodec")
    _sa = StateAdapter(ds.meta, cfg.max_state_dim)
    store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50)
    st = store.compute_stats()
    q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"])
    span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6))
    mid = 0.5 * (q01 + q99)

    imk = sorted(k for k in ds.meta.features if k.startswith("observation.images"))
    prim = next((k for k in imk if "front" in k or "base" in k), imk[0])

    def integ(d):
        return np.cumsum(d[:, :3], 0)

    n_eps = ds.num_episodes
    test = range(max(0, n_eps - 15), n_eps)

    def run(morph, prompt_prefix):
        errs, zf = [], []
        for ep in test:
            s = int(m.episodes["dataset_from_index"][ep])
            e = int(m.episodes["dataset_to_index"][ep])
            for idx in range(s, e - 1, 30):
                item = ds[idx]
                task = item.get("task") or ""
                text = f"{prompt_prefix} {task}" if prompt_prefix else task
                t = tok([text], padding=True, truncation=True, max_length=48, return_tensors="pt")
                img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0]
                stt = _sa(item["observation.state"])
                b = {"observation.images.cam0": img[None].cuda(),
                     "observation.images.cam1": torch.zeros_like(img)[None].cuda(),
                     "observation.state": stt[None].cuda(),
                     "observation.language.tokens": t["input_ids"].cuda(),
                     "observation.language.attention_mask": t["attention_mask"].bool().cuda(),
                     "morphology": morph[None].cuda(), "embodiment_id": torch.tensor([0]).cuda()}
                with torch.autocast("cuda", torch.bfloat16):
                    pr = pol.predict_action_chunk(b)[0].cpu().float().numpy()
                gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid
                pu = pr[:, :7] * span / 2 + mid
                errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000)
                zf.append(np.linalg.norm(integ(gu)[-1]) * 1000)
        return np.mean(errs), np.mean(zf)

    print(f"=== C-mega on held-out LeKiwi (n_eps_test={len(list(test))}) ===")
    zero_floor = None
    conditions = [
        ("numeric=lekiwi + text=lekiwi", mvec(desc["lekiwi"]), prompts["lekiwi"]),
        ("numeric=lekiwi + text=none  ", mvec(desc["lekiwi"]), None),
        ("numeric=none   + text=lekiwi", torch.zeros(16), prompts["lekiwi"]),
        ("numeric=none   + text=none  ", torch.zeros(16), None),
        ("numeric=so101(wrong,non-mobile) + text=lekiwi", mvec(desc["so101"]), prompts["lekiwi"]),
    ]
    for label, morph, prompt in conditions:
        err, zf = run(morph, prompt)
        if zero_floor is None:
            zero_floor = zf
        print(f"{label:48} endpoint {err:.1f}mm")
    print(f"\nzero-floor: {zero_floor:.1f}mm")
    print("prior C-diverse baseline (numeric-only, no mega training): 295.0mm")


if __name__ == "__main__":
    main()