from __future__ import annotations

import math
import sys
import time
from pathlib import Path

import numpy as np

from xr_rm_teleop.placo_ik_solver import PlacoIkSolver


CASES = {
    "left": [-78.81, 3.22, 67.96, 97.12, 95.08, -81.11, -74.55],
    "right": [-86.10, 22.80, -89.57, 93.98, -91.82, -87.32, -89.35],
}

TOOL_CHAINS = {
    "left": ("scissor_base_link", "scissor_link_7", 0.165),
    "right": ("omnipic_base_link", "omnipic_link_7", 0.14),
}


def rotation_z(angle: float) -> np.ndarray:
    cosine = math.cos(angle)
    sine = math.sin(angle)
    return np.asarray(
        [
            [cosine, -sine, 0.0],
            [sine, cosine, 0.0],
            [0.0, 0.0, 1.0],
        ]
    )


def angle_error(actual: np.ndarray, target: np.ndarray) -> float:
    cosine = np.clip((np.trace(target @ actual.T) - 1.0) * 0.5, -1.0, 1.0)
    return float(math.acos(cosine))


def main() -> None:
    urdf_path = Path(sys.argv[1]).resolve()
    for arm, joint_degrees in CASES.items():
        initial_joints = np.deg2rad(joint_degrees)
        drift_solver = PlacoIkSolver(str(urdf_path), 1.0 / 125.0, arm)
        joints = initial_joints.tolist()
        stationary_target = drift_solver.update_joint_state(joints)
        base_frame, flange_frame, tcp_length = TOOL_CHAINS[arm]
        world_to_base = drift_solver._robot.get_T_world_frame(base_frame)
        world_to_flange = drift_solver._robot.get_T_world_frame(flange_frame)
        base_to_flange = np.linalg.inv(world_to_base) @ world_to_flange
        flange_to_tcp = np.linalg.inv(base_to_flange) @ stationary_target
        assert np.allclose(flange_to_tcp[:3, 3], [0.0, 0.0, tcp_length])
        assert np.allclose(flange_to_tcp[:3, :3], np.eye(3), atol=1e-5)
        for _ in range(250):
            drift_solver.update_joint_state(joints)
            joints = drift_solver.solve(stationary_target)
        drift_degrees = float(
            np.max(np.abs(np.rad2deg(np.asarray(joints) - initial_joints)))
        )
        assert drift_degrees <= 0.05, (
            f"{arm} stationary target drifted {drift_degrees:.3f}deg"
        )

        solver = PlacoIkSolver(str(urdf_path), 1.0 / 125.0, arm)
        joints = initial_joints.tolist()
        current = solver.update_joint_state(joints)
        assert current.shape == (4, 4)
        target = current.copy()
        target[0, 3] += 0.01
        target[:3, :3] = rotation_z(0.05) @ target[:3, :3]

        solve_durations = []
        for _ in range(250):
            solver.update_joint_state(joints)
            started_at = time.perf_counter()
            joints = solver.solve(target)
            solve_durations.append(time.perf_counter() - started_at)

        actual = solver.update_joint_state(joints)
        position_error = np.linalg.norm(actual[:3, 3] - target[:3, 3])
        orientation_error = angle_error(actual[:3, :3], target[:3, :3])
        assert len(joints) == 7
        assert np.isfinite(joints).all()
        assert np.allclose(
            solver.base_configuration,
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
        )
        assert position_error <= 0.005
        assert orientation_error <= math.radians(2.0)
        print(
            f"{arm}: position_error={position_error:.6f}m, "
            f"orientation_error={math.degrees(orientation_error):.3f}deg, "
            f"stationary_drift={drift_degrees:.3f}deg, "
            f"solve_avg={1000.0 * np.mean(solve_durations):.3f}ms, "
            f"solve_max={1000.0 * max(solve_durations):.3f}ms, "
            f"solve_overruns={sum(value > 1.0 / 125.0 for value in solve_durations)}"
        )


if __name__ == "__main__":
    main()