"""Versioned, stateful action adapter mirrored by mobile/SafeActionController.ts.

Position targets are integrated at bounded speed, not remapped over full joint travel.
Both applied actions and target state are observable; no hidden policy-side filter.
"""

import numpy as np

from .kernel import TASK as T
from .kernel import clip, navigation_error


class SafeActionController:
    def __init__(self, config):
        self.config = config
        self.applied = np.zeros(T["actionSize"], dtype=np.float32)
        self.targets = np.zeros(T["actionSize"], dtype=np.float32)
        self.control = np.zeros(
            len(config["baseJoints"]) + len(config["armJoints"]) + len(config["gripperActuators"])
        )

    def reset(self, state, previous_control=None):
        previous_control = previous_control.copy() if previous_control is not None else None
        self.applied.fill(0)
        self.targets.fill(0)
        self.control.fill(0)
        n = len(self.config["baseJoints"])
        for i, j in enumerate(self.config["armJoints"]):
            if j["mode"] == "position":
                q = float(state[13 + i])
                if previous_control is not None and np.isfinite(previous_control[n + i]):
                    q = clip(
                        previous_control[n + i],
                        q - T["armTrackingError"],
                        q + T["armTrackingError"],
                    )
                q = clip(q, j["min"], j["max"])
                self.control[n + i] = q
                self.targets[3 + i] = 2 * (q - j["min"]) / (j["max"] - j["min"]) - 1
        opening = float(state[29])
        for i, g in enumerate(self.config["gripperActuators"]):
            index = n + len(self.config["armJoints"]) + i
            if (
                previous_control is not None
                and g.get("joint", self.config["gripperJoint"]) == self.config["gripperJoint"]
                and np.isfinite(previous_control[index])
            ):
                opening = clip(
                    (previous_control[index] - g["closed"]) / (g["open"] - g["closed"]), 0, 1
                )
                break
        self.targets[11] = 2 * opening - 1
        self._gripper(opening)

    def _gripper(self, opening):
        n = len(self.config["baseJoints"]) + len(self.config["armJoints"])
        for i, g in enumerate(self.config["gripperActuators"]):
            self.control[n + i] = g["closed"] + opening * (g["open"] - g["closed"])

    def apply(self, action, state, stage, lifted):
        action = np.asarray(action, dtype=np.float32)
        if action.shape != (T["actionSize"],) or not np.isfinite(action).all():
            raise ValueError("action must be finite [12]")
        distance, yaw = navigation_error(state)
        near = distance < T["navigationTolerance"] and abs(yaw) < T["navigationYawTolerance"]
        manipulate = stage != "navigate" and (near or lifted)
        c, dt = self.config, T["controlDt"]
        # Actual body-frame velocity commands have bounded acceleration.
        for i in range(3):
            limit = min(c["baseLimits"][i], T["baseSpeedLimits"][i])
            desired = 0 if near and not lifted else clip(float(action[i]))
            delta = T["baseAccelerationLimits"][i] * dt / limit
            self.applied[i] = clip(desired, self.applied[i] - delta, self.applied[i] + delta)
        largest = c["wheelLimit"]
        for i, row in enumerate(c["baseMix"]):
            self.control[i] = sum(
                row[j] * self.applied[j] * min(c["baseLimits"][j], T["baseSpeedLimits"][j])
                for j in range(3)
            )
            largest = max(largest, abs(self.control[i]))
        n = len(c["baseJoints"])
        self.control[:n] *= c["wheelLimit"] / largest
        for i, j in enumerate(c["armJoints"]):
            k = 3 + i
            speed = min(j["velocityLimit"], T["armSpeedLimit"])
            delta = T["armAccelerationLimit"] * dt / speed
            a = clip(float(action[k])) if manipulate else 0.0
            a = clip(a, self.applied[k] - delta, self.applied[k] + delta)
            if j["mode"] == "position":
                old = self.control[n + i]
                desired = clip(
                    old + a * speed * dt,
                    state[13 + i] - T["armTrackingError"],
                    state[13 + i] + T["armTrackingError"],
                )
                target = clip(clip(desired, old - speed * dt, old + speed * dt), j["min"], j["max"])
                # During navigation hold the reset/current target, not the gravity-sagged qpos.
                if not manipulate:
                    target = old
                self.control[n + i] = target
                self.applied[k] = (target - old) / (speed * dt)
                self.targets[k] = 2 * (target - j["min"]) / (j["max"] - j["min"]) - 1
            else:
                self.applied[k] = a if manipulate else 0
                self.control[n + i] = self.applied[k] * speed
        opening = (float(self.targets[11]) + 1) / 2
        target = clip(
            opening
            + (clip(float(action[11])) if manipulate and stage == "pick-place" else 0)
            * T["gripperOpeningRate"]
            * dt,
            0,
            1,
        )
        self.applied[11] = (target - opening) / (T["gripperOpeningRate"] * dt)
        self.targets[11] = 2 * target - 1
        self._gripper(target)
        return self.control