#!/usr/bin/env python3
"""Validate a converted xtrainer2 LeRobot dataset.

Two independent checks:

1. Decoded rotations must reproduce the official ``norm_stats.json`` in the
   20-D model space, which is what proves the quaternion convention is right.
2. StarVLA's dataloader must load a sample end to end.
"""

from __future__ import annotations

import argparse
import glob
import importlib.util
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from starVLA.dataloader.gr00t_lerobot.datasets import LeRobotSingleDataset
from starVLA.dataloader.gr00t_lerobot.transform.state_action import RotationTransform

REPO_ROOT = Path(__file__).resolve().parents[1]
OFFICIAL_STATS = "/apdcephfs_nj7/share_305204761/arianliu/data/benchmark/norm_stats.json"
ARM_BLOCKS = ((0, 3, 7), (8, 11, 15))


def load_data_config():
    """The examples tree is not a package, so load the registry by path."""
    path = REPO_ROOT / "examples/Xtrainer2/train_files/data_registry/data_config.py"
    spec = importlib.util.spec_from_file_location("_xtrainer2_data_config", path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module.Xtrainer2BimanualEefDataConfig()


def to_model_space(raw: np.ndarray, scramble_back_to_xyzw: bool = False) -> np.ndarray:
    """Expand stored 16-D poses into the 20-D layout used by the official stats."""
    transform = RotationTransform("quaternion", "rotation_6d")
    columns = []
    for position, rotation, gripper in ARM_BLOCKS:
        columns.append(raw[:, position : position + 3])
        block = raw[:, rotation : rotation + 4]
        if scramble_back_to_xyzw:
            block = block[:, [1, 2, 3, 0]]
        six = transform.forward(torch.from_numpy(np.ascontiguousarray(block)))
        columns.append(six.numpy())
        columns.append(raw[:, gripper : gripper + 1])
    return np.concatenate(columns, axis=1)


def _stats_error(model_space: np.ndarray, official: dict) -> tuple[float, np.ndarray]:
    ref_mean = np.array(official["state"]["mean"])[0]
    ref_lo = np.array(official["state"]["p_low"])[0]
    ref_hi = np.array(official["state"]["p_high"])[0]
    mean = model_space.mean(0)
    lo = np.quantile(model_space, 0.01, axis=0)
    hi = np.quantile(model_space, 0.99, axis=0)
    error = float(
        (np.abs(mean - ref_mean).mean() + np.abs(lo - ref_lo).mean() + np.abs(hi - ref_hi).mean())
        / 3
    )
    return error, mean


def read_states(dataset_root: Path) -> np.ndarray:
    frames = [
        pd.read_parquet(path)["observation.state"]
        for path in sorted(glob.glob(str(dataset_root / "data" / "chunk-*" / "*.parquet")))
    ]
    return np.stack([np.asarray(row, dtype=np.float32) for f in frames for row in f])


def check_convention(dataset_roots: list[Path]) -> tuple[float, float]:
    """Compare the stored convention against the wrong one, pooled over all tasks.

    The official stats cover the whole benchmark, so any single task is biased
    against them; only the pooled corpus is comparable, and only the gap between
    the two conventions is meaningful.
    """
    official = json.load(open(OFFICIAL_STATS))
    ref_mean = np.array(official["state"]["mean"])[0]
    raw = np.concatenate([read_states(root) for root in dataset_roots], axis=0)

    stored_error, mean = _stats_error(to_model_space(raw), official)
    wrong_error, _ = _stats_error(to_model_space(raw, scramble_back_to_xyzw=True), official)

    print(f"  pooled frames={len(raw)} over {len(dataset_roots)} datasets")
    print(f"  ours mean  {np.round(mean, 3)}")
    print(f"  ref  mean  {np.round(ref_mean, 3)}")
    print(f"  err as stored (wxyz)  : {stored_error:.4f}")
    print(f"  err if misread (xyzw) : {wrong_error:.4f}")
    return stored_error, wrong_error


def check_dataloader(dataset_root: Path) -> None:
    config = load_data_config()
    dataset = LeRobotSingleDataset(
        dataset_path=dataset_root,
        modality_configs=config.modality_config(),
        embodiment_tag=config.embodiment_tag,
        transforms=config.transform(),
        data_cfg={"lerobot_version": "v2.0", "action_mode": "abs", "include_state": True},
    )
    sample = dataset[0]
    print(f"  episodes={len(dataset.trajectory_lengths)}  samples={len(dataset)}")
    for key, value in sample.items():
        if isinstance(value, (np.ndarray, torch.Tensor)):
            print(f"  {key}: {tuple(value.shape)} {value.dtype}")
        else:
            print(f"  {key}: {value!r}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--root",
        type=Path,
        default=Path("/apdcephfs_nj7/share_305204761/arianliu/data/lerobot"),
    )
    parser.add_argument(
        "--min-margin",
        type=float,
        default=3.0,
        help="How many times better the stored convention must fit the official stats",
    )
    args = parser.parse_args()

    dataset_roots = sorted(p for p in args.root.iterdir() if p.is_dir())
    failures = []

    print("=== rotation convention vs official norm_stats ===")
    stored_error, wrong_error = check_convention(dataset_roots)
    margin = wrong_error / max(stored_error, 1e-9)
    print(f"  margin: {margin:.1f}x")
    if margin < args.min_margin:
        failures.append(f"convention margin {margin:.1f}x < {args.min_margin}x")

    for dataset_root in dataset_roots:
        print(f"\n=== {dataset_root.name}: StarVLA dataloader ===")
        check_dataloader(dataset_root)

    if failures:
        print("\nFAILED:")
        for failure in failures:
            print(f"  {failure}")
        return 1
    print("\nAll checks passed.")
    return 0


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