#!/usr/bin/env python3
"""Rewrite fruit_pick's observation.state from its native 15-dim layout
(eef_xyz[3] + eef_quat_xyzw[4] + joints[7] + gripper_width[1]) down to the
SAME 8-dim convention every other dataset in this repo uses (eef_pos[3] +
axis-angle rotation[3] + [gripper_width/2, -gripper_width/2]) -- per user
request, this training run must match place_cube_in_bowl/lift_new exactly.

Formula copied verbatim from
fastwam_train/lift2lerobot/convert_lift_hdf5_to_lerobot_v21.py::quat2axisangle
(itself from fastwam/experiments/libero/libero_utils.py / robosuite), so the
state convention is bit-identical to how place_cube_in_bowl and lift_new were
built -- NOT a from-scratch reimplementation.

Usage (run once, before staging/training):
    python fastwam_train/scripts/fix_fruit_pick_state_schema.py

Idempotent: skips if observation.state is already 8-dim. Backs up the
original parquet (.orig_15dim_state) and updates meta/info.json's declared
feature shape to match.
"""
import json
import math
import shutil
from pathlib import Path

import numpy as np
import pandas as pd

REPO = Path("/shared_work/amin-a/realworld-wam")
DATASET_DIR = REPO / "data/fruit_pick"
DATA_PARQUET = DATASET_DIR / "data/chunk-000/file-000.parquet"
INFO_JSON = DATASET_DIR / "meta/info.json"


def quat2axisangle(quat: np.ndarray) -> np.ndarray:
    """(x, y, z, w) vec4 -> (ax, ay, az) axis-angle exponential coordinates."""
    quat = quat.copy()
    if quat[3] > 1.0:
        quat[3] = 1.0
    elif quat[3] < -1.0:
        quat[3] = -1.0
    den = np.sqrt(1.0 - quat[3] * quat[3])
    if math.isclose(den, 0.0):
        return np.zeros(3)
    return (quat[:3] * 2.0 * math.acos(quat[3])) / den


def main():
    df = pd.read_parquet(DATA_PARQUET)
    sample_state = np.asarray(df["observation.state"].iloc[0])
    if sample_state.shape[-1] == 8:
        print(f"already 8-dim ({DATA_PARQUET}); nothing to do.")
        return
    assert sample_state.shape[-1] == 15, f"unexpected state dim {sample_state.shape[-1]}"

    backup = DATA_PARQUET.with_suffix(".parquet.orig_15dim_state")
    if not backup.exists():
        shutil.copy2(DATA_PARQUET, backup)
        print(f"backed up original -> {backup}")

    new_states = []
    for raw in df["observation.state"]:
        raw = np.asarray(raw, dtype=np.float64)
        eef_pos = raw[0:3]
        eef_quat_xyzw = raw[3:7]
        gripper_width = raw[14]
        ee_state = np.concatenate([eef_pos, quat2axisangle(eef_quat_xyzw)])
        gripper_state = np.array([gripper_width / 2.0, -gripper_width / 2.0])
        new_states.append(np.concatenate([ee_state, gripper_state]).astype(np.float32))

    df["observation.state"] = new_states
    df.to_parquet(DATA_PARQUET)
    print(f"rewrote {DATA_PARQUET.name}: observation.state 15-dim -> 8-dim ({len(df)} rows)")

    info = json.loads(INFO_JSON.read_text())
    info["features"]["observation.state"] = {
        "dtype": "float32",
        "shape": [8],
        "names": ["eef_x", "eef_y", "eef_z", "axis_x", "axis_y", "axis_z",
                   "gripper_pos", "gripper_neg"],
    }
    backup_info = INFO_JSON.with_suffix(".json.orig_15dim_state")
    if not backup_info.exists():
        shutil.copy2(INFO_JSON, backup_info)
    INFO_JSON.write_text(json.dumps(info, indent=4))
    print(f"updated {INFO_JSON} feature shape to [8]")


if __name__ == "__main__":
    main()