"""Math mirror of web_platform/src/mobile/TaskKernel.ts (SI, world, wxyz)."""

import json
import math
from pathlib import Path

import numpy as np

CONTRACTS = Path(__file__).resolve().parents[2] / "contracts"
TASK = json.loads((CONTRACTS / "mobile-manipulator-v2.json").read_text())
ROBOTS = json.loads((CONTRACTS / "mobile-robots-v1.json").read_text())


def clip(value, lo=-1.0, hi=1.0):
    return max(lo, min(hi, value))


def rotate(q, v):
    w, x, y, z = q
    vx, vy, vz = v
    tx, ty, tz = 2 * (y * vz - z * vy), 2 * (z * vx - x * vz), 2 * (x * vy - y * vx)
    return [
        vx + w * tx + y * tz - z * ty,
        vy + w * ty + z * tx - x * tz,
        vz + w * tz + x * ty - y * tx,
    ]


def canonical(q):
    n = math.hypot(*q)
    return np.asarray(q) * ((-1 if q[0] < 0 else 1) / n) if n > 1e-12 else np.array([1, 0, 0, 0])


def relative(parent, child):
    w, x, y, z = parent[3:] * np.array([1, -1, -1, -1])
    a, b, c, d = child[3:]
    q = [
        w * a - x * b - y * c - z * d,
        w * b + x * a + y * d - z * c,
        w * c - x * d + y * a + z * b,
        w * d + x * c - y * b + z * a,
    ]
    pos = np.clip(
        np.array(rotate([w, x, y, z], child[:3] - parent[:3])) / TASK["positionScale"], -1, 1
    )
    return np.concatenate((pos, canonical(q)))


def validate_config(c):
    primary = [
        g for g in c["gripperActuators"] if g.get("joint", c["gripperJoint"]) == c["gripperJoint"]
    ]
    if not primary or any(
        g["closed"] != c["gripperClosed"] or g["open"] != c["gripperOpen"] for g in primary
    ):
        raise ValueError("gripper observation/actuator stroke mismatch")
    joints = [
        *c["baseJoints"],
        *(j["name"] for j in c["armJoints"]),
        *set(
            [c["gripperJoint"], *(g.get("joint", c["gripperJoint"]) for g in c["gripperActuators"])]
        ),
    ]
    actuators = [
        *c["baseActuators"],
        *c["armActuators"],
        *(g["name"] for g in c["gripperActuators"]),
    ]
    if (
        not c["id"]
        or c["recipe"] not in ("lekiwi-v1", "lekiwi-bundle", "mjcf")
        or not 0 < len(c["armJoints"]) <= TASK["maxArmJoints"]
        or len(c["armJoints"]) != len(c["armActuators"])
        or not len(c["baseJoints"]) == len(c["baseActuators"]) == len(c["baseMix"])
        or not c["baseJoints"]
        or not c["gripperActuators"]
        or len(set(joints)) != len(joints)
        or len(set(actuators)) != len(actuators)
    ):
        raise ValueError("invalid RobotConfig topology")
    if (
        np.asarray(c["baseMix"]).shape != (len(c["baseJoints"]), 3)
        or not np.isfinite(c["baseMix"]).all()
    ):
        raise ValueError("invalid baseMix")
    if (
        len(c["baseLimits"]) != 3
        or not np.isfinite(c["baseLimits"]).all()
        or min(c["baseLimits"]) <= 0
        or not math.isfinite(c["wheelLimit"])
        or c["wheelLimit"] <= 0
        or len(c["eefOffset"]) != 3
        or not np.isfinite(c["eefOffset"]).all()
    ):
        raise ValueError("invalid scales")
    for j in c["armJoints"]:
        if (
            not np.isfinite([j["min"], j["max"], j["neutral"], j["velocityLimit"]]).all()
            or not j["min"] <= j["neutral"] <= j["max"]
            or j["min"] >= j["max"]
            or j["velocityLimit"] <= 0
            or j["mode"] not in ("position", "velocity")
        ):
            raise ValueError("invalid arm limits")
    for closed, opened in [
        (c["gripperClosed"], c["gripperOpen"]),
        *((g["closed"], g["open"]) for g in c["gripperActuators"]),
    ]:
        if not np.isfinite([closed, opened]).all() or closed == opened:
            raise ValueError("invalid gripper stroke")


