"""
Learned pick-and-place runtime for the demo2 arm (the DEPLOYED controller).

Loads the BC action-chunking model (arm_act.pt), and for a pick-place task:
perceive object + place pad -> obs [init arm qpos, pick, place] -> model ->
action chunk (CHUNK, 6) -> play OPEN-LOOP (set 5 arm ctrl + grasp each step).
Open-loop playback sidesteps BC covariate shift, exactly as the RX1 v2 runtime.

The scripted teacher (arm_pick_place.PickPlaceTeacher) is NOT used here — it only
generated the training data.
"""
from __future__ import annotations

import sys
from pathlib import Path
import numpy as np
import torch

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
from models.act_model import load_policy
from arm_pick_place import ArmIK, ARM_ACT
import grasp_planner as gp

CKPT = HERE / "model" / "arm_act.pt"
_GRASP_BIT = len(ARM_ACT)   # action col 5
_MAX_DCTRL = 3.0 * 0.002 * 5    # 3 rad/s over a 5-substep, 2ms-timestep control tick
_REACH_TOL = 0.04               # rad: arm "reached" this waypoint
_DWELL_CAP = 40                 # max control ticks per waypoint before forcing advance


class LearnedArmPolicy:
    """Open-loop learned pick-and-place over the demo2 arm."""

    def __init__(self, model, device="cpu"):
        self.model = model
        self.device = device
        self._policy, self._norm = load_policy(str(CKPT), device)
        self._ik = ArmIK(model)        # only for rest() + actuator indices
        self._chunk = None             # cached predicted trajectory (T,6)
        self._cursor = 0
        self._dwell = 0                # control ticks spent on the current waypoint
        self.phase = "idle"
        self.plan = None

    @staticmethod
    def available():
        return CKPT.exists()

    def plan_pickplace(self, data, pick_xyz, place_xyz, pick_body="object",
                       color=None):
        """Query the learned model once -> full open-loop action chunk."""
        arm_q = np.array([data.ctrl[a] for a in ARM_ACT], np.float32)
        obs = np.concatenate([arm_q, pick_xyz, place_xyz]).astype(np.float32)
        with torch.no_grad():
            out = self._policy(torch.from_numpy(self._norm.norm_obs(obs))
                               .float().unsqueeze(0).to(self.device))
        flat = out.squeeze(0).cpu().numpy().reshape(-1)
        chunk = self._norm.denorm_act(flat).reshape(-1, 1 + _GRASP_BIT)  # (T,6)
        self._chunk = chunk
        self._cursor = 0
        self._dwell = 0
        self._pick_body = pick_body        # weld targets the object being picked
        self.phase = "executing"
        self.plan = {
            "color": color,
            "pick": [round(float(x), 3) for x in pick_xyz],
            "place": [round(float(x), 3) for x in place_xyz],
            "chunk_len": int(len(chunk)),
        }
        return self.plan

    def step(self, model, data):
        """Advance one control step: slew the arm toward the current chunk target
        and advance the cursor only once it's reached. Returns (done, grasp_bool).

        The chunk is ~100 rows resampled from a much longer teacher trajectory, so
        replaying one row per control tick (no slew) drives the arm ~5x too fast
        and flings the object. Slewing toward each target — and advancing only on
        arrival — reproduces the smooth, self-paced motion the teacher used.
        """
        if self._chunk is None or self._cursor >= len(self._chunk):
            self.phase = "done"
            return True, False
        row = self._chunk[self._cursor]
        grasp = bool(row[_GRASP_BIT] > 0.5)

        prev = data.ctrl[:model.nu].copy()
        target = prev.copy()
        for k, a in enumerate(ARM_ACT):
            target[a] = row[k]
        target = self._ik.set_fingers(target, grasp)
        # slew-rate limit (≈ teacher's 3 rad/s over a 5-substep control tick)
        ctrl = prev + np.clip(target - prev, -_MAX_DCTRL, _MAX_DCTRL)
        data.ctrl[:model.nu] = ctrl
        gp.set_weld(model, data, eq_name="grasp_weld", active=grasp,
                    body2=getattr(self, "_pick_body", "object"))

        # advance to the next waypoint once the arm joints have essentially
        # reached this one (or after a dwell cap, so a hard target can't stall).
        arm_err = max(abs(ctrl[a] - row[k]) for k, a in enumerate(ARM_ACT))
        self._dwell += 1
        if arm_err < _REACH_TOL or self._dwell >= _DWELL_CAP:
            self._cursor += 1
            self._dwell = 0
        prog = self._cursor / max(1, len(self._chunk))
        self.plan = {**(self.plan or {}), "progress": round(prog, 3),
                     "grasp": grasp}
        return False, grasp


def make_arm_policy(model, device="cpu"):
    if LearnedArmPolicy.available():
        return LearnedArmPolicy(model, device)
    return None