from __future__ import annotations

import sysconfig
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import pinocchio as pin

from .types import JointLimits, physical_joint_limits


EXPECTED_JOINT_NAMES = tuple(f"joint_{index}" for index in range(1, 8))
FLANGE_FRAME = "link_7"


def default_urdf_path() -> Path:
    source_path = (
        Path(__file__).resolve().parents[2]
        / "kine_ctrl"
        / "urdf_rm75"
        / "RM75-B.urdf"
    )
    if source_path.is_file():
        return source_path
    installed_candidates = [
        Path(sysconfig.get_path("data"))
        / "share"
        / "rm75_ik"
        / "models"
        / "RM75-B.urdf"
    ]
    resolved = Path(__file__).resolve()
    if len(resolved.parents) > 4:
        installed_candidates.append(
            resolved.parents[4] / "share/rm75_ik/models/RM75-B.urdf"
        )
    for installed_path in installed_candidates:
        if installed_path.is_file():
            return installed_path
    raise FileNotFoundError("RM75-B.urdf was not found in source or installed data")


def validate_se3(value: pin.SE3, name: str = "pose") -> None:
    if not isinstance(value, pin.SE3):
        raise TypeError(f"{name} must be pinocchio.SE3")
    rotation = np.asarray(value.rotation)
    translation = np.asarray(value.translation)
    if rotation.shape != (3, 3) or translation.shape != (3,):
        raise ValueError(f"{name} has invalid dimensions")
    if not np.all(np.isfinite(rotation)) or not np.all(np.isfinite(translation)):
        raise ValueError(f"{name} must be finite")
    if not np.allclose(rotation.T @ rotation, np.eye(3), atol=1e-7):
        raise ValueError(f"{name} rotation must be orthonormal")
    if not np.isclose(np.linalg.det(rotation), 1.0, atol=1e-7):
        raise ValueError(f"{name} rotation determinant must be +1")


def pose_errors(current: pin.SE3, target: pin.SE3) -> Tuple[float, float]:
    validate_se3(current, "current")
    validate_se3(target, "target")
    position_error = float(np.linalg.norm(current.translation - target.translation))
    rotation_delta = current.rotation.T @ target.rotation
    orientation_error = float(np.linalg.norm(pin.log3(rotation_delta)))
    return position_error, orientation_error


class RM75Kinematics:
    """Pinocchio kinematics for one RM75-B.

    Instances own mutable Pinocchio data and are intentionally not thread-safe.
    Use one instance per arm/control thread.
    """

    def __init__(
        self,
        urdf_path: Optional[Path | str] = None,
        limits: Optional[JointLimits] = None,
    ) -> None:
        self.urdf_path = Path(urdf_path) if urdf_path is not None else default_urdf_path()
        if not self.urdf_path.is_file():
            raise FileNotFoundError(self.urdf_path)
        self.model = pin.buildModelFromUrdf(str(self.urdf_path))
        if self.model.nq != 7 or self.model.nv != 7:
            raise ValueError(
                f"expected RM75 model nq=nv=7, got nq={self.model.nq}, nv={self.model.nv}"
            )
        joint_names = tuple(self.model.names[1:])
        if joint_names != EXPECTED_JOINT_NAMES:
            raise ValueError(f"unexpected RM75 joint order: {joint_names}")
        frame_id = self.model.getFrameId(FLANGE_FRAME)
        if frame_id >= len(self.model.frames):
            raise ValueError(f"missing flange frame {FLANGE_FRAME!r}")
        self.flange_frame_id = frame_id
        self.limits = limits or physical_joint_limits()
        self.model.lowerPositionLimit[:7] = self.limits.lower
        self.model.upperPositionLimit[:7] = self.limits.upper
        self.data = self.model.createData()

    def validate_q(self, q_rad: np.ndarray, *, require_within_limits: bool = True) -> np.ndarray:
        q = np.asarray(q_rad, dtype=float)
        if q.shape != (7,):
            raise ValueError(f"RM75 configuration must have shape (7,), got {q.shape}")
        if not np.all(np.isfinite(q)):
            raise ValueError("RM75 configuration must be finite")
        if require_within_limits and not self.limits.contains(q):
            raise ValueError(f"configuration is outside {self.limits.name} joint limits")
        return q.copy()

    def forward(self, q_rad: np.ndarray, tool: Optional[pin.SE3] = None) -> pin.SE3:
        q = self.validate_q(q_rad)
        pin.framesForwardKinematics(self.model, self.data, q)
        flange = self.data.oMf[self.flange_frame_id]
        result = pin.SE3(flange.rotation.copy(), flange.translation.copy())
        if tool is not None:
            validate_se3(tool, "tool")
            result = result * tool
        return result

    def jacobian(self, q_rad: np.ndarray) -> np.ndarray:
        q = self.validate_q(q_rad)
        jacobian = pin.computeFrameJacobian(
            self.model,
            self.data,
            q,
            self.flange_frame_id,
            pin.ReferenceFrame.LOCAL,
        )
        return np.asarray(jacobian).copy()
