"""
mds_math.py — Classical Multi-Dimensional Scaling for UWB anchor calibration
             (Issue #602)

Given an N×N symmetric distance matrix D (in metres), recover the 2-D
Cartesian positions of N anchors using classical (metric) MDS.

Algorithm
─────────
  1. Compute the double-centred squared-distance matrix B = -½ H D² H
     where H = I - (1/N) 11ᵀ  (centering matrix).
  2. Eigen-decompose B = V Λ Vᵀ (symmetric).
  3. Keep the 2 largest positive eigenvalues (λ₁, λ₂).
  4. Positions X = V[:, :2] · diag(√λ₁, √λ₂).

The recovered configuration is up to a global rigid transformation
(rotation, reflection, translation).  The caller is responsible for
anchoring the coordinate frame (e.g. fix anchor-0 at origin, anchor-1
on the positive-X axis).

Usage
─────
  from saltybot_uwb_calibration.mds_math import mds_2d, anchor_frame_align

  D = np.array([[0, d01, d02, d03],
                [d01, 0, d12, d13],
                ...])
  positions, residual = mds_2d(D)          # (N,2) array, float RMS (m)
  positions = anchor_frame_align(positions) # align frame: anchor-0 at origin
"""

from __future__ import annotations

import numpy as np


def mds_2d(D: np.ndarray) -> tuple[np.ndarray, float]:
    """
    Recover 2-D anchor positions from pairwise distance matrix D.

    Parameters
    ----------
    D : (N, N) symmetric ndarray of float
        Pairwise distances in metres.  Diagonal must be 0.
        N must be >= 3 (4+ recommended for overdetermined system).

    Returns
    -------
    positions : (N, 2) ndarray
        Recovered x, y coordinates in metres.
    residual_rms : float
        RMS of (D_measured - D_recovered) in metres across all unique pairs.

    Raises
    ------
    ValueError
        If D is not square, not symmetric, or N < 3.
    """
    D = np.asarray(D, dtype=float)
    n = D.shape[0]
    if D.ndim != 2 or D.shape[1] != n:
        raise ValueError(f"D must be square, got shape {D.shape}")
    if n < 3:
        raise ValueError(f"Need at least 3 anchors, got {n}")
    if not np.allclose(D, D.T, atol=1e-3):
        raise ValueError("Distance matrix D is not symmetric")

    # ── Step 1: double-centre D² ──────────────────────────────────────────
    D2 = D ** 2
    H = np.eye(n) - np.ones((n, n)) / n
    B = -0.5 * H @ D2 @ H

    # ── Step 2: eigendecomposition ────────────────────────────────────────
    eigvals, eigvecs = np.linalg.eigh(B)      # eigvals in ascending order

    # ── Step 3: keep top-2 positive eigenvalues ───────────────────────────
    idx = np.argsort(eigvals)[::-1]           # descending
    top2_idx = idx[:2]
    lam = eigvals[top2_idx]
    lam = np.maximum(lam, 0.0)               # clamp numerical negatives

    # ── Step 4: recover positions ─────────────────────────────────────────
    V = eigvecs[:, top2_idx]
    positions = V * np.sqrt(lam)             # (N, 2)

    # ── Residual RMS ──────────────────────────────────────────────────────
    residual_rms = _residual(D, positions)

    return positions, residual_rms


def anchor_frame_align(positions: np.ndarray) -> np.ndarray:
    """
    Align recovered positions so that:
      - anchor-0 is at the origin (0, 0)
      - anchor-1 lies on the positive-X axis (y = 0)

    This removes the MDS ambiguity (translation + rotation + reflection).

    Parameters
    ----------
    positions : (N, 2) ndarray

    Returns
    -------
    (N, 2) ndarray  — aligned positions
    """
    positions = np.array(positions, dtype=float)
    if positions.shape[0] < 2:
        return positions

    # Translate so anchor-0 is at origin
    positions -= positions[0]

    # Rotate so anchor-1 is on positive-X axis
    dx, dy = positions[1, 0], positions[1, 1]
    angle = np.arctan2(dy, dx)
    cos_a, sin_a = np.cos(-angle), np.sin(-angle)
    R = np.array([[cos_a, -sin_a],
                  [sin_a,  cos_a]])
    positions = (R @ positions.T).T

    # Ensure anchor-1 has positive X (handle reflection)
    if positions[1, 0] < 0:
        positions[:, 0] *= -1.0

    return positions


def _residual(D_meas: np.ndarray, positions: np.ndarray) -> float:
    """Compute RMS of (measured - recovered) distances over all unique pairs."""
    n = D_meas.shape[0]
    errors = []
    for i in range(n):
        for j in range(i + 1, n):
            d_rec = float(np.linalg.norm(positions[i] - positions[j]))
            errors.append(D_meas[i, j] - d_rec)
    if not errors:
        return 0.0
    return float(np.sqrt(np.mean(np.array(errors) ** 2)))
