"""
trilateration.py — Pure-math helpers for UWB trilateration (Issue #546).

No ROS2 dependencies — importable in unit tests without a ROS2 install.

Algorithm: linear least-squares from N ≥ 3 anchors
────────────────────────────────────────────────────
Given anchors A_i = (x_i, y_i, z_i) with measured ranges r_i, the robot
position p = (x, y, z) satisfies:

    (x - x_i)² + (y - y_i)² + (z - z_i)² = r_i²

Subtract the equation for anchor 0 from each subsequent equation to
linearise:

    2(x_i - x_0)x + 2(y_i - y_0)y + 2(z_i - z_0)z
        = r_0² - r_i² + ‖a_i‖² - ‖a_0‖²

This yields A·p = b where A is (N-1)×3 and b is (N-1)×1.
Solve with numpy least-squares (lstsq).

2D mode (solve_z=False):
    z is fixed (robot_z parameter, default 0).
    Reduce to 2 unknowns (x, y): subtract z terms from b, drop the z column.
    Needs N ≥ 3 anchors.

Outlier rejection:
    Compute residual for each anchor: |r_meas - ‖p - a_i‖|.
    Reject anchors with residual > threshold.  Repeat if enough remain.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple

import numpy as np


AnchorPos = Tuple[float, float, float]   # (x, y, z) in map frame (metres)
RangeM    = float                         # measured range (metres)


# ── Core trilateration ────────────────────────────────────────────────────────

def trilaterate(
    anchors: List[AnchorPos],
    ranges:  List[RangeM],
    fixed_z: Optional[float] = None,
) -> Tuple[float, float, float]:
    """Least-squares trilateration from N ≥ 3 anchor ranges.

    Parameters
    ----------
    anchors   : list of (x, y, z) anchor positions in the map frame (metres)
    ranges    : measured distances from robot to each anchor (metres)
    fixed_z   : if not None, treat robot z as this value and solve 2D only

    Returns
    -------
    (x, y, z) : estimated robot position in the map frame (metres).
                When fixed_z is given, z = fixed_z.

    Raises
    ------
    ValueError  : fewer than 3 anchors provided
    RuntimeError: linear system is rank-deficient (e.g., collinear anchors)
    """
    n = len(anchors)
    if n < 3:
        raise ValueError(f"Need ≥ 3 anchors for trilateration, got {n}")
    if len(ranges) != n:
        raise ValueError("len(anchors) != len(ranges)")

    a = np.array(anchors, dtype=np.float64)   # (N, 3)
    r = np.array(ranges,  dtype=np.float64)   # (N,)

    # Reference anchor (index 0)
    a0 = a[0]
    r0 = r[0]
    norm_a0_sq = float(np.dot(a0, a0))

    if fixed_z is not None:
        # 2D mode: solve for (x, y), z is known
        z = fixed_z
        A_rows = []
        b_rows = []
        for i in range(1, n):
            ai = a[i]
            ri = r[i]
            norm_ai_sq = float(np.dot(ai, ai))
            # rhs adjusted for known z
            rhs = (r0**2 - ri**2 + norm_ai_sq - norm_a0_sq
                   + 2.0 * (ai[2] - a0[2]) * z)
            A_rows.append([2.0 * (ai[0] - a0[0]),
                           2.0 * (ai[1] - a0[1])])
            b_rows.append(rhs)

        A_mat = np.array(A_rows, dtype=np.float64)
        b_vec = np.array(b_rows, dtype=np.float64)
        result, _, rank, _ = np.linalg.lstsq(A_mat, b_vec, rcond=None)
        if rank < 2:
            raise RuntimeError("Rank-deficient 2D trilateration system — collinear anchors?")
        return float(result[0]), float(result[1]), z

    else:
        # 3D mode: solve for (x, y, z) — needs N ≥ 4 for well-conditioned solve
        A_rows = []
        b_rows = []
        for i in range(1, n):
            ai = a[i]
            ri = r[i]
            norm_ai_sq = float(np.dot(ai, ai))
            rhs = r0**2 - ri**2 + norm_ai_sq - norm_a0_sq
            A_rows.append([2.0 * (ai[0] - a0[0]),
                           2.0 * (ai[1] - a0[1]),
                           2.0 * (ai[2] - a0[2])])
            b_rows.append(rhs)

        A_mat = np.array(A_rows, dtype=np.float64)
        b_vec = np.array(b_rows, dtype=np.float64)
        result, _, rank, _ = np.linalg.lstsq(A_mat, b_vec, rcond=None)
        if rank < 3:
            raise RuntimeError("Rank-deficient 3D trilateration system — coplanar anchors?")
        return float(result[0]), float(result[1]), float(result[2])


# ── Outlier rejection ─────────────────────────────────────────────────────────

def reject_outliers(
    anchors:    List[AnchorPos],
    ranges:     List[RangeM],
    position:   Tuple[float, float, float],
    threshold_m: float = 0.4,
) -> List[int]:
    """Return indices of inlier anchors (residual ≤ threshold_m).

    Parameters
    ----------
    anchors     : anchor positions in map frame
    ranges      : measured ranges (metres)
    position    : current position estimate (x, y, z)
    threshold_m : max allowable range residual to be an inlier

    Returns
    -------
    inlier_indices : sorted list of indices whose residual is within threshold
    """
    px, py, pz = position
    inliers = []
    for i, (anchor, r_meas) in enumerate(zip(anchors, ranges)):
        ax, ay, az = anchor
        r_pred = math.sqrt((px - ax)**2 + (py - ay)**2 + (pz - az)**2)
        residual = abs(r_meas - r_pred)
        if residual <= threshold_m:
            inliers.append(i)
    return inliers


def residuals(
    anchors:  List[AnchorPos],
    ranges:   List[RangeM],
    position: Tuple[float, float, float],
) -> List[float]:
    """Compute per-anchor range residual (metres) for diagnostics."""
    px, py, pz = position
    res = []
    for anchor, r_meas in zip(anchors, ranges):
        ax, ay, az = anchor
        r_pred = math.sqrt((px - ax)**2 + (py - ay)**2 + (pz - az)**2)
        res.append(abs(r_meas - r_pred))
    return res


# ── 3D Kalman filter (constant-velocity, position-only observations) ──────────

class KalmanFilter3D:
    """Constant-velocity Kalman filter for 3D UWB position.

    State: [x, y, z, vx, vy, vz]ᵀ
    Observation: [x, y, z]  (position only)
    """

    def __init__(
        self,
        process_noise:    float = 0.1,
        measurement_noise: float = 0.15,
        dt: float = 0.1,
    ) -> None:
        self._q    = process_noise
        self._r    = measurement_noise
        self._dt   = dt
        self._x    = np.zeros(6)
        self._P    = np.diag([1.0, 1.0, 1.0, 0.5, 0.5, 0.5])
        self._init = False

    def _build_F(self, dt: float) -> np.ndarray:
        F = np.eye(6)
        F[0, 3] = dt
        F[1, 4] = dt
        F[2, 5] = dt
        return F

    def _build_Q(self, dt: float) -> np.ndarray:
        q  = self._q
        dt2 = dt * dt
        dt3 = dt2 * dt
        dt4 = dt3 * dt
        Q = np.zeros((6, 6))
        # Position-velocity cross terms (constant white-noise model)
        for i in range(3):
            Q[i,   i  ] = q * dt4 / 4.0
            Q[i,   i+3] = q * dt3 / 2.0
            Q[i+3, i  ] = q * dt3 / 2.0
            Q[i+3, i+3] = q * dt2
        return Q

    def predict(self, dt: float | None = None) -> None:
        if dt is not None:
            self._dt = dt
        F = self._build_F(self._dt)
        Q = self._build_Q(self._dt)
        self._x = F @ self._x
        self._P = F @ self._P @ F.T + Q

    def update(self, x_m: float, y_m: float, z_m: float) -> None:
        if not self._init:
            self._x[:3] = [x_m, y_m, z_m]
            self._init = True
            return

        H = np.zeros((3, 6))
        H[0, 0] = 1.0
        H[1, 1] = 1.0
        H[2, 2] = 1.0
        R = np.eye(3) * self._r

        z_vec = np.array([x_m, y_m, z_m])
        innov = z_vec - H @ self._x
        S     = H @ self._P @ H.T + R
        K     = self._P @ H.T @ np.linalg.inv(S)
        self._x = self._x + K @ innov
        self._P = (np.eye(6) - K @ H) @ self._P

    def position(self) -> Tuple[float, float, float]:
        return float(self._x[0]), float(self._x[1]), float(self._x[2])

    def velocity(self) -> Tuple[float, float, float]:
        return float(self._x[3]), float(self._x[4]), float(self._x[5])

    def reset(self) -> None:
        self._x    = np.zeros(6)
        self._P    = np.diag([1.0, 1.0, 1.0, 0.5, 0.5, 0.5])
        self._init = False
