#!/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)

    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}
        """
        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)
                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),
                }
        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