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
from xr_rm_teleop.realman_adapter import ArmPose


CASES = {
    "left": (
        [-79.55, -9.99, 71.01, 101.45, 95.07, -84.47, -74.52],
        [0.0, 0.0, 0.19, 0.0, 0.0, 0.0, 1.0],
    ),
    "right": (
        [-90.14, 3.76, -86.89, 87.89, -96.53, -79.62, -90.04],
        [0.0, 0.0, 0.16, 0.0, 0.0, 0.0, 1.0],
    ),
}


def angle_error(actual: list[float], target: list[float]) -> float:
    deltas = [
        math.atan2(math.sin(a - b), math.cos(a - b))
        for a, b in zip(actual, target)
    ]
    return math.sqrt(sum(value * value for value in deltas))


def main() -> None:
    urdf_path = Path(sys.argv[1]).resolve()
    for arm, (joint_degrees, tool_pose) in CASES.items():
        solver = PlacoIkSolver(str(urdf_path), tool_pose, 1.0 / 90.0)
        joints = np.deg2rad(joint_degrees).tolist()
        current = solver.update_joint_state(joints)
        target = ArmPose(
            current.x + 0.01,
            current.y,
            current.z,
            current.rx,
            current.ry,
            current.rz + 0.05,
        )

        solve_durations = []
        for _ in range(45):
            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(
            np.asarray(actual.xyz()) - np.asarray(target.xyz())
        )
        orientation_error = angle_error(actual.rpy(), target.rpy())
        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"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 / 90.0 for value in solve_durations)}"
        )


if __name__ == "__main__":
    main()