#!/usr/bin/env python3
"""
collect_demos.py  (v2) — generate behaviour-cloning demonstrations for the
learned 40-DoF action model.

Each episode: one cube is spawned at a random reachable spot, a random place
target is chosen (anywhere from table height up to a 3-cube stack height, and
over the box), and the privileged IK expert (models/expert.py) performs the
pick→place. Every control step we record:

    obs    = [ proprio(80) , pick_pos0(3) , place_pos(3) ]   # goal-conditioned
    action = [ joints(40)  , grasp(1) ]                       # what the expert did

Goal conditioning on (pick_pos0, place_pos) is what lets the cloned policy
generalise to pick / place-in-box / stack / sort — they are just different
place targets. Only successful episodes (final cube within tol of the target)
are kept.

    python scripts/collect_demos.py --n 400 --out data/demos.npz
"""
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,   # noqa: E402
                           RIGHT_ARM_ACT, LEFT_ARM_ACT)
from models.expert import PickPlaceExpert          # noqa: E402
from brain.perception import PerceptionModule       # noqa: E402

# Learned action space = arms only: 7 right + 7 left arm joints + 1 grasp bit.
# Torso/head are held at rest and the 18 finger joints follow the grasp bit, so
# the model never has to learn them. ARM_ACT indexes those 14 joints in ctrl.
ARM_ACT = list(RIGHT_ARM_ACT) + list(LEFT_ARM_ACT)   # 14 actuator indices

ROBOT_CFG = {
    "simulation": {"physics_timestep": 0.002, "n_substeps": 4},
    "vision": {"enabled": False},     # privileged exact state; no GL needed
    "actuator": {"max_joint_speed": 3.0, "control_clip": 2.5},
}

# Reachable workspace for the right arm (metres, world frame).
PICK_X = (0.28, 0.34); PICK_Y = (-0.20, -0.04); TABLE_Z = 0.445
PLACE_X = (0.27, 0.40); PLACE_Y = (-0.20, -0.02)


def _qadr(model, cube):
    j = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{cube}_free")
    return model.jnt_qposadr[j]


def _set_cube(env, cube, pos):
    a = _qadr(env.model, cube)
    env.data.qpos[a:a + 3] = pos
    env.data.qpos[a + 3:a + 7] = [1, 0, 0, 0]


def collect(n, out, seed=0, noise=0.04, base_frac=0.0):
    rng = np.random.default_rng(seed)
    env = RX1Env(ROBOT_CFG)
    ik = RX1IKSolver(env.model)
    nu = env._nu
    target_cube = "red_cube"
    target_bit = env._weld_bit[target_cube]
    others = [c for c in ("blue_cube", "green_cube", "yellow_cube")]
    # `base_frac` of episodes are STACKING demos: a base cube sits at the place
    # spot and the target is its top, so the model learns to place on top of a
    # cube (the contact + tighter precision stacking needs). The rest are
    # free-space pick/place/sort demos with the other cubes parked away.

    OBS, ACT, EP = [], [], []
    kept = 0
    for ep in range(n):
        env.reset()
        # park the unused cubes off-table so they never interfere or get perceived
        for i, c in enumerate(others):
            _set_cube(env, c, [1.0 + 0.1 * i, 1.0, 0.1])
        pick = np.array([rng.uniform(*PICK_X), rng.uniform(*PICK_Y), TABLE_Z])
        _set_cube(env, target_cube, pick)

        stacking = rng.random() < base_frac
        if stacking:
            # place a base cube on the table; the target is its top face
            base_xy = np.array([rng.uniform(*PLACE_X), rng.uniform(*PLACE_Y)])
            while np.linalg.norm(base_xy - pick[:2]) < 0.06:        # not under the pick
                base_xy = np.array([rng.uniform(*PLACE_X), rng.uniform(*PLACE_Y)])
            _set_cube(env, others[0], [base_xy[0], base_xy[1], TABLE_Z])
            place = np.array([base_xy[0], base_xy[1], TABLE_Z + CUBE_SIZE])
        else:
            place = np.array([rng.uniform(*PLACE_X), rng.uniform(*PLACE_Y),
                              TABLE_Z + rng.uniform(0.0, 3 * CUBE_SIZE)])
        mujoco.mj_forward(env.model, env.data)
        env.data.ctrl[:nu] = env.data.qpos[:nu]
        pick0 = env.perceive_objects()[target_cube]["pos"].copy()

        expert = PickPlaceExpert(ik, pick0, place)
        ep_obs, ep_act = [], []
        steps = 0
        while not expert.done and steps < 400:
            # Record proprio EXACTLY as the runtime policy will receive it: the
            # brain feeds the policy PerceptionModule._normalise_proprio(...), so
            # demos must match or the model sees out-of-distribution input.
            proprio = PerceptionModule._normalise_proprio(env._get_proprio())  # (80,)
            cube_live = env.perceive_objects()[target_cube]["pos"]
            action41, grasp = expert.act(env.data.qpos[:nu].copy(), cube_live)
            obs = np.concatenate([proprio, pick0, place]).astype(np.float32)
            # record only what the model outputs: 14 arm joints + grasp
            arm15 = np.concatenate([action41[ARM_ACT],
                                    [action41[nu]]]).astype(np.float32)
            ep_obs.append(obs); ep_act.append(arm15)
            # Drive the full body in sim. Inject exploration noise on the APPLIED
            # arm joints (DART): the state wanders off the expert's path and the
            # closed-loop expert corrects on the NEXT step, so the recorded
            # (state → clean action) pairs teach the policy to recover. This is
            # what makes the cloned policy robust to its own closed-loop drift.
            full = np.zeros(env.N_ACTUATORS, dtype=np.float32)
            full[:nu] = action41[:nu]
            full[target_bit] = action41[nu]               # grasp -> this cube's weld bit
            if noise > 0:
                full[ARM_ACT] += rng.normal(0, noise, len(ARM_ACT)).astype(np.float32)
            env.step(full)
            steps += 1

        final = env.perceive_objects()[target_cube]["pos"]
        err = float(np.linalg.norm(final - place))
        if err <= 0.06 and len(ep_obs) > 10:
            OBS.extend(ep_obs); ACT.extend(ep_act); EP.append(len(ep_obs))
            kept += 1
        if (ep + 1) % 10 == 0 or ep == n - 1:
            print(f"  episode {ep+1}/{n}  kept={kept}  last_err={err:.3f}m  steps={steps}")

    OBS = np.asarray(OBS, np.float32); ACT = np.asarray(ACT, np.float32)
    Path(out).parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out, obs=OBS, act=ACT, ep_len=np.asarray(EP, np.int32),
                        obs_dim=OBS.shape[1], act_dim=ACT.shape[1])
    print(f"\nsaved {OBS.shape[0]} transitions from {kept}/{n} successful episodes -> {out}")
    print(f"  obs_dim={OBS.shape[1]}  act_dim={ACT.shape[1]}")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--n", type=int, default=400)
    p.add_argument("--out", default="data/demos.npz")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--noise", type=float, default=0.04,
                   help="DART exploration noise (rad) on applied arm joints")
    p.add_argument("--base-frac", type=float, default=0.0,
                   help="fraction of episodes that are stacking demos (base cube "
                        "at the place spot, target = its top). Use ~0.5 to train "
                        "a stacking-capable model.")
    a = p.parse_args()
    collect(a.n, a.out, a.seed, a.noise, a.base_frac)