"""Pose comparison utilities for hardware integration tests.

Provides trajectory comparison using distance_between_poses() from
the reachy_mini SDK.  Returns per-frame and aggregate metrics (RMSE,
max, mean) for translation, rotation, and "magic distance".
"""

import numpy as np
from dataclasses import dataclass

from reachy_mini.utils.interpolation import distance_between_poses


@dataclass
class TrajectoryMetrics:
    """Aggregate metrics from comparing two trajectories."""

    n_frames: int
    # Magic distance (mm + degrees combined)
    magic_mean: float
    magic_rmse: float
    magic_max: float
    magic_p95: float
    # Translation (meters)
    trans_mean: float
    trans_rmse: float
    trans_max: float
    # Rotation (radians)
    rot_mean: float
    rot_rmse: float
    rot_max: float
    # Per-frame arrays
    magic_distances: np.ndarray
    trans_distances: np.ndarray
    rot_distances: np.ndarray

    def summary(self) -> str:
        return (
            f"Trajectory comparison ({self.n_frames} frames):\n"
            f"  Magic:  mean={self.magic_mean:.1f}  rmse={self.magic_rmse:.1f}  "
            f"max={self.magic_max:.1f}  p95={self.magic_p95:.1f}\n"
            f"  Trans:  mean={self.trans_mean*1000:.1f}mm  "
            f"rmse={self.trans_rmse*1000:.1f}mm  max={self.trans_max*1000:.1f}mm\n"
            f"  Rot:    mean={np.degrees(self.rot_mean):.1f}deg  "
            f"rmse={np.degrees(self.rot_rmse):.1f}deg  max={np.degrees(self.rot_max):.1f}deg"
        )


def frames_to_poses(frames: list[dict]) -> list[np.ndarray]:
    """Extract head poses (4x4 matrices) from recorded frames."""
    poses = []
    for f in frames:
        head = f["head"]
        poses.append(np.array(head, dtype=np.float64))
    return poses


def compare_trajectories(
    ref_times: list[float],
    ref_frames: list[dict],
    rec_times: list[float],
    rec_frames: list[dict],
) -> TrajectoryMetrics:
    """Compare two recorded trajectories frame-by-frame.

    Interpolates the recorded trajectory timestamps to match the reference
    timestamps using nearest-neighbor lookup, then computes per-frame
    distances using distance_between_poses().

    Args:
        ref_times: timestamps from the reference recording
        ref_frames: frame dicts from the reference recording
        rec_times: timestamps from the new recording
        rec_frames: frame dicts from the new recording

    Returns:
        TrajectoryMetrics with per-frame and aggregate distances.
    """
    ref_poses = frames_to_poses(ref_frames)
    rec_poses = frames_to_poses(rec_frames)
    rec_times_arr = np.array(rec_times)

    trans_dists = []
    rot_dists = []
    magic_dists = []

    for i, (t, ref_pose) in enumerate(zip(ref_times, ref_poses)):
        # Find nearest recorded frame by timestamp
        idx = int(np.argmin(np.abs(rec_times_arr - t)))
        rec_pose = rec_poses[idx]

        d_trans, d_rot, d_magic = distance_between_poses(ref_pose, rec_pose)
        trans_dists.append(d_trans)
        rot_dists.append(d_rot)
        magic_dists.append(d_magic)

    trans_arr = np.array(trans_dists)
    rot_arr = np.array(rot_dists)
    magic_arr = np.array(magic_dists)

    return TrajectoryMetrics(
        n_frames=len(ref_times),
        magic_mean=float(np.mean(magic_arr)),
        magic_rmse=float(np.sqrt(np.mean(magic_arr**2))),
        magic_max=float(np.max(magic_arr)),
        magic_p95=float(np.percentile(magic_arr, 95)),
        trans_mean=float(np.mean(trans_arr)),
        trans_rmse=float(np.sqrt(np.mean(trans_arr**2))),
        trans_max=float(np.max(trans_arr)),
        rot_mean=float(np.mean(rot_arr)),
        rot_rmse=float(np.sqrt(np.mean(rot_arr**2))),
        rot_max=float(np.max(rot_arr)),
        magic_distances=magic_arr,
        trans_distances=trans_arr,
        rot_distances=rot_arr,
    )