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": [-79.55, -9.99, 71.01, 101.45, 95.07, -84.47, -74.52],
    "right": [-90.14, 3.76, -86.89, 87.89, -96.53, -79.62, -90.04],
}


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)
        joints = initial_joints.tolist()
        stationary_target = drift_solver.update_joint_state(joints)
        flange = drift_solver._robot.get_T_world_frame("link_7")
        flange_to_tcp = np.linalg.inv(flange) @ stationary_target
        assert np.allclose(flange_to_tcp[:3, 3], [0.0, 0.0, 0.16])
        assert np.allclose(flange_to_tcp[:3, :3], np.eye(3))
        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)))
        )

        solver = PlacoIkSolver(str(urdf_path), 1.0 / 125.0)
        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()