from copy import deepcopy
import dataclasses
import io
import os
import pathlib
import time

import numpy as np
from PIL import Image
from scipy.spatial.transform import Rotation


os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")

JPEG_QUALITY = 90
SIDES = ("left", "right")
POSE_KEY = "observation.pose.left_hand_root_to_right_hand_root.absolute"
_K = ("observation.image.left", "observation.image.right", POSE_KEY,
      "observation.joint_states", "prompt")
CUP_PA1 = "Pick up the cup with your right hand and set it on the plate"
CUP_PA2 = "Pick up the cup with your left hand and return the cup to its original position"
CUP_HANDS = {CUP_PA1: ("right", 15), CUP_PA2: ("left", 14)}
CUP_ADVANCE = {CUP_PA1: 8, CUP_PA2: 4}
CUP_NOISE = {CUP_PA1: .15, CUP_PA2: .10}
PEN_PROMPTS = (
    "Pick up a pencil on the table and put it into the pen holder",
    "Pick up a ballpoint pen on the table and put it into the pen holder",
)


def _asset_id(ckpt):
    found = sorted(p.parent.name for p in (ckpt / "assets").glob("*/norm_stats.json"))
    if len(found) != 1:
        raise FileNotFoundError(f"expected one assets/<id>/norm_stats.json in {ckpt}, found {found}")
    return found[0]


def _m(r):
    a, b = r[..., :3], r[..., 3:6]
    c0 = a / np.linalg.norm(a, axis=-1, keepdims=True)
    b = b - np.sum(c0 * b, axis=-1, keepdims=True) * c0
    c1 = b / np.linalg.norm(b, axis=-1, keepdims=True)
    return np.stack([c0, c1, np.cross(c0, c1)], axis=-1)


def _d(x):

    rot = Rotation.from_matrix(_m(x[:, 3:9]))
    pr = Rotation.concatenate([Rotation.identity(), rot[:-1]])
    pp = np.concatenate([np.zeros((1, 3)), x[:-1, :3]])
    return np.concatenate([pr.inv().apply(x[:, :3] - pp), (pr.inv() * rot).as_quat(canonical=True)], axis=1)


def _image(img):

    from openpi.shared import image_tools
    a = np.asarray(image_tools.resize_with_pad(np.asarray(img, dtype=np.uint8), 224, 224))
    b = io.BytesIO(); Image.fromarray(a).save(b, format="JPEG", quality=JPEG_QUALITY)
    return np.asarray(Image.open(io.BytesIO(b.getvalue())).convert("RGB"))


def _build_transforms():
    from openpi import transforms as T
    from openpi.training.config import DataConfigFactory, ModelTransformFactory

    @dataclasses.dataclass(frozen=True)
    class _P(T.DataTransformFn):
        def __call__(self, d):
            return {k: d[k] for k in _K if k in d}

    @dataclasses.dataclass(frozen=True)
    class _I(T.DataTransformFn):
        n: int = 32

        def __call__(self, d):
            l, r = _image(d[_K[0]]), _image(d[_K[1]])
            p = np.array(d[_K[2]], dtype=np.float64)
            if p[6] < 0:
                p[3:] *= -1
            s = np.concatenate([p, np.asarray(d[_K[3]], dtype=np.float64)]).astype(np.float32)
            o = {"state": T.pad_to_dim(s, self.n),
                 "image": {"base_0_rgb": np.zeros_like(l), "left_wrist_0_rgb": l, "right_wrist_0_rgb": r},
                 "image_mask": {"base_0_rgb": np.False_, "left_wrist_0_rgb": np.True_, "right_wrist_0_rgb": np.True_}}
            if "prompt" in d:
                o["prompt"] = d["prompt"]
            return o

    @dataclasses.dataclass(frozen=True)
    class _O(T.DataTransformFn):
        def __call__(self, d):
            a = np.asarray(d["actions"], dtype=np.float64)
            return {"actions": np.concatenate([_d(a[:, 0:9]), _d(a[:, 9:18]), a[:, 18:20]], axis=1).astype(np.float32)}

    @dataclasses.dataclass(frozen=True)
    class _C(DataConfigFactory):
        def create(self, assets_dirs, model_config):
            return dataclasses.replace(
                self.create_base_config(assets_dirs, model_config),
                repack_transforms=T.Group(inputs=[_P()]),
                data_transforms=T.Group(inputs=[_I(model_config.action_dim)], outputs=[_O()]),
                model_transforms=ModelTransformFactory()(model_config))

    return _C


