"""Trajectory comparison metrics — RMSE, Absolute Trajectory Error, Relative Pose Error.

All inputs are Nx3 numpy arrays of ENU or local-tangent-plane positions in metres.
Callers convert from lat/lon via gps_denied.core.coordinates.CoordinateTransformer
before passing trajectories to these functions.
"""

from __future__ import annotations

import numpy as np


def trajectory_rmse(estimated: np.ndarray, ground_truth: np.ndarray) -> float:
    """Root-mean-square position error between two aligned trajectories.

    Both arrays must have the same shape (N, 3).
    No alignment is performed — caller is responsible for time synchronisation
    and (if needed) SE(3) alignment.
    """
    if estimated.shape != ground_truth.shape:
        raise ValueError(
            f"Shape mismatch: estimated={estimated.shape}, ground_truth={ground_truth.shape}"
        )
    diff = estimated - ground_truth
    sq = np.sum(diff * diff, axis=1)  # per-frame squared position error
    return float(np.sqrt(np.mean(sq)))


def absolute_trajectory_error(
    estimated: np.ndarray, ground_truth: np.ndarray
) -> dict[str, float]:
    """ATE: per-frame position error statistics."""
    if estimated.shape != ground_truth.shape:
        raise ValueError("shape mismatch")
    err = np.linalg.norm(estimated - ground_truth, axis=1)
    return {
        "rmse": float(np.sqrt(np.mean(err * err))),
        "mean": float(np.mean(err)),
        "median": float(np.median(err)),
        "max": float(np.max(err)),
        "std": float(np.std(err)),
    }


def relative_pose_error(
    estimated: np.ndarray, ground_truth: np.ndarray, delta: int = 1
) -> dict[str, float]:
    """RPE: error in relative motion over windows of `delta` frames."""
    if estimated.shape != ground_truth.shape:
        raise ValueError("shape mismatch")
    if delta < 1 or delta >= len(estimated):
        raise ValueError(f"delta must be in [1, N-1]; got delta={delta}, N={len(estimated)}")
    est_delta = estimated[delta:] - estimated[:-delta]
    gt_delta = ground_truth[delta:] - ground_truth[:-delta]
    err = np.linalg.norm(est_delta - gt_delta, axis=1)
    return {
        "rmse": float(np.sqrt(np.mean(err * err))),
        "mean": float(np.mean(err)),
        "max": float(np.max(err)),
    }
