#!/usr/bin/env python
"""Gate 1: per-dataset SO101 FK calibration audit.

Same recorded degrees can map to different physical poses across lerobot
zero-conventions → FK garbage that would poison B/C canonical labels while
leaving A untouched (a confound AGAINST the hypothesis). Drop datasets whose
FK produces non-physical EE trajectories.

Checks per dataset (sampled frames across episodes):
  - reach in [0.02, 0.40] m (SO101 max reach ~0.35)
  - z above a floor (> -0.20 m; base frame origin at arm mount)
  - trajectory smoothness (median consecutive EE step < 0.05 m at 30fps)
  - FK(state) vs FK(action-target) diffs correlate (both go through FK, deltas
    should track since action = commanded target of the same arm)

Writes ~/tinyvla_data/so101_fk_audit.json with per-dataset verdict.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np

DATA_ROOT = Path.home() / "tinyvla_data" / "so101_v3"
OUT = Path.home() / "tinyvla_data" / "so101_fk_audit.json"


def audit_one(name, root, fk, n_frames=200):
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    ds = LeRobotDataset(name, root=root)
    anames = ds.meta.features["action"]["shape"][0]
    if anames != 6:
        return {"verdict": "SKIP", "reason": f"action dim {anames} != 6"}
    # smoothness needs CONSECUTIVE frames → sample a few contiguous windows;
    # reach/corr can use sparse knots across the whole dataset
    n = len(ds)
    eps = ds.meta.episodes
    sparse = np.linspace(0, n - 1, min(n_frames, n)).astype(int)
    pos_state, pos_act = [], []
    for i in sparse:
        item = ds[int(i)]
        pos_state.append(fk.ee_pose(item["observation.state"].numpy())[:3, 3])
        pos_act.append(fk.ee_pose(item["action"].numpy())[:3, 3])
    ps = np.array(pos_state)
    pa = np.array(pos_act)
    reach = np.linalg.norm(ps, axis=1)

    # consecutive-frame EE steps within the first episode (real per-frame motion)
    e0, e1 = int(eps["dataset_from_index"][0]), int(eps["dataset_to_index"][0])
    consec = []
    prev = None
    for i in range(e0, min(e1, e0 + 150)):
        p = fk.ee_pose(ds[i]["observation.state"].numpy())[:3, 3]
        if prev is not None:
            consec.append(np.linalg.norm(p - prev))
        prev = p
    steps = np.array(consec) if consec else np.array([0.0])
    # correlation of state-vs-target displacement over sampled knots
    d_state = np.diff(ps, axis=0).flatten()
    d_act = np.diff(pa, axis=0).flatten()
    corr = float(np.corrcoef(d_state, d_act)[0, 1]) if d_state.std() > 1e-9 else 0.0

    ok_reach = bool(0.02 < reach.mean() < 0.40 and reach.max() < 0.50)
    ok_z = bool(ps[:, 2].min() > -0.20)
    ok_smooth = bool(np.median(steps) < 0.06)
    ok_corr = bool(corr > 0.5)
    verdict = "KEEP" if (ok_reach and ok_z and ok_smooth and ok_corr) else "DROP"
    return {
        "verdict": verdict,
        "reach_mean": round(float(reach.mean()), 3),
        "reach_max": round(float(reach.max()), 3),
        "z_min": round(float(ps[:, 2].min()), 3),
        "step_median": round(float(np.median(steps)), 4),
        "corr_state_target": round(corr, 3),
        "flags": {"reach": ok_reach, "z": ok_z, "smooth": ok_smooth, "corr": ok_corr},
    }


def main():
    from tinyvla.data.kinematics_so101 import SO101FK

    fk = SO101FK()
    results = {}
    roots = sorted(DATA_ROOT.iterdir())
    for r in roots:
        if not (r / "meta" / "info.json").exists():
            continue
        try:
            res = audit_one(r.name, r, fk)
        except Exception as e:
            res = {"verdict": "ERROR", "reason": f"{type(e).__name__}: {str(e)[:120]}"}
        results[r.name] = res
        print(f"{res['verdict']:6} {r.name}: {res}")
    OUT.write_text(json.dumps(results, indent=1))
    keep = sum(1 for v in results.values() if v["verdict"] == "KEEP")
    print(f"\nKEEP {keep}/{len(results)} -> {OUT}")


if __name__ == "__main__":
    main()