import math
from pathlib import Path
from types import SimpleNamespace
from xml.etree import ElementTree

import numpy as np
import pytest

from xr_rm_teleop.placo_ik_solver import (
    QP_ORIENTATION_TOLERANCE_RAD,
    QP_POSITION_TOLERANCE_M,
    PlacoIkSolver,
    _validated_transform,
)

DUAL_URDF_PATH = (
    Path(__file__).resolve().parents[1]
    / "models"
    / "dual_rm75"
    / "Dual_arm.urdf"
)
ARM_CASES = (
    (
        "left",
        [-78.81, 3.22, 67.96, 97.12, 95.08, -81.11, -74.55],
        list(range(14, 21)),
        list(range(13, 20)),
        "omnipic",
    ),
    (
        "right",
        [-86.10, 22.80, -89.57, 93.98, -91.82, -87.32, -89.35],
        list(range(7, 14)),
        list(range(6, 13)),
        "scissor",
    ),
)


def test_dual_urdf_has_expected_joints_meshes_and_tcp_frames() -> None:
    root = ElementTree.parse(DUAL_URDF_PATH).getroot()
    moving_joint_names = [
        joint.attrib["name"]
        for joint in root.findall("joint")
        if joint.attrib["type"] != "fixed"
    ]
    mesh_filenames = [
        mesh.attrib["filename"]
        for mesh in root.findall(".//mesh")
    ]
    fixed_joints = {
        "omnipic_base_mount_joint": (
            "dual_arm_base_link",
            "omnipic_base_link",
            None,
        ),
        "scissor_base_mount_joint": (
            "dual_arm_base_link",
            "scissor_base_link",
            None,
        ),
        "omnipic_OmniPic_tcp_fixed": (
            "omnipic_gripper_link",
            "omnipic_OmniPic_tcp",
            "0 0 0.14",
        ),
        "scissor_scissor_tcp_fixed": (
            "scissor_scissor_link",
            "scissor_scissor_tcp",
            "0 0 0",
        ),
        "scissor_scissor_fixed_joint": (
            "scissor_link_7",
            "scissor_scissor_link",
            "0 0 0.165",
        ),
    }

    assert moving_joint_names == [
        *[f"omnipic_joint_{index}" for index in range(1, 8)],
        *[f"scissor_joint_{index}" for index in range(1, 8)],
    ]
    assert all(filename.startswith("meshes/") for filename in mesh_filenames)
    for name, (parent, child, xyz) in fixed_joints.items():
        joint = root.find(f"joint[@name='{name}']")
        assert joint is not None
        assert joint.attrib["type"] == "fixed"
        assert joint.find("parent").attrib["link"] == parent
        assert joint.find("child").attrib["link"] == child
        if xyz is not None:
            assert joint.find("origin").attrib["xyz"] == xyz


def _dual_placo_solver(
    arm: str,
    joint_degrees: list[float],
) -> tuple[PlacoIkSolver, list[float]]:
    pytest.importorskip("placo")
    joints = [math.radians(value) for value in joint_degrees]
    return PlacoIkSolver(str(DUAL_URDF_PATH), 1.0 / 90.0, arm), joints


@pytest.mark.parametrize(
    "arm,joint_degrees,q_offsets,v_offsets,inactive_prefix",
    ARM_CASES,
)
def test_solver_uses_arm_specific_offsets(
    arm: str,
    joint_degrees: list[float],
    q_offsets: list[int],
    v_offsets: list[int],
    inactive_prefix: str,
) -> None:
    solver, _ = _dual_placo_solver(arm, joint_degrees)

    assert solver._q_offsets == q_offsets
    assert solver._v_offsets == v_offsets


@pytest.mark.parametrize(
    "arm,joint_degrees,q_offsets,v_offsets,inactive_prefix",
    ARM_CASES,
)
def test_joint_state_pose_is_relative_to_selected_arm_base(
    arm: str,
    joint_degrees: list[float],
    q_offsets: list[int],
    v_offsets: list[int],
    inactive_prefix: str,
) -> None:
    solver, joints = _dual_placo_solver(arm, joint_degrees)

    actual_pose = solver.update_joint_state(joints)
    world_base = solver._robot.get_T_world_frame(solver._base_frame)
    world_tcp = solver._robot.get_T_world_frame(solver._tcp_frame)

    assert actual_pose == pytest.approx(np.linalg.inv(world_base) @ world_tcp)


