#!/usr/bin/env python3
"""
robot_fk.py
-----------
Minimal forward kinematics engine for the robot.json format.

Matches the Blender hierarchy used by render_robot.py exactly:
  world_T_link = world_T_parent
                 @ Translate(mountPosition) @ Rotate_xyz(mountRotation)
                 @ Translate(jointOrigin)   @ Rotate_xyz(joint.rotation)
                 @ T_motion

  T_motion =  Rotate(axis, value_deg)   for revolute joints
              Translate(axis * value_mm) for linear joints

Units throughout: millimetres, degrees.

Public API
----------
fk = RobotFK.from_file("robot.json")

transforms = fk.compute({"x": 180, "y": 86, "z": -120,
                          "a": -60, "b": 22, "c": 91, "e": 10})
# → dict  link_name -> 4×4 np.ndarray (world frame, mm)

p_world = fk.marker_world(transforms, "Arm1", [0, -160, 35])
# → np.ndarray shape (3,), in mm

all_m = fk.all_markers_world(transforms)
# → dict  marker_id -> {"world_mm", "link", "local_mm"}

# Cumulative x-offset for a link at all-zero state
# (useful for 4a: x_slider = world_x - local_x - link_x_at_zero)
x0 = fk.link_x_at_zero_state("Arm1")   # → float mm
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

STATE_KEYS = ("x", "y", "z", "a", "b", "c", "e")


# ──────────────────────────────────────────────────────────────
# Low-level matrix helpers
# ──────────────────────────────────────────────────────────────

def _rot_axis_angle(axis: Sequence[float], angle_deg: float) -> np.ndarray:
    """3×3 rotation matrix via Rodrigues (axis need not be normalised)."""
    ax = np.asarray(axis, dtype=float)
    n = float(np.linalg.norm(ax))
    if n < 1e-12:
        return np.eye(3)
    ax = ax / n
    c = math.cos(math.radians(angle_deg))
    s = math.sin(math.radians(angle_deg))
    t = 1.0 - c
    x, y, z = ax
    return np.array([
        [t*x*x + c,   t*x*y - s*z, t*x*z + s*y],
        [t*x*y + s*z, t*y*y + c,   t*y*z - s*x],
        [t*x*z - s*y, t*y*z + s*x, t*z*z + c  ],
    ])


def _rot_xyz_euler(rx: float, ry: float, rz: float) -> np.ndarray:
    """XYZ Euler angles (degrees) → 3×3 — matches Blender XYZ Euler mode."""
    return (_rot_axis_angle([0, 0, 1], rz)
            @ _rot_axis_angle([0, 1, 0], ry)
            @ _rot_axis_angle([1, 0, 0], rx))


def make_T(R: np.ndarray, t: Sequence[float]) -> np.ndarray:
    """4×4 homogeneous transform."""
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = np.asarray(t, dtype=float)
    return T


def transform_point(T: np.ndarray, p: Sequence[float]) -> np.ndarray:
    """Apply 4×4 transform to a 3-D point."""
    h = np.array([p[0], p[1], p[2], 1.0])
    return (T @ h)[:3]


# ──────────────────────────────────────────────────────────────
# FK engine
# ──────────────────────────────────────────────────────────────

class RobotFK:
    """Forward kinematics for the robot.json format."""

    def __init__(self, robot_data: Dict[str, Any]):
        self.robot = robot_data
        self.links: Dict[str, Any] = robot_data.get("links", {})
        self._topo: List[str] = self._topological_sort()

    # ── construction ─────────────────────────────────────────

    @classmethod
    def from_file(cls, path) -> "RobotFK":
        with open(path, "r", encoding="utf-8") as f:
            return cls(json.load(f))

    def _topological_sort(self) -> List[str]:
        parent_map = {n: d.get("parent") for n, d in self.links.items()}
        visited, order = set(), []

        def visit(name: str) -> None:
            if name in visited:
                return
            visited.add(name)
            p = parent_map.get(name)
            if p and p in self.links:
                visit(p)
            order.append(name)

        for name in self.links:
            visit(name)
        return order

    # ── core computation ──────────────────────────────────────

    def compute(self, joint_values: Dict[str, float]) -> Dict[str, np.ndarray]:
        """
        Compute link world transforms for the given joint state.

        Parameters
        ----------
        joint_values : dict  variable_name -> value
            Linear joints (x, e): mm
            Revolute joints (y, z, a, b, c): degrees

        Returns
        -------
        dict  link_name -> 4×4 np.ndarray (world frame, mm)
        """
        state = {k: 0.0 for k in STATE_KEYS}
        for k, v in joint_values.items():
            if k in state:
                state[k] = float(v)

        transforms: Dict[str, np.ndarray] = {}

        for link_name in self._topo:
            d = self.links[link_name]
            parent = d.get("parent")
            T_parent = transforms.get(parent, np.eye(4)) if parent else np.eye(4)

            # 1 · Mount (static in parent frame)
            mp  = d.get("mountPosition", [0, 0, 0])
            mr  = d.get("mountRotation", [0, 0, 0])
            T_m = make_T(_rot_xyz_euler(*mr), mp)

            # 2 · Joint origin (pivot point in mount frame)
            ji  = d.get("jointToParent", {}) or {}
            jp  = ji.get("origin", [0, 0, 0])
            jr  = ji.get("rotation", [0, 0, 0])
            T_j = make_T(_rot_xyz_euler(*jr), jp)

            # 3 · Motion
            jtype = str(ji.get("type", "fixed")).lower()
            var   = str(ji.get("variable", "")).lower()
            axis  = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
            val   = state.get(var, 0.0)

            if jtype == "revolute":
                T_mot = make_T(_rot_axis_angle(axis, val), [0, 0, 0])
            elif jtype == "linear":
                n = float(np.linalg.norm(axis))
                u = axis / n if n > 1e-12 else np.array([1.0, 0, 0])
                T_mot = make_T(np.eye(3), u * val)
            else:
                T_mot = np.eye(4)

            transforms[link_name] = T_parent @ T_m @ T_j @ T_mot

        return transforms

    # ── marker helpers ────────────────────────────────────────

    @staticmethod
    def marker_world(transforms: Dict[str, np.ndarray],
                     link_name: str,
                     local_pos: Sequence[float]) -> np.ndarray:
        """Transform a local marker position → world (mm)."""
        return transform_point(transforms.get(link_name, np.eye(4)), local_pos)

    # ── Marker-Eckpunkte (Mehrpunkt-Residuen, doc/Homing_5_Pose_MultiPoint_Weighted.md Schritt 3) ──

    @staticmethod
    def _shortest_arc_R(normal: Sequence[float]) -> np.ndarray:
        """
        Rotationsmatrix [0,0,1] → normal (kürzester Bogen).

        Repliziert THREE.Quaternion.setFromUnitVectors(vFrom=[0,0,1], vTo=n)
        exakt (inkl. antiparallelem Sonderfall), damit die hier erzeugten
        Ecken zur visuell verifizierten Orientierungs-Konvention in
        boardViewer.html passen (qNormalLoc dort).
        """
        n = np.asarray(normal, dtype=float)
        nn = float(np.linalg.norm(n))
        n = n / nn if nn > 1e-12 else np.array([0.0, 0.0, 1.0])
        # three.js: quat = (cross([0,0,1], n), dot([0,0,1], n) + 1), dann normalisieren.
        r = float(n[2]) + 1.0
        if r < 1e-12:
            # antiparallel (n == [0,0,-1]): three.js else-Zweig für vFrom=[0,0,1] → (0,-1,0,0)
            qx, qy, qz, qw = 0.0, -1.0, 0.0, 0.0
        else:
            qx, qy, qz, qw = -float(n[1]), float(n[0]), 0.0, r   # cross([0,0,1], n) = (-ny, nx, 0)
        ql = math.sqrt(qx * qx + qy * qy + qz * qz + qw * qw)
        qx, qy, qz, qw = qx / ql, qy / ql, qz / ql, qw / ql
        return np.array([
            [1 - 2 * (qy * qy + qz * qz), 2 * (qx * qy - qz * qw),     2 * (qx * qz + qy * qw)],
            [2 * (qx * qy + qz * qw),     1 - 2 * (qx * qx + qz * qz), 2 * (qy * qz - qx * qw)],
            [2 * (qx * qz - qy * qw),     2 * (qy * qz + qx * qw),     1 - 2 * (qx * qx + qy * qy)],
        ])

    @staticmethod
    def _marker_plane_corners(half: float) -> np.ndarray:
        """
        Die 4 Eckpunkte in der Marker-Ebene (+Z = Normale), in DERSELBEN
        Reihenfolge wie die von 3b_corner_marker_poses.py triangulierten
        `corners_m` (ArUco 0..3, im Uhrzeigersinn von der Vorderseite gesehen).

        Ecke 0 zeigt in Richtung (+h, +h) — dieselbe Konvention wie der visuell
        kalibrierte Orientierungszeiger (1,1,0) in boardViewer.html. Damit passt
        sie zu den manuell/visuell gesetzten `spin`-Werten der ARM-Marker, die
        im Homing die Gelenkwinkel bestimmen (gegen echte corners_m am
        Seed-Pose verifiziert, test_robot_fk_corners.py, RMS ~1 mm).

        Hinweis: Die Board-Referenzmarker (Set A0) sind ~90° anders kalibriert,
        ihre Eckreihenfolge passt unter dieser Konvention NICHT — egal, weil
        Board der Root-Link ist: ihr Eck-Residuum ist konstant bzgl. der
        Gelenkvariablen und beeinflusst die Schätzung nicht (der robuste
        Huber-Verlust dämpft es als Ausreißer). Siehe Doc-Notiz.

        Die Drehrichtung 0→1→2→3 ist so, dass 3b's `corner_plane_normal()`
        (outward = -cross(e01,e02)) wieder +Z liefert — identisch zur
        Beobachtungs-Konvention.
        """
        h = float(half)
        return np.array([
            [ h,  h, 0.0],   # 0
            [ h, -h, 0.0],   # 1
            [-h, -h, 0.0],   # 2
            [-h,  h, 0.0],   # 3
        ])

    @classmethod
    def marker_corners_local(cls,
                             position: Sequence[float],
                             normal: Sequence[float],
                             size_mm: float,
                             spin_deg: float = 0.0) -> np.ndarray:
        """
        Die 4 Marker-Ecken im LINK-lokalen Frame (mm), Reihenfolge wie `corners_m`.

        Orientierung = Spin um die Normale ∘ Minimal-Rotation [0,0,1]→Normale,
        exakt wie boardViewer.html (qSpinLoc.multiply(qNormalLoc)).
        """
        n = np.asarray(normal, dtype=float)
        R = _rot_axis_angle(n, float(spin_deg)) @ cls._shortest_arc_R(n)
        plane = cls._marker_plane_corners(float(size_mm) / 2.0)   # (4,3)
        return np.asarray(position, dtype=float) + (R @ plane.T).T  # (4,3)

    @classmethod
    def marker_corners_world(cls,
                             transforms: Dict[str, np.ndarray],
                             link_name: str,
                             position: Sequence[float],
                             normal: Sequence[float],
                             size_mm: float,
                             spin_deg: float = 0.0) -> np.ndarray:
        """Die 4 Marker-Ecken im Weltframe (mm), Reihenfolge wie `corners_m`."""
        T = transforms.get(link_name, np.eye(4))
        local = cls.marker_corners_local(position, normal, size_mm, spin_deg)
        return np.array([transform_point(T, c) for c in local])

    def all_markers_world(self,
                          transforms: Dict[str, np.ndarray]
                          ) -> Dict[int, Dict[str, Any]]:
        """
        Returns
        -------
        dict  marker_id -> {world_mm, local_mm, link, normal_world, corners_world}

        corners_world: (4,3) Welt-mm in `corners_m`-Reihenfolge (für den
        marker_observation="corner_points"-Modus in 5_pose_estimation.py).
        """
        default_size = float((self.robot.get("markerDefaults", {}) or {}).get("size", 25.0))
        result: Dict[int, Dict[str, Any]] = {}
        for lname, ldata in self.links.items():
            T = transforms.get(lname, np.eye(4))
            R = T[:3, :3]
            for m in ldata.get("markers", []):
                mid = int(m.get("id", -1))
                if mid < 0 or "position" not in m:
                    continue
                loc = np.array(m["position"], dtype=float)
                nor = np.array(m.get("normal", [0, 0, 1]), dtype=float)
                size_mm = float(m.get("size", default_size))
                spin_deg = float(m.get("spin", 0.0))
                local_corners = self.marker_corners_local(loc, nor, size_mm, spin_deg)
                result[mid] = {
                    "world_mm":    transform_point(T, loc),
                    "local_mm":    loc,
                    "link":        lname,
                    "normal_world": (R @ nor) / max(np.linalg.norm(R @ nor), 1e-12),
                    "corners_world": np.array([transform_point(T, c) for c in local_corners]),
                }
        return result

    # ── x-axis invariant helpers (used by 4a) ────────────────

    def link_x_at_zero_state(self, link_name: str) -> float:
        """
        Return the world x-coordinate of the link-frame origin
        when ALL joint variables are zero.

        For any link reachable via only x-axis rotations (Arm1, Ellbow, Arm2),
        this value is constant regardless of the actual revolute angles.
        Adding the slider value x_mm gives the true link origin x:
            link_origin_world_x = x_mm + link_x_at_zero_state(link_name)

        And for any marker in that link:
            marker_world_x = x_mm + link_x_at_zero_state(link_name) + marker_local_x
        """
        T = self.compute({k: 0.0 for k in STATE_KEYS})
        return float(T[link_name][0, 3])

    def joint_origin_world(self,
                            link_name: str,
                            joint_state: Dict[str, float]) -> np.ndarray:
        """
        World position of the pivot that link_name rotates / slides around.
        """
        d = self.links[link_name]
        parent = d.get("parent")
        T_all = self.compute(joint_state)
        T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)

        mp  = d.get("mountPosition", [0, 0, 0])
        mr  = d.get("mountRotation", [0, 0, 0])
        T_m = make_T(_rot_xyz_euler(*mr), mp)

        ji  = d.get("jointToParent", {}) or {}
        jp  = ji.get("origin", [0, 0, 0])
        jr  = ji.get("rotation", [0, 0, 0])
        T_j = make_T(_rot_xyz_euler(*jr), jp)

        return transform_point(T_parent @ T_m @ T_j, [0, 0, 0])

    def joint_axis_world(self,
                          link_name: str,
                          joint_state: Dict[str, float]) -> np.ndarray:
        """
        Joint axis of link_name expressed in world frame.
        """
        d = self.links[link_name]
        parent = d.get("parent")
        T_all = self.compute(joint_state)
        T_parent = T_all.get(parent, np.eye(4)) if parent else np.eye(4)

        mp  = d.get("mountPosition", [0, 0, 0])
        mr  = d.get("mountRotation", [0, 0, 0])
        T_m = make_T(_rot_xyz_euler(*mr), mp)

        ji  = d.get("jointToParent", {}) or {}
        jp  = ji.get("origin", [0, 0, 0])
        jr  = ji.get("rotation", [0, 0, 0])
        T_j = make_T(_rot_xyz_euler(*jr), jp)

        R_to_pivot = (T_parent @ T_m @ T_j)[:3, :3]
        axis_local = np.asarray(ji.get("axis", [1, 0, 0]), dtype=float)
        world = R_to_pivot @ axis_local
        n = float(np.linalg.norm(world))
        return world / n if n > 1e-12 else world

    # ── utility ───────────────────────────────────────────────

    def chain(self, link_name: str) -> List[str]:
        """Return chain from root to link_name (inclusive)."""
        out, cur = [], link_name
        while cur:
            out.append(cur)
            cur = self.links.get(cur, {}).get("parent")
        return list(reversed(out))

    def board_marker_world_positions(self, length_scale: float = 1.0) -> Dict[int, np.ndarray]:
        """
        Return known world positions for all Board markers (in mm).
        Board is the root, so its marker positions ARE world positions.
        length_scale: 1/1000 if robot.json uses mm.
        """
        board = self.links.get("Board", {})
        result: Dict[int, np.ndarray] = {}
        for m in board.get("markers", []):
            mid = int(m.get("id", -1))
            if mid >= 0 and "position" in m:
                p = np.array(m["position"], dtype=float) * length_scale
                result[mid] = p
        return result