def decode_action(config, action, output=None):
    """Legacy v1 fixture oracle only; v2 environments use SafeActionController."""
    action = np.asarray(action, dtype=np.float32)
    if action.shape != (TASK["actionSize"],) or not np.isfinite(action).all():
        raise ValueError("action must be finite [12]")
    nbase, narm = len(config["baseJoints"]), len(config["armJoints"])
    if output is None:
        output = np.zeros(nbase + narm + len(config["gripperActuators"]), dtype=np.float64)
    largest = config["wheelLimit"]
    for i, row in enumerate(config["baseMix"]):
        output[i] = sum(row[j] * clip(float(action[j])) * config["baseLimits"][j] for j in range(3))
        largest = max(largest, abs(output[i]))
    output[:nbase] *= config["wheelLimit"] / largest
    for i, j in enumerate(config["armJoints"]):
        a = clip(float(action[3 + i]))
        output[nbase + i] = (
            j["min"] + (a + 1) * 0.5 * (j["max"] - j["min"])
            if j["mode"] == "position"
            else a * j["velocityLimit"]
        )
    opening = (clip(float(action[11])) + 1) * 0.5
    for i, g in enumerate(config["gripperActuators"]):
        output[nbase + narm + i] = g["closed"] + opening * (g["open"] - g["closed"])
    return output


STAGES = ("navigate", "reach", "pick-place")


def navigation_error(s):
    distance = math.hypot(
        s[37] - TASK["navigationOffset"][0] - s[0], s[38] - TASK["navigationOffset"][1] - s[1]
    )
    w, x, y, z = s[3:7]
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z))
    return distance, yaw


