"""
ranging_math.py — Pure triangulation and Kalman filter helpers for UWB follow-me.

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

2-Anchor geometry
─────────────────
Anchors are mounted on the robot stem, separated by `sep` metres along the Y-axis:

    A0  (0, +sep/2, anchor_z)  ← port side
    A1  (0, -sep/2, anchor_z)  ← starboard side

The person tag is at (x_t, y_t, tag_z) in base_link.

Step 1 — height compensation:
    dz = anchor_z - tag_z
    r0_h = sqrt(max(0, r0² - dz²))
    r1_h = sqrt(max(0, r1² - dz²))

Step 2 — lateral offset (Y):
    y_t = (r1_h² - r0_h²) / (2 * sep)

Step 3 — forward distance (X):
    x_t = sqrt(max(0, r0_h² - (y_t - sep/2)²))

Returns (x_t, y_t); caller should treat negative x_t as 0.
"""

import math


# ── Bearing helper ────────────────────────────────────────────────────────────

def bearing_from_pos(x: float, y: float) -> float:
    """
    Compute horizontal bearing to a point (x, y) in base_link.

    Parameters
    ----------
    x : forward distance (metres)
    y : lateral offset (metres, positive = left of robot / CCW)

    Returns
    -------
    bearing_rad : bearing in radians, range -π .. +π
                  positive = target to the left (CCW)
                  0 = directly ahead
    """
    return math.atan2(y, x)


# ── Triangulation ─────────────────────────────────────────────────────────────

def triangulate_2anchor(
    r0: float,
    r1: float,
    sep: float,
    anchor_z: float = 0.0,
    tag_z: float = 0.0,
) -> tuple:
    """
    Triangulate person position in base_link from two anchor ranges.

    Parameters
    ----------
    r0 : range to anchor-0 (port,  y = +sep/2) in metres
    r1 : range to anchor-1 (starboard, y = -sep/2) in metres
    sep : anchor separation in metres (centre-to-centre)
    anchor_z : height of anchors above ground (metres, default 0)
    tag_z    : height of person tag above ground (metres, default 0)

    Returns
    -------
    (x_t, y_t) : person position in metres (base_link X = forward, Y = left)
                 x_t is clamped to ≥ 0 (can't be behind anchors geometrically).
    """
    if sep <= 0.0:
        raise ValueError(f"sep must be > 0, got {sep}")
    if r0 < 0.0 or r1 < 0.0:
        raise ValueError(f"ranges must be ≥ 0, got r0={r0}, r1={r1}")

    # Height correction — bring to horizontal plane
    dz = anchor_z - tag_z
    dz2 = dz * dz
    r0_h2 = max(0.0, r0 * r0 - dz2)
    r1_h2 = max(0.0, r1 * r1 - dz2)
    r0_h = math.sqrt(r0_h2)

    # Lateral (Y) offset — from circle-circle intersection formula
    y_t = (r1_h2 - r0_h2) / (2.0 * sep)

    # Forward (X) distance
    dx2 = r0_h2 - (y_t - sep / 2.0) ** 2
    x_t = math.sqrt(max(0.0, dx2))

    return x_t, y_t


# ── Simple 2-D Kalman filter ──────────────────────────────────────────────────

