#!/usr/bin/env python3
"""
collect_rx1_lerobot.py — record RX1 stacking demonstrations as a LeRobotDataset
for fine-tuning SmolVLA (or any lerobot policy).

Each episode: two cubes are teleported to random reachable spots (optionally
with distractor cubes on the table), and the weld-servo controller from
scripts/stack_demo.py stacks one on the other. Every RECORD_EVERY-th control
step we record:

    observation.images.ego       (256×256 RGB, head camera — real render)
    observation.images.top       (256×256 RGB, top-down workspace camera)
    observation.state            (8,)  7 right-arm joint positions + weld state
    action                       (8,)  7 right-arm position targets + grasp bit
    task                         "stack the <a> cube on the <b> cube"

Action/state semantics match the runtime wrapper in rx1_vla_demo.py: the
non-arm joints are commanded to their current pose (compliant regime — see
models/expert.py for why the torso must lean), the right arm is driven by the
recorded targets, and the grasp bit drives the target cube's weld.

Only successful episodes (top cube settles on the base within tolerance) are
kept, matching scripts/collect_demos.py's convention.

    MUJOCO_GL=egl python scripts/collect_rx1_lerobot.py \
        --episodes 60 --root /path/to/datasets/rx1_stack
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import mujoco

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from env.rx1_env import RX1Env                       # noqa: E402
from models.rx1_ik import RX1IKSolver, CUBE_SIZE, RIGHT_ARM_ACT  # noqa: E402
import scripts.stack_demo as sd                       # noqa: E402

IMG = 256
FPS = 31                  # env control dt is 8 ms; we record every 4th step
RECORD_EVERY = 4
# Same reachable workspace as scripts/collect_demos.py, trimmed at the far-x
# end where table-height servoing stalls (see stack_demo.py's tower note).
SPAWN_X = (0.28, 0.34)
SPAWN_Y = (-0.18, -0.06)
# Distractors only need to be on the table (x 0.14-0.58, y ±0.4 in the scene),
# not inside the arm's reach box — constraining all four cubes to the small
# workspace makes ≥MIN_SEP placements nearly unsatisfiable.
DISTRACT_X = (0.20, 0.50)
DISTRACT_Y = (-0.30, 0.12)
TABLE_Z = 0.445
MIN_SEP = 0.08            # min cube separation so the hand never clips a neighbour
COLORS = ("red", "blue", "green", "yellow")
TORSO_YAW, TORSO_ROLL = 0, 2   # torso actuator indices (see models/rx1_ik.py)

FEATURES = {
    "observation.images.ego": {
        "dtype": "image", "shape": (IMG, IMG, 3),
        "names": ["height", "width", "channels"],
    },
    "observation.images.top": {
        "dtype": "image", "shape": (IMG, IMG, 3),
        "names": ["height", "width", "channels"],
    },
    "observation.state": {
        "dtype": "float32", "shape": (8,),
        "names": [f"right_arm_{i}" for i in range(7)] + ["grasp"],
    },
    "action": {
        "dtype": "float32", "shape": (8,),
        "names": [f"right_arm_{i}" for i in range(7)] + ["grasp"],
    },
}


class RecordingEnv:
    """Proxy around RX1Env: intercepts step() to record (obs_t, action_t)
    pairs at FPS while the stack_demo controller runs unmodified.

    DART-style exploration (`noise` > 0): the EXECUTED right-arm targets get
    Gaussian noise while the RECORDED label stays clean. The state drifts off
    the expert path and the closed-loop servo corrects it on the next step,
    so the dataset teaches recovery from drift — without it a cloned policy
    compounds its own errors (same trick as scripts/collect_demos.py)."""

    def __init__(self, env: RX1Env, noise: float = 0.0, rng=None):
        self._env = env
        self._r = mujoco.Renderer(env.model, IMG, IMG)
        self._n = 0
        self._bit = None          # action index of the target cube's weld bit
        self._eq = None           # equality id of that weld (live grasp state)
        self._noise = noise
        self._rng = rng or np.random.default_rng()
        self.frames: list[dict] = []

    def __getattr__(self, name):
        return getattr(self._env, name)

    def set_target(self, cube: str):
        self._bit = self._env._weld_bit[cube]
        self._eq = dict(self._env._cube_welds)[cube]
        self._n = 0
        self._welded_once = False
        self.frames = []

    def _rgb(self, camera: str) -> np.ndarray:
        self._r.update_scene(self._env.data, camera=camera)
        return self._r.render().copy()

    def step(self, action):
        if self._bit is not None and self._n % RECORD_EVERY == 0:
            qpos = self._env.data.qpos
            state = np.concatenate([
                qpos[RIGHT_ARM_ACT],
                [float(self._env.data.eq_active[self._eq])],
            ]).astype(np.float32)
            act = np.concatenate([
                np.asarray(action)[RIGHT_ARM_ACT],
                [float(np.asarray(action)[self._bit])],
            ]).astype(np.float32)
            self.frames.append({
                "observation.images.ego": self._rgb("ego"),
                "observation.images.top": self._rgb("top"),
                "observation.state": state,
                "action": act,
            })
        self._n += 1
        action = np.array(action, dtype=np.float32, copy=True)
        # Hold torso yaw/roll at zero: fully-compliant torso lets the arm's
        # weight drag the robot sideways over an episode (it visibly slumps
        # right), which shifts the arm base laterally. Pitch stays compliant —
        # the arm cannot reach the table from an upright back (models/expert.py).
        # The runtime wrapper (rx1_vla_demo.py --hold-back) applies the same
        # regime, so train and test dynamics match.
        action[TORSO_YAW] = 0.0
        action[TORSO_ROLL] = 0.0
        if self._noise > 0 and self._bit is not None:
            welded = action[self._bit] > 0.5
            self._welded_once |= bool(welded)
            # Perturb only during approach and carry. After the release the
            # cube is sitting placed — noise there just knocks it off and
            # teaches nothing (the episode is effectively over).
            if welded or not self._welded_once:
                action[RIGHT_ARM_ACT] += self._rng.normal(0, self._noise, 7)
        return self._env.step(action)


def _set_cube(env, cube, pos):
    j = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, f"{cube}_free")
    a = env.model.jnt_qposadr[j]
    env.data.qpos[a:a + 3] = pos
    env.data.qpos[a + 3:a + 7] = [1, 0, 0, 0]
    env.data.qvel[env.model.jnt_dofadr[j]:env.model.jnt_dofadr[j] + 6] = 0


def _sample_spots(rng, n_distractors):
    """[base, top] inside the reachable workspace plus distractors anywhere
    on the table, all ≥ MIN_SEP apart. Whole-set restarts, not greedy: a
    centrally-placed first spot can make the second unsatisfiable forever."""
    for _ in range(100):
        spots = []
        ok = True
        for i in range(2 + n_distractors):
            rx, ry = (SPAWN_X, SPAWN_Y) if i < 2 else (DISTRACT_X, DISTRACT_Y)
            for _ in range(40):
                p = np.array([rng.uniform(*rx), rng.uniform(*ry)])
                if all(np.linalg.norm(p - q) >= MIN_SEP for q in spots):
                    spots.append(p)
                    break
            else:
                ok = False
                break
        if ok:
            return spots
    return None


def collect(n_episodes: int, root: str, repo_id: str, seed: int,
            noise: float = 0.0, video: bool = False, append: bool = False):
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    if append:
        # Reopen an existing dataset and add episodes (verified: reopen +
        # add_frame/save_episode/finalize works on a finalized v3 dataset).
        # Set HF_DATASETS_IN_MEMORY_MAX_SIZE so the reopen doesn't materialize
        # a dataset-sized Arrow cache on disk.
        dataset = LeRobotDataset(repo_id, root=root)
        print(f"appending to {root}: {dataset.num_episodes} existing episodes")
    else:
        features = FEATURES
        if video:
            features = {k: ({**v, "dtype": "video"} if v["dtype"] == "image" else v)
                        for k, v in FEATURES.items()}
        dataset = LeRobotDataset.create(
            repo_id=repo_id, fps=FPS, features=features, root=root,
            robot_type="rx1", use_videos=video)

    rng = np.random.default_rng(seed)
    env = RecordingEnv(RX1Env(sd.ROBOT_CFG), noise=noise, rng=rng)
    ik = RX1IKSolver(env.model)
    view = sd._Viewer(env, False, 0)

    kept = ep = 0
    while kept < n_episodes:
        ep += 1
        env.reset()
        base, top, *rest = rng.permutation(COLORS)
        with_distractors = rng.random() < 0.7
        spots = _sample_spots(rng, 2 if with_distractors else 0)
        if spots is None:
            print(f"episode {ep}: spot sampling failed — retrying")
            continue
        _set_cube(env, f"{base}_cube", [*spots[0], TABLE_Z])
        _set_cube(env, f"{top}_cube", [*spots[1], TABLE_Z])
        for i, c in enumerate(rest):
            if with_distractors:
                _set_cube(env, f"{c}_cube", [*spots[2 + i], TABLE_Z])
            else:
                _set_cube(env, f"{c}_cube", [1.0 + 0.1 * i, 1.0, 0.1])   # off-table
        mujoco.mj_forward(env.model, env.data)
        env.data.ctrl[:env._nu] = env.data.qpos[:env._nu]

        task = f"stack the {top} cube on the {base} cube"
        env.set_target(f"{top}_cube")
        base_pos = env.perceive_objects()[f"{base}_cube"]["pos"]
        place = np.array([base_pos[0], base_pos[1], base_pos[2] + CUBE_SIZE])
        err, servo_err = sd.stack_one(env, ik, f"{top}_cube", place, view)

        settled = env.perceive_objects()[f"{top}_cube"]["pos"]
        ok = (np.linalg.norm(settled[:2] - place[:2]) <= sd.XY_TOL
              and abs(settled[2] - place[2]) <= sd.Z_TOL)
        print(f"episode {ep}: '{task}'  servo={servo_err * 100:.1f}cm "
              f"settled={err * 100:.1f}cm  frames={len(env.frames)}  "
              f"{'KEEP' if ok else 'DISCARD'}  ({kept + int(ok)}/{n_episodes})")
        if not ok:
            continue
        for frame in env.frames:
            dataset.add_frame({**frame, "task": task})
        dataset.save_episode()
        kept += 1

    dataset.finalize()
    print(f"\nsaved {kept} episodes ({ep - kept} discarded) -> {root}")


def main():
    ap = argparse.ArgumentParser(description="Collect RX1 stacking demos as a LeRobotDataset")
    ap.add_argument("--episodes", type=int, default=60)
    ap.add_argument("--root", required=True)
    ap.add_argument("--repo-id", default="local/rx1_stack")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--noise", type=float, default=0.0,
                    help="DART exploration noise std (rad) on executed arm targets")
    ap.add_argument("--video", action="store_true",
                    help="store camera streams as video instead of per-frame images")
    ap.add_argument("--append", action="store_true",
                    help="add episodes to an existing dataset at --root")
    args = ap.parse_args()
    collect(args.episodes, args.root, args.repo_id, args.seed, args.noise,
            args.video, args.append)


if __name__ == "__main__":
    main()