class TaskKernel:
    def __init__(self, config, stage="pick-place"):
        validate_config(config)
        if stage not in STAGES:
            raise ValueError("invalid training stage")
        self.stage = stage
        self.last_action = np.zeros(TASK["actionSize"], dtype=np.float32)
        self.targets = np.zeros(TASK["actionSize"], dtype=np.float32)
        self.config = config
        self.observation = np.zeros(TASK["observationSize"], dtype=np.float32)
        self.reset()

    def reset(self):
        self.steps = self.settle = 0
        self.has_lifted = self.terminated = self.truncated = False
        self.reward = self.action_rate = 0.0
        self.previous_distance = None
        self.last_action.fill(0)
        self.targets.fill(0)
        self.info = {
            "reward_components": dict.fromkeys(
                [
                    "reach",
                    "lift",
                    "transport",
                    "success",
                    "navigation",
                    "action_rate",
                    "joint_velocity",
                    "safety",
                ],
                0.0,
            ),
            "is_success": False,
            "stage": "navigate" if self.stage == "navigate" else "reach",
            "safety_stop": "",
            "navigation_distance": 0.0,
            "max_joint_velocity": 0.0,
        }

    def record_action(self, applied, targets):
        self.action_rate = sum(
            (float(applied[i]) - float(self.last_action[i])) ** 2 for i in range(TASK["actionSize"])
        )
        self.last_action[:] = applied
        self.targets[:] = targets

    def observe(self, state):
        s, o, t = state, self.observation, TASK
        o.fill(0)
        o[:3] = np.clip(s[:3] / t["positionScale"], -1, 1)
        o[3:7] = canonical(s[3:7])
        o[7:10] = np.clip(s[7:10] / t["linearVelocityScale"], -1, 1)
        o[10:13] = np.clip(s[10:13] / t["angularVelocityScale"], -1, 1)
        for i, j in enumerate(self.config["armJoints"]):
            o[13 + i] = clip(2 * (s[13 + i] - j["min"]) / (j["max"] - j["min"]) - 1)
            o[21 + i] = clip(s[21 + i] / t["jointVelocityScale"])
            o[29 + i] = 1
        o[37] = clip(2 * s[29] - 1)
        o[38:41] = np.clip(s[30:33] / t["positionScale"], -1, 1)
        o[41:45] = canonical(s[33:37])
        o[45:52] = relative(s[30:37], s[37:44])
        o[52:59] = relative(s[47:54], s[37:44])
        o[59:62] = np.clip(s[47:50] / t["positionScale"], -1, 1)
        o[62:66] = canonical(s[50:54])
        o[66], o[67] = self.has_lifted, self.settle / t["settleSteps"]
        o[68:80] = self.last_action
        o[80:92] = self.targets
        return o

    def evaluate(self, s, safety_stop="", peak_velocity=0.0):
        if self.terminated or self.truncated:
            raise RuntimeError("episode ended; reset required")
        t = TASK
        reach = math.hypot(*(s[37:40] - s[30:33]))
        goal = math.hypot(*(s[37:40] - s[47:50]))
        lift = clip((s[39] - t["objectStart"][2]) / t["liftHeight"], 0, 1)
        if self.stage == "pick-place" and lift >= 1 and reach < t["graspDistance"] and s[29] < 0.4:
            self.has_lifted = True
        distance, yaw = navigation_error(s)
        near = distance < t["navigationTolerance"] and abs(yaw) < t["navigationYawTolerance"]
        stopped = (
            math.hypot(*s[7:10]) < t["navigationSpeedTolerance"] and math.hypot(*s[10:13]) < 0.1
        )
        settled = (
            self.has_lifted
            and goal < t["goalTolerance"]
            and s[29] > t["releaseOpening"]
            and reach > t["graspDistance"]
            and math.hypot(*s[44:47]) < t["settleSpeed"]
        )
        if self.stage == "navigate":
            settled = near and stopped
        elif self.stage == "reach":
            settled = (
                near and stopped and reach < t["graspDistance"] and math.hypot(*s[21:29]) < 0.15
            )
        self.settle = min(t["settleSteps"], self.settle + 1) if settled else 0
        success = self.settle == t["settleSteps"]
        r = self.info["reward_components"]
        r["reach"] = t["controlDt"] * t["reachWeight"] * math.exp(-t["reachGain"] * reach)
        r["lift"] = t["controlDt"] * t["liftWeight"] * lift
        r["transport"] = (
            t["controlDt"] * t["transportWeight"] * math.exp(-t["transportGain"] * goal)
            if self.has_lifted
            else 0.0
        )
        if self.stage == "navigate":
            r["reach"] = r["lift"] = r["transport"] = 0.0
        elif self.stage == "reach":
            r["lift"] = r["transport"] = 0.0
        progress = 0 if self.previous_distance is None else self.previous_distance - distance
        self.previous_distance = distance
        r["navigation"] = (
            (
                t["navigationProgressWeight"] * progress
                - t["controlDt"] * (distance + 0.1 * abs(yaw))
            )
            if not self.has_lifted
            else 0.0
        )
        r["action_rate"] = -t["actionRateWeight"] * self.action_rate
        r["joint_velocity"] = (
            -t["controlDt"] * t["jointVelocityWeight"] * sum(float(v) ** 2 for v in s[21:29])
        )
        peak_velocity = max(peak_velocity, max(abs(s[21:29])))
        if not safety_stop and peak_velocity > t["jointSpeedStop"]:
            safety_stop = "joint_velocity"
        if not safety_stop and (
            1 - 2 * (s[4] ** 2 + s[5] ** 2) < 0.5 or max(abs(s[:2])) > t["positionScale"]
        ):
            safety_stop = "base_pose"
        if safety_stop:
            success = False
        r["success"] = t["successBonus"] if success else 0.0
        r["safety"] = -t["safetyPenalty"] if safety_stop else 0.0
        self.reward = sum(r.values())
        self.terminated = success or bool(safety_stop)
        self.info["safety_stop"] = safety_stop
        self.info["navigation_distance"] = distance
        self.info["max_joint_velocity"] = max(self.info["max_joint_velocity"], peak_velocity)
        self.steps += 1
        self.truncated = self.steps >= t["maxSteps"] and not success
        self.info["is_success"] = success
        self.info["stage"] = (
            "safety-stop"
            if safety_stop
            else "success"
            if success
            else "navigate"
            if self.stage == "navigate" or (not near and not self.has_lifted)
            else "transport"
            if self.has_lifted
            else "lift"
            if reach < t["graspDistance"]
            else "reach"
        )
        self.observe(s)
        return self.observation, self.reward, self.terminated, self.truncated, self.info