#!/usr/bin/env python3
"""Check that a gr00t_views dataset's `action` column really is in the LeRobot order its
meta/modality.json claims.

RoboCasa's simulator emits actions as
    [eef_pos(3), eef_rot(3), gripper(1), base(3), torso(1), base_mode(1)]
while a gr00t_views dataset declares
    [base_motion(4), control_mode(1), eef_pos(3), eef_rot(3), gripper_close(1)]

Any builder that copies simulator actions through without reordering produces a dataset whose
numbers are right but whose columns are mislabelled. Training then puts arm motion on the
base-motion channel and a constant on the gripper channel, so the policy drives the robot away
and never closes the gripper. The loss still converges -- nothing errors -- which is why this
has to be checked explicitly.

The two layouts are told apart by which dimensions are constant, no reference dataset needed:

    LeRobot    dims 0-3 all zero, dim 4 constant, dim 11 two-valued (the gripper)
    robosuite  dims 0-5 continuous, dim 6 two-valued (the gripper), dims 7-10 zero, dim 11 constant

    python check_action_layout.py --dataset <gr00t_views dataset> [--episodes 20]
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys

import numpy as np
import pyarrow.parquet as pq

LEROBOT_NAMES = ["base_x", "base_y", "base_th", "torso", "ctrl_mode",
                 "eef_x", "eef_y", "eef_z", "rot_x", "rot_y", "rot_z", "gripper"]
ROBOSUITE_NAMES = ["eef_x", "eef_y", "eef_z", "rot_x", "rot_y", "rot_z",
                   "gripper", "base_x", "base_y", "base_th", "torso", "base_mode"]
# robosuite index -> LeRobot position; apply as actions[:, FIX] to repair a dataset.
FIX = [7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5, 6]


def load_actions(dataset: str, n_episodes: int) -> np.ndarray:
    paths = sorted(glob.glob(os.path.join(dataset, "data", "chunk-*", "*.parquet")))
    if not paths:
        sys.exit(f"no parquet under {dataset}/data/chunk-*/")
    rows = []
    for p in paths[:n_episodes]:
        rows.append(np.stack(pq.read_table(p, columns=["action"]).column("action").to_pylist()))
    return np.vstack(rows), len(paths), min(n_episodes, len(paths))


def classify(a: np.ndarray) -> tuple[str, list[str]]:
    """Which layout do the constant/continuous dimensions match?"""
    const = [i for i in range(a.shape[1]) if len(np.unique(np.round(a[:, i], 4))) == 1]
    zero = [i for i in const if abs(a[:, i]).max() < 1e-6]
    twoval = [i for i in range(a.shape[1]) if len(np.unique(np.round(a[:, i], 4))) == 2]
    notes = [f"  all-zero dims : {zero}", f"  constant dims : {const}",
             f"  two-valued    : {twoval}  (the gripper is two-valued)"]
    lerobot = set([0, 1, 2, 3]).issubset(zero) and 4 in const and 11 in twoval
    robosuite = set([7, 8, 9, 10]).issubset(zero) and 11 in const and 6 in twoval
    if lerobot and not robosuite:
        return "lerobot", notes
    if robosuite and not lerobot:
        return "robosuite", notes
    return "unknown", notes


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--dataset", required=True, help="gr00t_views dataset directory")
    ap.add_argument("--episodes", type=int, default=20, help="episodes to sample (default 20)")
    a = ap.parse_args()

    actions, n_total, n_read = load_actions(a.dataset, a.episodes)
    print(f"dataset : {a.dataset}")
    print(f"sampled : {n_read} of {n_total} episodes, {actions.shape[0]} frames x {actions.shape[1]} dims\n")

    if actions.shape[1] != 12:
        print(f"FAIL  expected 12 action dims, found {actions.shape[1]}")
        return 2

    verdict, notes = classify(actions)
    names = LEROBOT_NAMES if verdict == "lerobot" else ROBOSUITE_NAMES
    for i in range(12):
        col = actions[:, i]
        uniq = len(np.unique(np.round(col, 4)))
        tag = "  <- constant" if uniq == 1 else ("  <- two-valued" if uniq == 2 else "")
        label = names[i] if verdict != "unknown" else "?"
        print(f"  dim{i:2d} {label:10s} {col.min():8.3f} .. {col.max():7.3f}   "
              f"|mean| {abs(col).mean():6.3f}   uniq {uniq:5d}{tag}")
    print()
    print("\n".join(notes))
    print()

    modality = os.path.join(a.dataset, "meta", "modality.json")
    if os.path.exists(modality):
        act = json.load(open(modality)).get("action", {})
        print("  modality.json action slices: "
              + ", ".join(f"{k} [{v['start']}:{v['end']}]" for k, v in act.items()))
        print()

    if verdict == "lerobot":
        print("PASS  action column is in LeRobot order, matching modality.json.")
        return 0
    if verdict == "robosuite":
        print("FAIL  action column is in RoboCasa/robosuite order but modality.json declares")
        print("      LeRobot order. Training would read arm motion as base motion and a")
        print("      constant as the gripper.")
        print(f"      Repair: actions = actions[:, {FIX}]")
        print("      Fix the builder and rebuild; do not patch the parquet in place unless you")
        print("      also confirm nothing else consumed it.")
        return 1
    print("UNKNOWN  neither layout matched. Inspect the dimensions above by hand:")
    print("      LeRobot expects dims 0-3 zero, dim 4 constant, dim 11 two-valued.")
    print("      An all-constant or degenerate episode sample can also land here --")
    print("      re-run with a larger --episodes before concluding.")
    return 3


if __name__ == "__main__":
    raise SystemExit(main())