class KalmanFilter2D:
    """
    Constant-velocity Kalman filter for a 2-D position estimate.

    State vector: [x, y, vx, vy]ᵀ

    Parameters
    ----------
    process_noise : Q scalar — how much we trust the motion model
    measurement_noise : R scalar — how noisy the UWB measurement is
    dt : expected update interval in seconds (used for F matrix)
    """

    def __init__(
        self,
        process_noise: float = 0.1,
        measurement_noise: float = 0.3,
        dt: float = 0.05,
    ):
        self._q = process_noise
        self._r = measurement_noise
        self._dt = dt

        # State: [x, y, vx, vy]
        self._x = [0.0, 0.0, 0.0, 0.0]

        # Covariance matrix P (4×4, stored as flat list, row-major)
        init_pos_var = 1.0
        init_vel_var = 0.5
        self._P = [
            init_pos_var, 0.0, 0.0, 0.0,
            0.0, init_pos_var, 0.0, 0.0,
            0.0, 0.0, init_vel_var, 0.0,
            0.0, 0.0, 0.0, init_vel_var,
        ]
        self._initialised = False

    # ── Matrix helpers (minimal, no numpy — importable without GPU deps) ──────

    @staticmethod
    def _mat_add(A, B):
        return [a + b for a, b in zip(A, B)]

    @staticmethod
    def _mat_sub(A, B):
        return [a - b for a, b in zip(A, B)]

    @staticmethod
    def _mat_mul(A, B, n):
        """n×n matrix multiply."""
        C = [0.0] * (n * n)
        for i in range(n):
            for k in range(n):
                for j in range(n):
                    C[i * n + j] += A[i * n + k] * B[k * n + j]
        return C

    @staticmethod
    def _mat_T(A, n):
        """n×n transpose."""
        T = [0.0] * (n * n)
        for i in range(n):
            for j in range(n):
                T[j * n + i] = A[i * n + j]
        return T

    def _build_F(self):
        dt = self._dt
        return [
            1.0, 0.0, dt,  0.0,
            0.0, 1.0, 0.0, dt,
            0.0, 0.0, 1.0, 0.0,
            0.0, 0.0, 0.0, 1.0,
        ]

    def _build_Q(self):
        q = self._q
        dt = self._dt
        dt2 = dt * dt
        dt3 = dt2 * dt
        dt4 = dt3 * dt
        # Continuous white-noise model discretised for constant-velocity
        return [
            q * dt4 / 4, 0.0,         q * dt3 / 2, 0.0,
            0.0,         q * dt4 / 4, 0.0,         q * dt3 / 2,
            q * dt3 / 2, 0.0,         q * dt2,     0.0,
            0.0,         q * dt3 / 2, 0.0,         q * dt2,
        ]

    def predict(self, dt: float = None):
        """Predict step — advances state by dt (defaults to self._dt)."""
        if dt is not None:
            self._dt = dt
        F = self._build_F()
        Q = self._build_Q()
        n = 4
        x = self._x
        # x = F @ x
        new_x = [
            F[0] * x[0] + F[1] * x[1] + F[2] * x[2] + F[3] * x[3],
            F[4] * x[0] + F[5] * x[1] + F[6] * x[2] + F[7] * x[3],
            F[8] * x[0] + F[9] * x[1] + F[10] * x[2] + F[11] * x[3],
            F[12] * x[0] + F[13] * x[1] + F[14] * x[2] + F[15] * x[3],
        ]
        # P = F @ P @ F.T + Q
        FP  = self._mat_mul(F, self._P, n)
        FT  = self._mat_T(F, n)
        FPF = self._mat_mul(FP, FT, n)
        self._P = self._mat_add(FPF, Q)
        self._x = new_x

    def update(self, x_meas: float, y_meas: float):
        """
        Update step with a 2-D position measurement.

        H = [[1,0,0,0],[0,1,0,0]]  (observe position only)
        """
        if not self._initialised:
            self._x[0] = x_meas
            self._x[1] = y_meas
            self._initialised = True
            return

        # Innovation
        innov_x = x_meas - self._x[0]
        innov_y = y_meas - self._x[1]

        # S = H @ P @ H.T + R  (2×2, position rows/cols only)
        P = self._P
        r = self._r
        # H picks rows 0,1 and cols 0,1 from P
        S00 = P[0]  + r
        S01 = P[1]
        S10 = P[4]
        S11 = P[5]  + r

        # K = P @ H.T @ S⁻¹  (4×2 Kalman gain)
        det = S00 * S11 - S01 * S10
        if abs(det) < 1e-12:
            return
        inv_det = 1.0 / det
        # S⁻¹
        Si00 =  S11 * inv_det
        Si01 = -S01 * inv_det
        Si10 = -S10 * inv_det
        Si11 =  S00 * inv_det

        # P @ H.T = first two columns of P (H.T selects cols 0,1)
        PH = [P[0], P[1], P[4], P[5], P[8], P[9], P[12], P[13]]  # 4×2

        # K (4×2) = PH @ Si
        K = [
            PH[0] * Si00 + PH[1] * Si10,   PH[0] * Si01 + PH[1] * Si11,
            PH[2] * Si00 + PH[3] * Si10,   PH[2] * Si01 + PH[3] * Si11,
            PH[4] * Si00 + PH[5] * Si10,   PH[4] * Si01 + PH[5] * Si11,
            PH[6] * Si00 + PH[7] * Si10,   PH[6] * Si01 + PH[7] * Si11,
        ]

        # x = x + K @ innov
        self._x[0] += K[0] * innov_x + K[1] * innov_y
        self._x[1] += K[2] * innov_x + K[3] * innov_y
        self._x[2] += K[4] * innov_x + K[5] * innov_y
        self._x[3] += K[6] * innov_x + K[7] * innov_y

        # P = (I - K @ H) @ P  (Joseph form for numerical stability skipped —
        # this is a low-latency embedded context where simplicity wins)
        n = 4
        # K @ H selects first two columns → rows of P
        KH = [0.0] * 16
        for row in range(4):
            KH[row * 4 + 0] = K[row * 2 + 0]
            KH[row * 4 + 1] = K[row * 2 + 1]
        # I - KH
        IKH = [-KH[i] for i in range(16)]
        for i in range(4):
            IKH[i * 4 + i] += 1.0
        self._P = self._mat_mul(IKH, self._P, n)

    def position(self) -> tuple:
        """Return current (x, y) position estimate."""
        return self._x[0], self._x[1]

    def velocity(self) -> tuple:
        """Return current (vx, vy) velocity estimate."""
        return self._x[2], self._x[3]

    def reset(self):
        """Reset filter to uninitialised state."""
        self._x = [0.0, 0.0, 0.0, 0.0]
        init_pos_var = 1.0
        init_vel_var = 0.5
        self._P = [
            init_pos_var, 0.0, 0.0, 0.0,
            0.0, init_pos_var, 0.0, 0.0,
            0.0, 0.0, init_vel_var, 0.0,
            0.0, 0.0, 0.0, init_vel_var,
        ]
        self._initialised = False