class CupContinuousAdvance:
    def __init__(self, advance_rows):
        if type(advance_rows) is not int or not 1 <= advance_rows <= 16:
            raise ValueError('Continuous place advance must be in 1..16 rows')
        self.advance_rows = advance_rows
        self.state = dict(prompt=None, clock=0, saw_open=False, armed=False, deadline=None,
                          waveform=None, latched=False, cursor=0, done=False)
        self._generation = 0
        self._owner = object()

    def propose(self, actions, obs, adopt_rows):
        a = np.asarray(actions)
        if a.shape != (32, 16) or a.dtype != np.float32 or not np.isfinite(a).all():
            raise ValueError('Continuous place advance expects finite float32[32,16]')
        if type(adopt_rows) is not int or not 1 <= adopt_rows <= 32:
            raise ValueError('Continuous place advance requires 1..32 adopted rows')
        prompt = obs.get('prompt', '').strip().rstrip('.')
        if prompt not in CUP_HANDS:
            raise ValueError('Continuous place advance requires task1 PA1/PA2')
        fingers = np.asarray(obs['observation.joint_states'], dtype=float)
        if fingers.shape != (2,) or not np.isfinite(fingers).all():
            raise ValueError('Expected two finite measured finger angles')
        side, column = CUP_HANDS[prompt]
        hand = column - 14
        state = deepcopy(self.state)
        if state['prompt'] != prompt:
            state = dict(prompt=prompt, clock=0, saw_open=False, armed=False, deadline=None,
                         waveform=None, latched=False, cursor=0, done=False)
        output = a.copy()
        if not state['done'] and not state['armed']:
            if fingers[hand] > .53:
                state['saw_open'] = True
            elif state['saw_open'] and 3 * 2 * np.pi / 4096 <= fingers[hand] <= .53:
                state['armed'] = True

        raw = a[:, column]


        threshold = min(.53, float(raw[0]) + .05)
        opening = None
        if np.any(raw >= .56):
            for i in range(31):
                if raw[i] >= threshold and raw[i + 1] >= threshold:
                    opening = i + 1
                    break

        if state['armed'] and fingers[hand] >= .56 and raw[0] >= .56:
            state.update(latched=False, done=True, deadline=None, waveform=None)

        start = None
        if state['armed'] and not state['latched'] and not state['done']:
            if opening is None:
                state.update(deadline=None, waveform=None)
            else:
                if state['deadline'] is None:
                    state['deadline'] = state['clock'] + opening
                    state['waveform'] = raw[opening - 1:].tolist()
                desired = state['deadline'] - state['clock'] - self.advance_rows
                candidate = max(desired, 1)
                if candidate <= adopt_rows and candidate < opening:
                    start = candidate

        natural_start = (opening if state['armed'] and not state['done']
                         and not state['latched'] and start is None
                         and opening is not None and opening <= adopt_rows else None)

        report = dict(enabled=True, kind='cup_continuous_place_advance_v1',
                      hand=side, advance_rows=self.advance_rows, clock=state['clock'],
                      forecast_opening_row=opening, opening_threshold_rad=threshold,
                      anchored_opening_row=state['deadline'], forced_start_row=start,
                      natural_start_row=natural_start,
                      already_latched=state['latched'], modified=False,
                      waveform='continuous values from the first reserved forecast',
                      pick_advance=False, height_or_robot_state_used=False)
        release_row = start or natural_start
        hold_end = ((release_row - 1) // 4 + 1) * 4 if release_row else 0
        states = []
        for index in range(adopt_rows):
            if start == index + 1:
                state.update(latched=True, cursor=0, open_floor=float(fingers[hand]))
            if state['latched']:
                wave = state['waveform']
                if not wave:
                    raise RuntimeError('Missing reserved continuous opening waveform')
                value = max(float(fingers[hand]), state.get('open_floor', 0.),
                            wave[min(state['cursor'], len(wave)-1)])
                output[index, column] = value
                state.update(cursor=state['cursor']+1, open_floor=value)
                report['modified'] = True
            if release_row and release_row <= index + 1 <= hold_end:
                output[index, :14] = np.tile([0., 0., 0., 0., 0., 0., 1.], 2)
                report['modified'] = True
            state['clock'] += 1
            states.append(deepcopy(state))
        return dict(actions=output, states=states, report=report,
                    adopt_rows=adopt_rows, _owner=self._owner, _generation=self._generation)

    def commit(self, proposal, executed_rows):
        if proposal.get('_owner') is not self._owner or proposal.get('_generation') != self._generation:
            raise ValueError('Stale continuous place advance proposal')
        if type(executed_rows) is not int or not 0 <= executed_rows <= proposal['adopt_rows']:
            raise ValueError('Invalid verified executed prefix')
        if executed_rows:
            self.state = deepcopy(proposal['states'][executed_rows - 1])
            self._generation += 1


PEN_POLICY = dict(kind='u1_pen_delayed_bidirectional_32_retention_narrow_retry_v2',
    lookahead_rows=16, retained_forecast_rows=32, command_row_1based=8,
    delay_completed_chunks=1, close_below_rad=.07, release_at_or_above_rad=.27,
    preclose_open_range_rad=[.07, .27], close_cancel_opening_delta_rad=.02,
    close_target_rad=0., open_target='first_remaining_raw_release_value',
    latch_max_completed_chunks=2, latch_release='natural_prefix_catchup_or_countermand_or_forecast_withdrawal_or_bound',
    initial_closed_is_verified_grasp=False, narrow_retry='current_not_holding_plus_future_close_and_no_pending_open_or_open_latch_or_prefix_open', state_commit='verified_full_eight_rows_only',
    raw_four_stage_widths_preserved=True, ground_truth_inputs=False)


def policy_for_rows(rows=8, *, place_advance_rows=0, close_lookahead=True, close_defer_only=False):
    if type(rows) is not int or rows not in (8, 16):
        raise ValueError('U1 pen lookahead execution rows must be8 or16')
    if type(close_defer_only) is not bool or (close_defer_only and (not close_lookahead or rows != 8)):
        raise ValueError('Confirmed defer-only requires close lookahead enabled and adopt8')
    if type(close_lookahead) is not bool:
        raise ValueError('close_lookahead must be boolean')
    if type(place_advance_rows) is not int or not 0 <= place_advance_rows <= 16:
        raise ValueError('Pen place advance must be an integer in0..16')
    policy = (deepcopy(PEN_POLICY) if rows == 8 else
        dict(deepcopy(PEN_POLICY), kind='u1_pen_delayed_bidirectional_32_retention_adopt16_v3',
             lookahead_rows=32, command_row_1based=16, state_commit='verified_full_sixteen_rows_only'))
    if place_advance_rows:
        policy.update(kind=policy['kind']+'_place_advance_v1', place_advance_rows=place_advance_rows,
            release_rule='first raw >=0.27 forecast absolute executed-row deadline minus advance; earliest current row1',
            release_target='first current raw opening value; existing bounded open latch',
            release_forecast_rows=32, release_withdrawal_cancels=True,
            pick_advance_changed=False, arm_actions_changed=False)
    if not close_lookahead:
        policy.update(kind=policy['kind']+'_open_only_v1', close_lookahead_enabled=False,
            close_rule='raw model timing only; no forced-close reservation or latch',
            normal_model_closing_preserved=True)
    if close_defer_only:
        policy.update(kind=policy['kind']+'_confirmed_close_defer_v1', close_defer_only=True,
            close_defer_near_rows=12, close_defer_confirmations=2, close_defer_reserve_rows=16,
            close_defer_deadline='first eligible raw close forecast; absolute verified executed-row clock',
            close_defer_trigger='deadline already expired, raw close within first12 and delayed for two consecutive completed observations',
            close_defer_command_row=rows, normal_model_closing_preserved=True)
    return policy


class PenGraspSchedule:
    def __init__(self, rows=8, *, place_advance_rows=0, close_lookahead=True, close_defer_only=False):
        self.rows = rows
        self.policy = policy_for_rows(rows, place_advance_rows=place_advance_rows, close_lookahead=close_lookahead, close_defer_only=close_defer_only)
        self.close_lookahead = close_lookahead
        self.close_defer_only = close_defer_only
        self.close_defer_state = {s: None for s in SIDES}
        self.place_advance_rows = place_advance_rows
        self.release_deadlines = {s: None for s in SIDES}
        self.executed_row_clock = 0
        self.pending = {s: None for s in SIDES}
        self.phases = {s: None for s in SIDES}
        self.verified_close = {s: False for s in SIDES}
        self.seen_open = {s: False for s in SIDES}
        self.latches = {s: None for s in SIDES}
        self.last_committed_cycle = -1
        self._generation = 0
        self._owner = id(self)

    def propose(self, cycle, predictions, targets, measured_angles, readings, limits, *, allow_clipping):
        prediction = np.asarray(predictions)
        selected = np.asarray(targets, dtype=float).copy()
        measured = np.asarray(measured_angles, dtype=float)
        if (type(cycle) is not int or cycle <= self.last_committed_cycle
                or prediction.shape != (32, 16) or prediction.dtype.kind != 'f'
                or not np.isfinite(prediction).all() or selected.shape != (self.rows, 2)
                or not np.isfinite(selected).all() or measured.shape != (2,)
                or not np.isfinite(measured).all()):
            raise ValueError('Invalid U1 pen transition inputs/cycle')
        next_close_defer = deepcopy(self.close_defer_state)
        next_release_deadlines = deepcopy(self.release_deadlines)
        next_pending = {s: None for s in SIDES}
        next_phases = deepcopy(self.phases)
        next_verified = deepcopy(self.verified_close)
        next_seen = deepcopy(self.seen_open)
        next_latches = {s: None for s in SIDES}
        report = dict(enabled=True, policy=deepcopy(self.policy), cycle=cycle,
            state_committed=False, pending_before=deepcopy(self.pending),
            place_advance_rows=self.place_advance_rows, executed_row_clock=self.executed_row_clock,
            close_defer_only=self.close_defer_only, close_defer_before=deepcopy(self.close_defer_state),
            release_deadlines_before=deepcopy(self.release_deadlines),
            phases_before=deepcopy(self.phases), latches_before=deepcopy(self.latches),
            verified_close_before=deepcopy(self.verified_close), hands={},
            _owner=self._owner, _generation=self._generation)
        for col, side in enumerate(SIDES):
            values = prediction[:, 14 + col]
            adopted = selected[:, col].copy()
            angle = float(measured[col])
            holding = readings.get(side, {}).get('contact_force_hold') is True
            verified = bool(self.verified_close[side] or holding)
            seen_open = bool(self.seen_open[side] or angle >= .07)
            if angle >= .27 and not holding:
                verified = False
            state = ('closed_holding' if holding else 'wide_open' if angle >= .27
                     else 'preclose_open' if angle >= .07 and not verified
                     else 'release_started' if angle >= .07
                     else 'closed' if verified else 'initial_closed_unverified')

            phase = ('closed' if holding or (angle < .07 and verified)
                     else 'release_started' if verified else 'wide_open' if angle >= .27
                     else 'preclose_open' if angle >= .07 else 'initial_closed_unverified')
            closes = np.flatnonzero(values < .07)
            opens = np.flatnonzero(values >= .27)
            close16 = closes[closes < self.policy['lookahead_rows']]
            open16 = opens[opens < self.policy['lookahead_rows']]
            prefix_close = np.flatnonzero(adopted < .07)
            prefix_open = np.flatnonzero(adopted >= .27)
            open_then_close = bool(len(prefix_open) and np.any(adopted[prefix_open[0]+1:] < .07))


            pending_before = self.pending[side]
            latch_before = self.latches[side]
            retry_preclose = bool(verified and not holding and .07 <= angle < .27
                and len(closes) and not len(prefix_close) and not len(prefix_open)
                and not (pending_before and pending_before['direction'] == 'open')
                and not (latch_before and latch_before['direction'] == 'open')
                and (not len(opens) or closes[0] < opens[0]))
            if retry_preclose:
                phase = 'preclose_retry'
            hand = dict(measured_angle_rad=angle, contact_force_hold=holding,
                source_state=state, source_phase=phase, verified_close_history=verified, retry_after_verified_close=retry_preclose,
                close_prediction_rows_1based=(closes + 1).tolist(),
                open_prediction_rows_1based=(opens + 1).tolist(),
                original_selected_row_rad=float(adopted[self.rows-1]), selected_row_rad=float(adopted[self.rows-1]),
                model_row_rad=float(values[self.rows-1]), command_row_1based=self.rows, event='none', override=False,
                scheduled_in_adopted_plan=False, calibrated_clipping=False,
                override_rows_1based=[], latch_event='none', close_lookahead_enabled=self.close_lookahead)

            hand[f'original_selected_row{self.rows}_rad'] = float(adopted[self.rows-1])
            hand[f'selected_row{self.rows}_rad'] = float(adopted[self.rows-1])
            hand[f'model_row{self.rows}_rad'] = float(values[self.rows-1])

            def apply(rows, requested):
                lo, hi = (float(limits[side][k]) for k in ('min_angle_rad', 'max_angle_rad'))
                if not np.isfinite([lo, hi]).all() or lo > hi:
                    raise ValueError('Invalid measured gripper limits')
                if not lo <= requested <= hi and not allow_clipping:
                    raise ValueError(side + ': pen transition target outside calibrated limits')
                target = float(np.clip(requested, lo, hi))
                selected[rows, col] = target
                hand.update(override=True, requested_target_rad=float(requested),
                    calibrated_clipping=bool(hand['calibrated_clipping'] or target != requested))
                hand['override_rows_1based'].extend((np.arange(self.rows)[rows] + 1).tolist())
                return target

            pending = self.pending[side]
            latch = self.latches[side]
            if not self.close_lookahead:
                if pending is not None and pending['direction'] == 'close':
                    pending = None
                    hand['discarded_close_reservation'] = True
                if latch is not None and latch['direction'] == 'close':
                    latch = None
                    hand['discarded_close_latch'] = True
            if latch is not None:
                direction = latch['direction']
                remaining = closes if direction == 'close' else opens
                catchup = prefix_close if direction == 'close' else prefix_open
                countermand = bool(len(prefix_open)) if direction == 'close' else open_then_close
                if countermand:
                    hand['latch_event'] = 'released_explicit_countermand'
                elif not len(remaining):
                    hand['latch_event'] = 'released_forecast_withdrawn'
                elif latch['completed_chunks'] >= self.policy['latch_max_completed_chunks']:
                    hand['latch_event'] = 'released_bounded_duration'
                else:
                    stop = int(catchup[0]) if len(catchup) else self.rows
                    if stop:
                        apply(slice(0, stop), latch['target_rad'])
                    hand['latch_event'] = 'natural_prefix_catchup' if len(catchup) else 'maintain_short_latch'
                    if not len(catchup):
                        next_latches[side] = dict(latch, completed_chunks=latch['completed_chunks'] + 1)

                hand['event'] = 'existing_latch'
            elif pending is not None:
                direction = pending['direction']
                hand.update(direction=direction, detected_cycle=pending['detected_cycle'], due_cycle=pending['due_cycle'])
                if cycle != pending['due_cycle']:
                    hand['event'] = 'cancelled_missed_due_cycle'
                elif direction == 'close' and (holding or (angle < .07 and verified)):
                    hand['event'] = 'already_closed_measured'
                elif direction == 'open' and state == 'wide_open':
                    hand['event'] = 'already_open_measured'
                elif direction == 'close' and phase not in ('preclose_open', 'preclose_retry'):
                    hand['event'] = 'cancelled_source_state_changed'
                elif direction == 'open' and not verified:
                    hand['event'] = 'cancelled_source_state_changed'
                elif not len(closes if direction == 'close' else opens):
                    hand['event'] = 'cancelled_prediction_retracted_32'
                elif direction == 'close' and np.any(adopted > pending['anchor_angle_rad'] + .02):
                    hand['event'] = 'cancelled_opening_command_before_deadline'
                elif direction == 'open' and (open_then_close or (phase == 'release_started' and len(prefix_close))):
                    hand['event'] = 'cancelled_closing_command_before_deadline'
                elif direction == 'close' and len(prefix_close):
                    hand['event'] = 'normal_policy_close_in_chunk'
                elif direction == 'open' and len(prefix_open):
                    hand['event'] = 'normal_policy_open_in_chunk'
                else:
                    requested = 0. if direction == 'close' else float(values[opens[0]])
                    actual = apply(slice(self.rows-1, self.rows), requested)
                    hand['event'] = 'apply_due_transition'
                    next_latches[side] = dict(direction=direction, target_rad=actual, completed_chunks=0)
            else:
                direction = None
                if self.close_lookahead and not self.close_defer_only and phase in ('preclose_open', 'preclose_retry') and len(close16):
                    if np.any(adopted > angle + .02):
                        hand['event'] = 'not_armed_opening_command'
                    elif len(prefix_close):
                        hand['event'] = 'normal_policy_close_in_chunk'
                    else:
                        direction = 'close'
                elif not self.place_advance_rows and verified and (holding or angle < .27) and len(open16):
                    if open_then_close:
                        hand['event'] = 'not_armed_open_then_close'
                    elif len(prefix_open):
                        hand['event'] = 'normal_policy_open_in_chunk'
                    else:
                        direction = 'open'
                if direction:
                    next_pending[side] = dict(direction=direction, detected_cycle=cycle,
                        due_cycle=cycle+1, anchor_angle_rad=angle)
                    hand.update(event='arm_next_chunk', direction=direction,
                        detected_cycle=cycle, due_cycle=cycle+1)

            if self.close_defer_only:
                first_close = int(closes[0])+1 if len(closes) else None
                check = dict(enabled=True, first_close_row=first_close, near_rows=32,
                    confirmations_required=1, previous=deepcopy(self.close_defer_state[side]),
                    modified=False, event='ineligible')
                eligible = bool(phase in ('preclose_open','preclose_retry') and not holding
                    and len(closes) and not len(prefix_close) and not len(prefix_open)
                    and not (pending_before and pending_before['direction']=='open')
                    and not (latch_before and latch_before['direction']=='open')
                    and not np.any(adopted > angle+.02)
                    and (not len(opens) or closes[0] < opens[0]))


                if not eligible or latch is not None:
                    next_close_defer[side] = None
                    check['event'] = 'natural_close_in_prefix' if len(prefix_close) else 'reset_ineligible_or_latch'
                else:
                    old = self.close_defer_state[side]
                    deferred = deepcopy(old) if old is not None else None
                    if deferred is None and first_close <= 32:
                        deferred = dict(deadline=self.executed_row_clock+first_close, confirmations=0)
                        check['event'] = 'reserve_first_close_deadline'
                    if deferred is not None:
                        expected = deferred['deadline']-self.executed_row_clock
                        expired = self.executed_row_clock >= deferred['deadline']
                        near = first_close <= 32
                        delayed = first_close >= expected+2
                        confirm = bool(expired and near and delayed)
                        deferred['confirmations'] = deferred['confirmations']+1 if confirm else 0
                        check.update(deadline=deferred['deadline'], expected_close_row=expected,
                            deadline_expired=expired, near_close=near, delayed=delayed,
                            confirmations=deferred['confirmations'], event='observe_close_deferral')
                        if deferred['confirmations'] >= 1:
                            actual = apply(slice(self.rows-1,self.rows),0.)
                            hand.update(event='apply_confirmed_deferred_close',direction='close')
                            check.update(event='apply_confirmed_deferred_close',modified=True,
                                command_row_1based=self.rows,requested_target_rad=0.,selected_target_rad=actual)
                            next_pending[side] = None
                            next_latches[side] = dict(direction='close',target_rad=actual,completed_chunks=0)
                            deferred = None
                    next_close_defer[side] = deferred
                hand['close_defer'] = check

            if self.place_advance_rows:
                advance = dict(enabled=True, advance_rows=self.place_advance_rows,
                    forecast_open_row=None if not len(opens) else int(opens[0])+1,
                    original_deadline=self.release_deadlines[side], target_row=None,
                    modified=False, event='not_in_release_state')


                if not len(opens) or open_then_close:
                    next_release_deadlines[side] = None
                    advance['event'] = 'forecast_withdrawn' if not len(opens) else 'prefix_open_then_close'
                elif latch is not None and latch['direction'] == 'open':
                    next_release_deadlines[side] = None
                    advance['event'] = 'existing_open_latch'
                elif verified and (holding or angle < .27):
                    first_open = int(opens[0])+1
                    deadline = self.release_deadlines[side]
                    if deadline is None:
                        deadline = self.executed_row_clock + first_open
                    next_release_deadlines[side] = deadline


                    expected = min(deadline-self.executed_row_clock, first_open)
                    target_row = max(1, expected-self.place_advance_rows)
                    advance.update(original_deadline=deadline, expected_open_row=expected,
                        target_row=target_row, event='reserved_open_deadline')
                    if target_row <= self.rows and target_row < first_open:
                        requested = float(values[opens[0]])
                        actual = apply(slice(target_row-1,self.rows), requested)
                        advance.update(modified=True,event='apply_advanced_open',requested_target_rad=requested,
                            selected_target_rad=actual, previous_transition_event=hand['event'])
                        hand.update(event='apply_advanced_open',direction='open')
                        next_pending[side] = None
                        next_latches[side] = dict(direction='open',target_rad=actual,completed_chunks=0)
                        next_release_deadlines[side] = None
                    elif len(prefix_open):
                        advance['event'] = 'natural_open_in_prefix'
                        next_release_deadlines[side] = None
                else:
                    next_release_deadlines[side] = None
                hand['place_advance'] = advance

            final = float(selected[self.rows-1, col])


            next_verified[side] = bool(verified or (final < .07 and seen_open))
            if final >= .27:
                next_verified[side] = False
            next_seen[side] = seen_open
            next_phases[side] = ('closed' if next_verified[side] and final < .07
                else 'release_started' if next_verified[side]
                else 'wide_open' if final >= .27 else 'preclose_open' if final >= .07
                else 'initial_closed_unverified')
            hand.update(selected_row_rad=final, phase_after_selected_chunk=next_phases[side],
                verified_close_after_selected_chunk=next_verified[side],
                override_rows_1based=sorted(set(hand['override_rows_1based'])))
            hand[f'selected_row{self.rows}_rad'] = final
            report['hands'][side] = hand
        report.update(proposed_close_defer=next_close_defer,
            proposed_release_deadlines=next_release_deadlines,
            proposed_pending=next_pending, proposed_phases=next_phases,
            proposed_verified_close=next_verified, proposed_seen_open=next_seen,
            proposed_latches=next_latches)
        return selected, report

    def annotate_plan(self, report, adopted_rows):
        report['adopted_rows'] = adopted_rows
        for hand in report['hands'].values():
            hand['scheduled_in_adopted_plan'] = bool(hand['override'] and adopted_rows == self.rows)

    def commit(self, report, *, executed_rows):
        if (report.get('_owner') != self._owner or report.get('_generation') != self._generation
                or report.get('state_committed') or report['cycle'] <= self.last_committed_cycle):
            raise ValueError('Pen transition proposal already committed/stale/foreign')
        if type(executed_rows) is not int or not 0 <= executed_rows <= self.rows:
            raise ValueError('Invalid verified executed rows')
        full = executed_rows == self.rows
        if full:
            self.close_defer_state = deepcopy(report['proposed_close_defer'])
            self.release_deadlines = deepcopy(report['proposed_release_deadlines'])
            self.pending = deepcopy(report['proposed_pending'])
            self.phases = deepcopy(report['proposed_phases'])
            self.verified_close = deepcopy(report['proposed_verified_close'])
            self.seen_open = deepcopy(report['proposed_seen_open'])
            self.latches = deepcopy(report['proposed_latches'])
        else:
            self.close_defer_state = {s: None for s in SIDES}
            self.release_deadlines = {s: None for s in SIDES}
            self.pending = {s: None for s in SIDES}
            self.latches = {s: None for s in SIDES}
        self.executed_row_clock += executed_rows
        self.last_committed_cycle = report['cycle']
        self._generation += 1
        report.update(state_committed=True, executed_rows=executed_rows,
            commit_result='full_chunk' if full else 'shortened_drop_pending',
            close_defer_after=deepcopy(self.close_defer_state),
            release_deadlines_after=deepcopy(self.release_deadlines),
            executed_row_clock_after=self.executed_row_clock,
            pending_after=deepcopy(self.pending), phases_after=deepcopy(self.phases),
            latches_after=deepcopy(self.latches), verified_close_after=deepcopy(self.verified_close))


ARM_POLICY = dict(
    kind="u1_pen_grasp_axis_pose_retention_v9",
    correction_scale=1.25, approach_axis_hand=[1., 0., 0.],
    grasp_width_min_rad=3 * 2 * np.pi / 4096,
    closed_command_threshold_rad=.07, raw_forecast_horizon_rows=32,
    progress_epsilon_m=1e-6, reference_max_translation_m=.005,
    reference_max_rotation_rad=.05,
    coordinate_reference="opposite hand, assumed stationary; contract inter-hand pose only",
    direction="negative active-hand X at width-gated grasp, frozen in opposite-hand frame",
    translation="configured scale times the old/latest target deficit along estimated up only",
    history_source="raw predictions, never corrected predictions",
    orientation="retain missing rotation progress around a frozen forecast axis; same correction scale as height",
    rotation_axis_min_angle_rad=np.pi/180, rotation_progress_epsilon_rad=1e-5,
    rotation_blend_in="smoothstep over executed prefix; full deficit at prefix endpoint",
    grippers="unchanged; any independent client lookahead stays separate",
    fixed_postgrasp_row_cutoff=False, world_inputs=False, robot_fk_inputs=False,
    calibrated_transform_inputs=False, current_or_contact_inputs=False,
    lift_success_classifier=False,
)


def regime(x):
    return "closed" if x < .07 else "open" if x >= .27 else "narrow"


def pair(x):
    a = np.asarray([x[s] for s in SIDES] if isinstance(x, dict) else x, dtype=float)
    if a.shape != (2,) or not np.isfinite(a).all():
        raise ValueError("Two finite finger values required")
    return a


def matrix(pose):
    p = np.asarray(pose, dtype=float)
    if p.shape != (7,) or not np.isfinite(p).all():
        raise ValueError("Finite xyz/xyzw pose required")
    q = p[3:]; norm = np.linalg.norm(q)
    if abs(norm - 1.) > .02:
        raise ValueError("Unit xyzw quaternion required")
    x, y, z, w = q / norm
    t = np.eye(4)
    t[:3, :3] = [[1-2*(y*y+z*z), 2*(x*y-z*w), 2*(x*z+y*w)],
                [2*(x*y+z*w), 1-2*(x*x+z*z), 2*(y*z-x*w)],
                [2*(x*z-y*w), 2*(y*z+x*w), 1-2*(x*x+y*y)]]
    t[:3, 3] = p[:3]
    return t


def inverse(t):
    result = np.eye(4)
    result[:3, :3] = t[:3, :3].T
    result[:3, 3] = -result[:3, :3] @ t[:3, 3]
    return result


def angle(r):
    return float(np.arccos(np.clip((np.trace(r)-1.)/2., -1., 1.)))


def quaternion(r):

    trace=float(np.trace(r))
    if trace > 0:
        s=2*np.sqrt(max(0.,trace+1.))
        q=np.array([(r[2,1]-r[1,2])/s,(r[0,2]-r[2,0])/s,(r[1,0]-r[0,1])/s,s/4])
    else:
        i=int(np.argmax(np.diag(r)));j=(i+1)%3;k=(i+2)%3
        s=2*np.sqrt(max(0.,1.+r[i,i]-r[j,j]-r[k,k]))
        q=np.zeros(4);q[i]=s/4;q[j]=(r[i,j]+r[j,i])/s;q[k]=(r[i,k]+r[k,i])/s
        q[3]=(r[k,j]-r[j,k])/s
    q/=np.linalg.norm(q)
    return -q if q[3]<0 else q


def axis_rotation(axis, radians):
    x,y,z=axis
    skew=np.array([[0.,-z,y],[z,0.,-x],[-y,x,0.]])
    return np.eye(3)+np.sin(radians)*skew+(1-np.cos(radians))*(skew@skew)


def twist_angle(r, axis, reference=0.):

    q=quaternion(r);v=float(q[:3]@axis);w=float(q[3])
    if np.hypot(v,w)<1e-8:
        raise ValueError("Rotation twist is singular")
    value=2*np.arctan2(v,w)
    return float(reference+(value-reference+np.pi)%(2*np.pi)-np.pi)


def twist_path(rotations, base, axis, reference):
    result=[]
    for r in rotations:
        reference=twist_angle(r@base.T,axis,reference);result.append(reference)
    return np.asarray(result)


def forecast(raw, starts):
    poses = np.empty((len(raw), 2, 4, 4))
    for h in range(2):
        p = starts[h].copy()
        for i, row in enumerate(raw):
            p = p @ matrix(row[h*7:h*7+7])
            poses[i, h] = p
    return poses


class GraspAxisRetention:


    def __init__(self, execution_rows=None, *, correction_scale=None):
        if execution_rows is not None and (type(execution_rows) is not int
                                          or execution_rows not in (8, 16)):
            raise ValueError("Execution rows must be 8, 16, or None")
        scale = ARM_POLICY["correction_scale"] if correction_scale is None else correction_scale
        if isinstance(scale, (bool, np.bool_)) or not isinstance(scale, (int, float, np.integer, np.floating)) or not np.isfinite(scale) or not 0 <= scale <= 2:
            raise ValueError("correction_scale must be a finite number in [0,2]")
        self.scale = float(scale)
        self.n = execution_rows
        self.policy = dict(ARM_POLICY, execution_rows=execution_rows, correction_scale=self.scale,
                           max_old_age_chunks=None if execution_rows is None else 32//execution_rows-1)
        self.reset()

    def reset(self):
        self.history = []
        self.clock = self.generation = 0
        self.held_last = [False, False]
        self.held_rows = [0, 0]
        self.grasp_directions = [None, None]
        self.retention_done = [False, False]
        self.rotation_bases = [None, None]
        self.rotation_axes = [None, None]
        self.rotation_retention_done = [False, False]
        self.rotation_progress = [0., 0.]
        self.observed_execution_rows = None
        self.pending = None
        self.last_prompt = None
        self.last_report = None

    def propose(self, raw, obs, targets=None, previous_applied=None):
        raw = np.asarray(raw)
        if raw.shape != (32, 16) or raw.dtype.kind != "f" or not np.isfinite(raw).all():
            raise ValueError("Finite floating (32,16) raw predictions required")
        relative = matrix(obs[POSE_KEY])
        current = [inverse(relative), relative]
        measured = pair(obs["observation.joint_states"])
        previous = measured if previous_applied is None else pair(previous_applied)
        n = self.n or 32
        selected = np.asarray(raw[:n, 14:16] if targets is None else targets, dtype=float)
        if selected.shape != (n, 2) or not np.isfinite(selected).all():
            raise ValueError("Selected gripper targets must match the returned prefix")
        poses = forecast(raw, current)


        relative_motion = forecast(raw, [np.eye(4), np.eye(4)])
        ref_rows = self.n or 16
        stable = []
        for h in range(2):
            other = 1-h
            translation = np.max(np.linalg.norm(relative_motion[:ref_rows, other, :3, 3], axis=1))
            rotation = max(angle(t[:3, :3]) for t in relative_motion[:ref_rows, other])
            stable.append(bool(translation <= ARM_POLICY["reference_max_translation_m"]
                               and rotation <= ARM_POLICY["reference_max_rotation_rad"]))
        held, directions, done, resets, counts, reasons = [], [], [], [], [], []
        for h in range(2):


            gate = selected[:, h] if self.n else selected[:1, h]
            sequence = [regime(previous[h])] + [regime(v) for v in gate]
            transition = any(a != b for a, b in zip(sequence, sequence[1:]))
            closing = all(x == "closed" for x in sequence)
            width = ARM_POLICY["grasp_width_min_rad"]-1e-12 <= measured[h] < .07
            active = bool(closing and width and stable[h])
            held.append(active); resets.append(transition)
            counts.append(self.held_rows[h] if active and self.held_last[h] else 0)
            direction = self.grasp_directions[h] if active and self.held_last[h] else None
            if active and direction is None:
                direction = -current[h][:3, :3][:, 0].copy()
            directions.append(direction)
            done.append(self.retention_done[h] if active and self.held_last[h] else False)
            why = []
            if transition: why.append("selected_grip_transition")
            if not closing or not width: why.append("width_gate_inactive")
            if not stable[h]: why.append("opposite_hand_prediction_moves")
            if done[h]: why.append("postgrasp_lift_phase_complete")
            reasons.append(why)
        output = raw.copy(); reports = {}; rotation_states = []
        for h, side in enumerate(SIDES):
            direction = directions[h]; start = current[h][:3, 3]
            fresh = poses[:n, h, :3, 3]
            additions = np.zeros(n); best = 0.; chosen = None
            pending_forward = False; old_peak = 0.; fresh_peak = 0.
            if held[h] and not done[h]:
                progress = (fresh-start) @ direction
                closed_end = next((i for i in range(32) if regime(raw[i, 14+h]) != "closed"), 32)
                if closed_end:
                    fresh_peak = max(0., float(np.max((poses[:closed_end, h, :3, 3]-start) @ direction)))
                for entry in self.history:
                    offset = self.clock-entry["birth"]
                    if resets[h] or not entry["valid"][h] or not 1 <= offset < 32:
                        continue
                    length = min(n, 32-offset)
                    if self.n and length != n: continue
                    mismatch = next((i for i in range(length)
                        if regime(entry["grips"][offset+i, h]) != regime(selected[i, h])
                        or regime(selected[i, h]) != "closed"), length)
                    if self.n and mismatch != length:
                        reasons[h].append("old_grip_regime_mismatch"); continue
                    length = mismatch
                    if not length: continue
                    old_progress = (entry["poses"][offset:offset+length, h, :3, 3]-start) @ direction
                    peak = float(np.max(old_progress)); old_peak = max(old_peak, peak)
                    if peak <= ARM_POLICY["progress_epsilon_m"]: continue
                    pending_forward = True
                    delta = np.zeros(n)
                    delta[:length] = np.maximum(np.maximum(old_progress, 0.)-progress[:length], 0.) * self.scale
                    amount = float(np.max(delta))
                    if amount > best:
                        best, additions, chosen = amount, delta, (entry, offset)
                end = min(n, closed_end)
                if (not pending_forward and fresh_peak <= ARM_POLICY["progress_epsilon_m"]
                        and end and progress[end-1] < -ARM_POLICY["progress_epsilon_m"]):
                    done[h] = True
                    reasons[h].append("no_unexecuted_upward_forecast; return_to_latest_model")
            modified = False
            rotation=self._rotation_proposal(h,current,poses,selected,held[h],resets[h],raw[:,14+h])
            rotation_states.append(rotation["state"])
            if chosen is not None or rotation["modified"]:
                prior_p = start; prior_r = current[h][:3, :3]
                for i in range(n):
                    retained = fresh[i] if chosen is None else fresh[i] + additions[i]*direction
                    output[i, h*7:h*7+3] = prior_r.T @ (retained-prior_p)
                    target_r=rotation["rotations"][i]
                    if rotation["modified"]:
                        q=quaternion(prior_r.T@target_r)
                        if q@raw[i,h*7+3:h*7+7]<0:q=-q
                        output[i,h*7+3:h*7+7]=q
                    prior_p, prior_r = retained, target_r
                modified = chosen is not None and best>0
            reports[side] = dict(
                blended=False, translation_retained=modified, grasp_by_width=held[h],
                measured_width_rad=float(measured[h]), previous_applied_command_rad=float(previous[h]),
                width_threshold_rad=ARM_POLICY["grasp_width_min_rad"], postgrasp_executed_rows=counts[h],
                reset_reasons=reasons[h], retention_phase_complete=done[h],
                reference_hand=SIDES[1-h], reference_assumed_stationary=stable[h],
                direction_frame="opposite hand at grasp, assumed stationary",
                correction_direction=None if direction is None else direction.tolist(),
                correction_direction_hand=None if direction is None else (current[h][:3, :3].T @ direction).tolist(),
                pending_old_forward_forecast=pending_forward,
                old_forward_peak_m=old_peak, fresh_forward_peak_m=fresh_peak,
                old_sources=[] if chosen is None else [dict(generation=chosen[0]["generation"],
                    age_chunks=self.generation-chosen[0]["generation"], consumed_rows=chosen[1])],
                added_forward_m=additions.tolist() if chosen is not None else [],
                max_target_shift_m=best, **rotation["report"])
        assert np.array_equal(output[:, 14:], raw[:, 14:])
        for h in range(2):
            if not reports[SIDES[h]]["rotation_retained"]:
                assert np.array_equal(output[:, h*7+3:h*7+7], raw[:, h*7+3:h*7+7])
        assert np.array_equal(output[n:], raw[n:])
        report = dict(enabled=True, policy=self.policy, generation=self.generation,
                      executed_row_clock=self.clock, modified=not np.array_equal(output, raw),
                      hands=reports, retained_prefix_rows=n, committed=False,
                      future_rows="unchanged raw predictions" if self.n else "finite old-tail overlap only")
        return dict(actions=output, report=report, generation=self.generation,
                    poses=poses, grips=raw[:, 14:].copy(), selected=selected.copy(),
                    reset=resets, held=held, held_rows=counts, directions=directions, phase_done=done,
                    rotation_states=rotation_states,
                    relative_path=self._relative_path(relative, output))

    def _rotation_proposal(self,h,current,poses,selected,held,reset,raw_grips):
        n=self.n or 32;fresh=poses[:n,h,:3,:3]
        carry=held and self.held_last[h]
        base=self.rotation_bases[h] if carry else current[h][:3,:3].copy() if held else None
        axis=self.rotation_axes[h] if carry else None
        complete=self.rotation_retention_done[h] if carry else False
        previous=self.rotation_progress[h] if carry else 0.
        state=dict(base=base,axis=axis,done=complete,progress=previous)
        report=dict(rotation_retained=False,rotation_phase_complete=complete,
            rotation_axis_frame="opposite hand, assumed stationary",rotation_axis=None,
            rotation_progress_rad=previous,added_rotation_rad=[],rotation_old_sources=[],
            max_target_rotation_change_rad=0.,rotation_reset_reasons=[])
        result=dict(rotations=fresh.copy(),modified=False,state=state,report=report)
        if not held:return result
        if complete:
            report["rotation_reset_reasons"].append("postgrasp_rotation_phase_complete")
            return result


        closed_end=next((i for i,v in enumerate(raw_grips) if regime(v)!="closed"),32)
        if axis is None:
            candidates=[quaternion(r@base.T) for r in poses[:closed_end,h,:3,:3]]
            if candidates:
                q=max(candidates,key=lambda q:2*np.arctan2(np.linalg.norm(q[:3]),q[3]))
                magnitude=2*np.arctan2(np.linalg.norm(q[:3]),q[3])
                if magnitude>=ARM_POLICY["rotation_axis_min_angle_rad"]:
                    axis=q[:3]/np.linalg.norm(q[:3]);state["axis"]=axis
            if axis is None:
                report["rotation_reset_reasons"].append("waiting_for_rotation_forecast_axis")
                return result
        report["rotation_axis"]=axis.tolist()
        try:
            present=twist_angle(current[h][:3,:3]@base.T,axis,previous)
            fresh_angles=twist_path(poses[:,h,:3,:3],base,axis,present)
            state["progress"]=present;report["rotation_progress_rad"]=present
            best=0.;chosen=None;additions=np.zeros(n);pending=False
            old_peak=0.;epsilon=ARM_POLICY["rotation_progress_epsilon_rad"]
            for entry in self.history:
                offset=self.clock-entry["birth"]
                if reset or not entry["valid"][h] or not 1<=offset<32:continue
                length=min(n,32-offset)
                if self.n and length!=n:continue
                mismatch=next((i for i in range(length) if regime(entry["grips"][offset+i,h])!=regime(selected[i,h])
                               or regime(selected[i,h])!="closed"),length)
                if self.n and mismatch!=length:continue
                length=mismatch
                if not length:continue
                old_angles=twist_path(entry["poses"][:,h,:3,:3],base,axis,present)
                progress=old_angles[offset:offset+length]-present
                peak=float(np.max(progress));old_peak=max(old_peak,peak)
                if peak<=epsilon:continue
                pending=True;delta=np.zeros(n)
                delta[:length]=np.maximum(np.maximum(progress,0.)-(fresh_angles[:length]-present),0.)*self.scale
                amount=float(np.max(delta))
                if amount>best:best,chosen,additions=amount,(entry,offset),delta
            fresh_peak=max(0.,float(np.max(fresh_angles[:closed_end]-present))) if closed_end else 0.
            end=min(n,closed_end)
            if not pending and fresh_peak<=epsilon and end and fresh_angles[end-1]-present<=epsilon:
                state["done"]=True;report["rotation_phase_complete"]=True
                report["rotation_reset_reasons"].append("no_unexecuted_rotation_forecast; return_to_latest_model")
            report.update(pending_old_rotation_forecast=pending,old_rotation_peak_rad=old_peak,
                          fresh_rotation_peak_rad=fresh_peak)
            if chosen is not None:
                if best>=np.pi:raise ValueError("Rotation correction reaches ambiguous half-turn")
                blend_rows=self.n or self.observed_execution_rows or 16
                t=np.minimum(np.arange(1,n+1,dtype=float)/blend_rows,1.)
                unramped=additions.copy()
                additions=additions*(3*t*t-2*t*t*t)
                for i in range(n):result["rotations"][i]=axis_rotation(axis,additions[i])@fresh[i]
                result["modified"]=True
                report.update(rotation_retained=True,added_rotation_rad=additions.tolist(),
                    unramped_added_rotation_rad=unramped.tolist(),rotation_blend_rows=blend_rows,
                    max_target_rotation_change_rad=float(np.max(additions)),rotation_old_sources=[dict(
                    generation=chosen[0]["generation"],age_chunks=self.generation-chosen[0]["generation"],consumed_rows=chosen[1])])
        except ValueError as exc:
            state["done"]=True;report["rotation_phase_complete"]=True
            report["rotation_reset_reasons"].append(str(exc))
        return result

    @staticmethod
    def _relative_path(relative, output):
        p = forecast(output, [np.eye(4), relative])
        return np.asarray([inverse(row[0]) @ row[1] for row in p])

    def commit(self, proposal, executed_rows, applied=None, *, verified=True):
        if proposal["generation"] != self.generation:
            raise ValueError("Stale grasp-axis proposal")
        limit = self.n or 32
        if type(executed_rows) is not int or not 0 <= executed_rows <= limit:
            raise ValueError("Invalid executed row count")
        if not executed_rows or (self.n and executed_rows != self.n):
            self.history.clear(); self.held_last = [False, False]; self.held_rows = [0, 0]
            self.grasp_directions = [None, None]; self.retention_done = [False, False]
            self.rotation_bases = [None, None]; self.rotation_axes = [None, None]
            self.rotation_retention_done = [False, False]; self.rotation_progress = [0., 0.]
            self.clock += executed_rows; self.generation += 1
            proposal["report"].update(committed=True, executed_rows=executed_rows, history_reset="shortened_prefix")
            return
        end = [regime(v) for v in proposal["selected"][executed_rows-1]]
        valid = [x == "closed" and proposal["report"]["hands"][side]["reference_assumed_stationary"]
                 for side,x in zip(SIDES,end)]
        if applied is not None:
            actual = pair(applied)
            valid = [v and regime(a) == e for v, a, e in zip(valid, actual, end)]
        reset = list(proposal["reset"])
        if self.n is None:
            for h in range(2):
                if any(regime(x) != "closed" for x in proposal["selected"][:executed_rows, h]):
                    reset[h] = True; proposal["held"][h] = False
        kept = []
        for entry in self.history:
            entry = dict(entry, valid=[a and b and not r and h for a,b,r,h in
                         zip(entry["valid"], valid, reset, proposal["held"])])
            if any(entry["valid"]) and self.clock+executed_rows-entry["birth"] < 32:
                kept.append(entry)
        kept.append(dict(birth=self.clock, generation=self.generation,
                         poses=proposal["poses"].copy(), grips=proposal["grips"].copy(), valid=valid))
        self.history = [e for e in kept if any(e["valid"])]
        self.held_last = [h and v for h,v in zip(proposal["held"], valid)]
        self.held_rows = [c+executed_rows if h else 0 for c,h in zip(proposal["held_rows"],self.held_last)]
        self.grasp_directions = [None if not h or d is None else d.copy()
                                for h,d in zip(self.held_last,proposal["directions"])]
        self.retention_done = [bool(h and d) for h,d in zip(self.held_last,proposal["phase_done"])]
        self.rotation_bases=[s["base"].copy() if held and s["base"] is not None else None
            for held,s in zip(self.held_last,proposal["rotation_states"])]
        self.rotation_axes=[s["axis"].copy() if held and s["axis"] is not None else None
            for held,s in zip(self.held_last,proposal["rotation_states"])]
        self.rotation_retention_done=[bool(held and s["done"]) for held,s in zip(self.held_last,proposal["rotation_states"])]
        self.rotation_progress=[s["progress"] if held else 0. for held,s in zip(self.held_last,proposal["rotation_states"])]
        self.clock += executed_rows; self.generation += 1
        proposal["report"].update(committed=True, executed_rows=executed_rows,
            endpoint_verified=verified, commit_boundary="verified_client_execution" if verified else "observed_relative_pose_alignment",
            postgrasp_executed_rows_after=dict(zip(SIDES,self.held_rows)),
            history_generations=[e["generation"] for e in self.history])


def _slow_half_speed(actions):


    out = np.repeat(actions, 2, axis=0)
    for start in (0, 7):
        rotation = Rotation.from_quat(actions[:, start+3:start+7])
        half = Rotation.from_rotvec(rotation.as_rotvec() * .5)
        translation = actions[:, start:start+3] * .5
        out[0::2, start:start+3] = translation
        out[1::2, start:start+3] = half.inv().apply(translation)
        out[0::2, start+3:start+7] = half.as_quat(canonical=True)
        out[1::2, start+3:start+7] = half.as_quat(canonical=True)
    return out.astype(np.float32)


def _pa1_pick_compression(actions, measured):
    measured = np.asarray(measured, dtype=float)
    if measured[1] <= .56 or np.max(np.abs(actions[0, 14:16]-measured)) > .05:
        return None
    if np.max(np.ptp(actions[:, 14:16], axis=0)) > .03 or np.min(actions[:, 15]) < .56:
        return None
    distances = np.stack([np.linalg.norm(actions[:, k:k+3], axis=1) for k in (0, 7)], axis=1)
    angles = np.stack([Rotation.from_quat(actions[:, k+3:k+7]).magnitude() for k in (0, 7)], axis=1)
    if np.max(distances[:4].sum(axis=0)) >= .002 or np.max(angles[:4].sum(axis=0)) >= np.deg2rad(1.):
        return None
    groups, counts = [], []
    index = 0
    while index < len(actions) and len(groups) < 8:
        end = index+1
        length = distances[index].copy()
        angle = angles[index].copy()
        while end < min(index+8, len(actions)):
            next_length = length+distances[end]
            next_angle = angle+angles[end]
            if np.max(next_length) > .002 or np.max(next_angle) > np.deg2rad(.5):
                break
            length, angle = next_length, next_angle
            end += 1
        combined = actions[end-1].astype(np.float64).copy()
        for start in (0, 7):
            position = np.zeros(3)
            rotation = Rotation.identity()
            for row in actions[index:end]:
                position += rotation.apply(row[start:start+3])
                rotation = rotation*Rotation.from_quat(row[start+3:start+7])
            combined[start:start+3] = position
            combined[start+3:start+7] = rotation.as_quat(canonical=True)
        groups.append(combined)
        counts.append(end-index)
        index = end
    consumed = sum(counts[:4])
    if consumed <= 4:
        return None
    while len(groups) < 8:
        hold = groups[-1].copy()
        hold[:14] = np.tile([0., 0., 0., 0., 0., 0., 1.], 2)
        groups.append(hold)
    return _slow_half_speed(np.asarray(groups, dtype=np.float32)), consumed


def _cup_tip_pivot(actions, column, release_row, hold_angle):
    out = actions.copy()
    opening = 2*(release_row-1)
    chunk = opening//8*8
    stop = max(opening, chunk+4)
    start = stop-4
    delay = stop-opening
    if delay:
        out[stop:, column] = actions[opening:len(actions)-delay, column]
    out[start:stop, column] = hold_angle
    rotation = Rotation.from_rotvec([0., -np.deg2rad(10.)/4., 0.])
    pivot = np.array([.09, 0., 0.])
    translation = pivot-rotation.apply(pivot)
    offset = 0 if column == 14 else 7
    out[start:stop, offset:offset+3] = translation
    out[start:stop, offset+3:offset+7] = rotation.as_quat(canonical=True)
    return out


def _pen_preclose_approach(actions, measured):
    closes = np.flatnonzero(actions[:16, 15] < .07)
    if not len(closes) or measured[1] < .07:
        return None
    first = int(closes[0])
    stop = max(first, 4)
    start = stop-4
    delay = stop-first
    out = actions.copy()
    if delay:
        out[stop:, 15] = actions[first:len(actions)-delay, 15]
        out[start:stop, 15] = measured[1]
    rotation = Rotation.identity()
    for i in range(start, stop):
        out[i, 7:10] += rotation.inv().apply([.005/4., 0., 0.])
        rotation = rotation*Rotation.from_quat(out[i, 10:14])
    return out


def _pen_extra_lift(actions, direction):
    out = actions.copy()
    rotation = Rotation.identity()
    for i in range(16):
        out[i, 7:10] += rotation.inv().apply(np.asarray(direction) * .0025)
        rotation = rotation * Rotation.from_quat(out[i, 10:14])
    return out


def _pa2_initial_lift_direction(actions):
    position = np.zeros(3)
    rotation = Rotation.identity()
    for row in actions[:8]:
        if row[14] >= .56:
            break
        position += rotation.apply(row[:3])
        rotation = rotation * Rotation.from_quat(row[3:7])
    length = np.linalg.norm(position)
    return position / length if length > 1e-6 else None


def _pa2_scale_lift(actions, direction, release_row):
    out = actions.copy()
    rotation = Rotation.identity()
    limit = min(16, 2 * (release_row - 1)) if release_row is not None else 16
    finished = False
    for i in range(limit):
        if out[i, 14] >= .56:
            finished = True
            break
        projected = float(rotation.apply(out[i, :3]) @ direction)
        if projected < -1e-5:
            finished = True
            break
        if projected > 0:
            out[i, :3] += rotation.inv().apply(.5 * projected * direction)
        rotation = rotation * Rotation.from_quat(out[i, 3:7])
    return out, finished or release_row is not None


class Policy:
    def __init__(self, checkpoint_dir: str):
        from openpi.models import pi0_config
        from openpi.policies import policy_config
        from openpi.training.config import AssetsConfig, DataConfig, TrainConfig

        checkpoint = pathlib.Path(checkpoint_dir)
        if not (checkpoint / "params").is_dir():
            checkpoint = checkpoint / "step_042000"
        config = TrainConfig(
            name="u1_42000_submission",
            model=pi0_config.Pi0Config(pi05=True, action_dim=32, action_horizon=32),
            data=_build_transforms()(
                assets=AssetsConfig(asset_id=_asset_id(checkpoint)),
                base_config=DataConfig(prompt_from_task=True)),
        )
        self._model = policy_config.create_trained_policy(config, str(checkpoint))
        self.reset()

        rng = self._model._rng
        image = np.random.default_rng(0).integers(0, 256, (480, 640, 3), dtype=np.uint8)
        obs = {_K[0]: image, _K[1]: image.copy(),
               POSE_KEY: np.array([.23, -.05, .27, -.58, -.09, .72, .34], np.float32),
               _K[3]: np.array([.24, .78], np.float32)}
        started = time.monotonic()
        try:
            for prompt in (PEN_PROMPTS[0], CUP_PA1, CUP_PA2):
                self._sample(dict(obs, prompt=prompt), prompt)
        finally:
            self._model._rng = rng
            self.reset()
        print(f"[policy] pre-warmed in {time.monotonic()-started:.1f}s", flush=True)

    def reset(self):

        self._prompt = None
        self._second_half = None
        self._pending = None
        self._cycle = 0
        self._previous_grip = None
        self._cup = None
        self._cup_seen_open = False
        self._cup_grasped = False
        self._cup_tilted = False
        self._pa2_lift_axis = None
        self._pa2_lift_done = False
        self._pen_approach_done = False
        self._pen_lift_added = False
        self._pen = PenGraspSchedule(8, close_defer_only=True)
        self._arm = GraspAxisRetention(execution_rows=8, correction_scale=1.0)

    def _sample(self, obs, prompt):
        clean = {key: obs[key] for key in _K if key in obs}
        mapped = prompt
        if prompt.startswith("Pick 2nd ") and " part and place into cell" in prompt:
            mapped = "Pick 1st " + prompt[len("Pick 2nd "):]
        clean["prompt"] = mapped
        sigma = CUP_NOISE.get(prompt, 1.0)
        kwargs = {}
        if sigma != 1.0:
            import jax

            _, key = jax.random.split(self._model._rng)
            noise = jax.random.normal(key, (1, 32, 32))
            kwargs["noise"] = noise * sigma
        actions = np.asarray(self._model.infer(clean, **kwargs)["actions"], dtype=np.float32)
        if actions.shape != (32, 16) or not np.isfinite(actions).all():
            raise ValueError("Expected finite float32 model actions [32,16]")
        return actions.copy()

    def _finish_previous_pair(self):


        if self._pending is None:
            return
        kind, proposals, applied = self._pending
        if kind == "cup":
            self._cup.commit(proposals, proposals['adopt_rows'])
        else:
            transition, arm = proposals
            self._pen.commit(transition, executed_rows=8)
            self._arm.commit(arm, 8, applied=applied, verified=False)
            self._previous_grip = applied.copy()
        self._pending = None
        self._cycle += 1

    def infer(self, obs: dict) -> np.ndarray:
        prompt = obs.get("prompt", "").strip().rstrip(".")
        if prompt != self._prompt:
            self.reset()
            self._prompt = prompt
            if prompt in CUP_HANDS:
                self._cup = CupContinuousAdvance(CUP_ADVANCE[prompt])
        if self._second_half is not None:
            result = self._second_half
            self._second_half = None
            return result.copy()

        self._finish_previous_pair()
        raw = self._sample(obs, prompt)
        tip_release = None
        lift_direction = None
        if prompt in CUP_HANDS:
            finger = float(obs[_K[3]][CUP_HANDS[prompt][1]-14])
            if finger > .53:
                self._cup_seen_open = True
            elif self._cup_seen_open and finger >= 3*2*np.pi/4096:
                self._cup_grasped = True
            if prompt == CUP_PA1 and not self._cup_grasped:
                compressed = _pa1_pick_compression(raw, obs[_K[3]])
                if compressed is not None:
                    output, consumed = compressed
                    proposal = self._cup.propose(raw, dict(obs, prompt=prompt), consumed)
                    self._pending = ('cup', proposal, None)
                    return output.copy()
            proposal = self._cup.propose(raw, dict(obs, prompt=prompt), 8)
            controlled = proposal["actions"]
            self._pending = ("cup", proposal, None)
            tip_release = proposal["report"]["forced_start_row"] or proposal["report"]["natural_start_row"]
        elif prompt in PEN_PROMPTS:
            measured = np.asarray(obs[_K[3]], dtype=float)
            if measured[1] >= .27 and np.any(raw[:8, 15] >= .27):
                self._pen_approach_done = False
                self._pen_lift_added = False


            limits = {side: dict(min_angle_rad=-1e10, max_angle_rad=1e10) for side in SIDES}
            targets, transition = self._pen.propose(
                self._cycle, raw, raw[:8, 14:16], measured, {}, limits,
                allow_clipping=False)
            arm = self._arm.propose(raw, obs, targets,
                                    previous_applied=self._previous_grip)
            hand = arm["report"]["hands"]["right"]
            if not self._pen_lift_added and hand["translation_retained"]:
                lift_direction = hand["correction_direction_hand"]
            controlled = arm["actions"].copy()
            controlled[:8, 14:16] = targets
            self._pending = ("pen", (transition, arm), targets[-1].copy())
        else:
            controlled = raw

        slow = _slow_half_speed(controlled)
        if prompt == CUP_PA2 and self._cup_grasped and not self._pa2_lift_done:
            relative_rotation = Rotation.from_quat(np.asarray(obs[POSE_KEY])[3:])
            if self._pa2_lift_axis is None:
                direction = _pa2_initial_lift_direction(slow)
                if direction is not None:
                    self._pa2_lift_axis = relative_rotation.inv().apply(direction)
            if self._pa2_lift_axis is not None:
                direction = relative_rotation.apply(self._pa2_lift_axis)
                slow, self._pa2_lift_done = _pa2_scale_lift(slow, direction, tip_release)
        if lift_direction is not None:
            slow = _pen_extra_lift(slow, lift_direction)
            self._pen_lift_added = True
        if prompt in PEN_PROMPTS and not self._pen_approach_done:
            approached = _pen_preclose_approach(slow, np.asarray(obs[_K[3]], dtype=float))
            if approached is not None:
                slow = approached
                self._pen_approach_done = True
        if tip_release is not None and self._cup_grasped and not self._cup_tilted:
            column = CUP_HANDS[prompt][1]
            slow = _cup_tip_pivot(slow, column, tip_release, float(obs[_K[3]][column-14]))
            self._cup_tilted = True


        self._second_half = slow[8:24].copy()
        return slow[:16].copy()