@pytest.mark.parametrize(
    "arm,joint_degrees,q_offsets,v_offsets,inactive_prefix",
    ARM_CASES,
)
def test_qp_solve_converges_without_moving_inactive_arm(
    arm: str,
    joint_degrees: list[float],
    q_offsets: list[int],
    v_offsets: list[int],
    inactive_prefix: str,
) -> None:
    solver, joints = _dual_placo_solver(arm, joint_degrees)
    start_pose = solver.update_joint_state(joints)
    inactive_q_offsets = [
        solver._robot.get_joint_offset(f"{inactive_prefix}_joint_{index}")
        for index in range(1, 8)
    ]
    inactive_before = solver._robot.state.q[inactive_q_offsets].copy()
    target_pose = start_pose.copy()
    target_pose[0, 3] += 0.01

    result = solver.solve(target_pose)
    reached_pose = solver.update_joint_state(result)
    position_error = np.linalg.norm(
        target_pose[:3, 3] - reached_pose[:3, 3]
    )
    rotation_delta = (
        target_pose[:3, :3] @ reached_pose[:3, :3].T
    )
    orientation_error = math.acos(
        float(
            np.clip(
                (np.trace(rotation_delta) - 1.0) * 0.5,
                -1.0,
                1.0,
            )
        )
    )

    assert np.asarray(result).shape == (7,)
    assert np.isfinite(result).all()
    assert position_error <= QP_POSITION_TOLERANCE_M
    assert orientation_error <= QP_ORIENTATION_TOLERANCE_RAD
    assert solver._robot.state.q[inactive_q_offsets] == pytest.approx(
        inactive_before
    )


def test_solver_rejects_unknown_arm() -> None:
    pytest.importorskip("placo")

    with pytest.raises(ValueError, match="arm must be left or right"):
        PlacoIkSolver(str(DUAL_URDF_PATH), 1.0 / 90.0, "middle")


def test_qp_solve_accepts_position_error_within_two_millimeters() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._actual_joints = np.zeros(7)
    solver._robot = SimpleNamespace(
        state=SimpleNamespace(q=np.zeros(14))
    )
    solver._frame_task = SimpleNamespace(T_world_frame=None)
    solver._target_errors = lambda: (1.5e-3, 0.0)

    result = solver.solve(np.eye(4))

    assert result == pytest.approx([0.0] * 7)


def test_qp_solve_rejects_position_error_above_two_millimeters() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._actual_joints = np.zeros(7)
    solver._robot = SimpleNamespace(
        state=SimpleNamespace(q=np.zeros(14)),
        update_kinematics=lambda: None,
    )
    solver._frame_task = SimpleNamespace(T_world_frame=None)
    solver._solver = SimpleNamespace(solve=lambda update: None)
    solver._validate_result = lambda result, previous: None
    solver._target_errors = lambda: (2.1e-3, 0.0)

    with pytest.raises(RuntimeError, match="QP did not converge after 30"):
        solver.solve(np.eye(4))


def test_validated_transform_accepts_finite_se3_and_returns_a_copy() -> None:
    transform = np.eye(4)
    transform[:3, 3] = [0.3, -0.1, 0.2]

    actual = _validated_transform(transform)

    assert actual == pytest.approx(transform)
    assert actual is not transform


@pytest.mark.parametrize(
    "transform",
    [
        np.eye(3),
        np.full((4, 4), np.nan),
        np.vstack([np.eye(3, 4), [0.0, 0.0, 0.0, 2.0]]),
        np.diag([2.0, 1.0, 1.0, 1.0]),
    ],
)
def test_validated_transform_rejects_invalid_se3(transform: np.ndarray) -> None:
    with pytest.raises(ValueError):
        _validated_transform(transform)


def test_qp_result_rejects_nan_position_and_velocity_violations() -> None:
    solver = object.__new__(PlacoIkSolver)
    solver._joint_limits = np.asarray([[-1.0, 1.0]] * 7)
    solver._velocity_limits = np.ones(7)
    solver._dt = 0.1
    solver._actual_joints = np.zeros(7)

    with pytest.raises(ValueError, match="finite"):
        solver._validate_result(np.full(7, np.nan))
    with pytest.raises(ValueError, match="position"):
        solver._validate_result(np.full(7, 2.0))
    with pytest.raises(ValueError, match="velocity"):
        solver._validate_result(np.full(7, 0.2))