"""
_uwb_tracker.py — UWB DW3000 anchor/tag ranging + bearing estimation (Issue #365).

Pure-Python library (no ROS2 / hardware dependencies) so it can be fully unit-tested
on a development machine without the physical ESP32-UWB-Pro anchors attached.

Hardware layout
---------------
Two DW3000 anchors are mounted on the SaltyBot chassis, separated by a ~25 cm
baseline, both facing forward:

    anchor0 (left, x = −baseline/2)   anchor1 (right, x = +baseline/2)
           │←──────── baseline ────────→│
                        ↑
                   robot forward

The wearable tag is carried by the Tee (person being followed).

Bearing estimation
------------------
Given TWR (two-way ranging) distances d0 and d1 from the two anchors to the tag,
and the known baseline B between the anchors, we compute bearing θ using the
law of cosines:

    cos α = (d0² - d1² + B²) / (2 · B · d1)   [angle at anchor1]

Bearing is measured from the perpendicular bisector of the baseline (i.e. the
robot's forward axis):

    θ = 90° − α

Positive θ = tag is to the right, negative = tag is to the left.

The formula degrades when the tag is very close to the baseline or at extreme
angles. We report a confidence value that penalises these conditions.

Serial protocol (ESP32 → Orin via USB-serial)
---------------------------------------------
Each anchor sends newline-terminated ASCII frames at ≥10 Hz:

    RANGE,<anchor_id>,<tag_id>,<distance_mm>\n

Example:
    RANGE,0,T0,1532\n    → anchor 0, tag T0, distance 1532 mm
    RANGE,1,T0,1748\n    → anchor 1, tag T0, distance 1748 mm

A STATUS frame is sent once on connection:
    STATUS,<anchor_id>,OK\n

The node opens two serial ports (one per anchor).  A background thread reads
each port and updates a shared RangingState.

Kalman filter
-------------
A simple 1-D constant-velocity Kalman filter smooths the bearing output.
State: [bearing_deg, bearing_rate_dps].
"""

from __future__ import annotations

import math
import re
import threading
import time
from dataclasses import dataclass, field
from typing import Optional, Tuple

import numpy as np

# ── Constants ─────────────────────────────────────────────────────────────────

# Fix quality codes (written to UwbTarget.fix_quality)
FIX_NONE   = 0
FIX_SINGLE = 1   # only one anchor responding
FIX_DUAL   = 2   # both anchors responding — full bearing estimate

# Maximum age of a ranging measurement before it is considered stale (seconds)
_STALE_S = 0.5

# ── Serial protocol parsing ────────────────────────────────────────────────────

_RANGE_RE  = re.compile(r'^RANGE,(\d+),(\w+),([\d.]+)\s*$')
_STATUS_RE = re.compile(r'^STATUS,(\d+),(\w+)\s*$')


@dataclass
class RangeFrame:
    """Decoded RANGE frame from a DW3000 anchor."""
    anchor_id:   int
    tag_id:      str
    distance_m:  float
    timestamp:   float = field(default_factory=time.monotonic)


def parse_frame(line: str) -> Optional[RangeFrame]:
    """
    Parse one ASCII line from the anchor serial stream.

    Returns a RangeFrame on success, None for STATUS/unknown/malformed lines.

    Parameters
    ----------
    line : raw ASCII line (may include trailing whitespace / CR)

    Examples
    --------
    >>> f = parse_frame("RANGE,0,T0,1532")
    >>> f.anchor_id, f.tag_id, f.distance_m
    (0, 'T0', 1.532)
    >>> parse_frame("STATUS,0,OK") is None
    True
    >>> parse_frame("GARBAGE") is None
    True
    """
    m = _RANGE_RE.match(line.strip())
    if m is None:
        return None
    anchor_id  = int(m.group(1))
    tag_id     = m.group(2)
    distance_m = float(m.group(3)) / 1000.0  # mm → m
    return RangeFrame(anchor_id=anchor_id, tag_id=tag_id, distance_m=distance_m)


# ── Two-anchor bearing geometry ────────────────────────────────────────────────

def bearing_from_ranges(
    d0:         float,
    d1:         float,
    baseline_m: float,
) -> Tuple[float, float]:
    """
    Compute bearing to tag from two anchor distances.

    Parameters
    ----------
    d0 : distance from anchor-0 (left, at x = −baseline/2) to tag (m)
    d1 : distance from anchor-1 (right, at x = +baseline/2) to tag (m)
    baseline_m : separation between the two anchors (m)

    Returns
    -------
    (bearing_deg, confidence)
        bearing_deg : signed bearing in degrees (positive = right)
        confidence  : 0.0–1.0 quality estimate

    Notes
    -----
    Derived from law of cosines applied to the triangle
    (anchor0, anchor1, tag):
        cos(α) = (d0² − d1² + B²) / (2·B·d0)   where α is angle at anchor0
    Forward axis is the perpendicular bisector of the baseline, so:
        bearing = 90° − α_at_anchor1
    We use the symmetric formula through the midpoint:
        x_tag = (d0² - d1²) / (2·B)              [lateral offset from midpoint]
        y_tag = sqrt(d0² - (x_tag + B/2)²)       [forward distance]
        bearing = atan2(x_tag, y_tag) * 180/π
    """
    if baseline_m <= 0.0:
        raise ValueError(f'baseline_m must be positive, got {baseline_m}')
    if d0 <= 0.0 or d1 <= 0.0:
        raise ValueError(f'distances must be positive, got d0={d0} d1={d1}')

    B = baseline_m
    # Lateral offset of tag from midpoint of baseline (positive = towards anchor1 = right)
    x = (d0 * d0 - d1 * d1) / (2.0 * B)

    # Forward distance (Pythagorean)
    # anchor0 is at x_coord = -B/2 relative to midpoint
    y_sq = d0 * d0 - (x + B / 2.0) ** 2
    if y_sq < 0.0:
        # Triangle inequality violated — noisy ranging; clamp to zero
        y_sq = 0.0
    y = math.sqrt(y_sq)

    bearing_deg = math.degrees(math.atan2(x, max(y, 1e-3)))

    # ── Confidence ───────────────────────────────────────────────────────────
    # 1. Range agreement penalty — if |d0-d1| > sqrt(d0²+d1²) * factor, tag
    #    is almost directly on the baseline extension (extreme angle, poor geometry)
    mean_d = (d0 + d1) / 2.0
    diff   = abs(d0 - d1)
    # Geometric dilution: confidence falls as |bearing| approaches 90°
    bearing_penalty = math.cos(math.radians(bearing_deg))  # 1.0 at 0°, 0.0 at 90°
    bearing_penalty = max(0.0, bearing_penalty)

    # 2. Distance sanity: if tag is unreasonably close (< B) confidence drops
    dist_penalty = min(1.0, mean_d / max(B, 0.1))

    confidence = float(np.clip(bearing_penalty * dist_penalty, 0.0, 1.0))

    return bearing_deg, confidence


def bearing_single_anchor(
    d0:         float,
    baseline_m: float = 0.25,
) -> Tuple[float, float]:
    """
    Fallback bearing when only one anchor is responding.

    With a single anchor we cannot determine lateral position, so we return
    bearing=0 (directly ahead) with a low confidence proportional to 1/distance.
    This keeps the follow-me controller moving toward the target even in
    degraded mode.

    Returns (0.0, confidence) where confidence ≤ 0.3.
    """
    # Single-anchor confidence: ≤ 0.3, decreases with distance
    confidence = float(np.clip(0.3 * (2.0 / max(d0, 0.5)), 0.0, 0.3))
    return 0.0, confidence


# ── Kalman filter for bearing smoothing ───────────────────────────────────────

class BearingKalman:
    """
    1-D constant-velocity Kalman filter for bearing smoothing.

    State: [bearing_deg, bearing_rate_dps]
    """

    def __init__(
        self,
        process_noise_deg2:  float = 10.0,   # bearing process noise (deg²/frame)
        process_noise_rate2: float = 100.0,  # rate process noise
        meas_noise_deg2:     float = 4.0,    # bearing measurement noise (deg²)
    ) -> None:
        self._x = np.zeros(2, dtype=np.float64)    # [bearing, rate]
        self._P = np.diag([100.0, 1000.0])          # initial covariance
        self._Q = np.diag([process_noise_deg2, process_noise_rate2])
        self._R = np.array([[meas_noise_deg2]])
        self._F = np.array([[1.0, 1.0],             # state transition (dt=1 frame)
                            [0.0, 1.0]])
        self._H = np.array([[1.0, 0.0]])            # observation: bearing only
        self._initialised = False

    def predict(self) -> float:
        """Predict next bearing; returns predicted bearing_deg."""
        self._x = self._F @ self._x
        self._P = self._F @ self._P @ self._F.T + self._Q
        return float(self._x[0])

    def update(self, bearing_deg: float) -> float:
        """
        Update with a new bearing measurement.

        On first call, seeds the filter state.
        Returns smoothed bearing_deg.
        """
        if not self._initialised:
            self._x[0] = bearing_deg
            self._initialised = True
            return bearing_deg

        self.predict()

        y = np.array([[bearing_deg]]) - self._H @ self._x.reshape(-1, 1)
        S = self._H @ self._P @ self._H.T + self._R
        K = self._P @ self._H.T @ np.linalg.inv(S)
        self._x = self._x + (K @ y).ravel()
        self._P = (np.eye(2) - K @ self._H) @ self._P
        return float(self._x[0])

    @property
    def bearing_rate_dps(self) -> float:
        """Current bearing rate estimate (degrees/second at nominal 10 Hz)."""
        return float(self._x[1]) * 10.0  # per-frame rate → per-second


# ── Shared ranging state (thread-safe) ────────────────────────────────────────

@dataclass
class _AnchorState:
    distance_m: float = 0.0
    timestamp:  float = 0.0
    valid:      bool  = False


class UwbRangingState:
    """
    Thread-safe store for the most recent ranging measurement from each anchor.

    Updated by the serial reader threads, consumed by the ROS2 publish timer.
    """

    def __init__(
        self,
        baseline_m:    float = 0.25,
        stale_timeout: float = _STALE_S,
    ) -> None:
        self.baseline_m    = baseline_m
        self.stale_timeout = stale_timeout
        self._lock         = threading.Lock()
        self._anchors      = [_AnchorState(), _AnchorState()]
        self._kalman       = BearingKalman()

    def update_anchor(self, anchor_id: int, distance_m: float) -> None:
        """Record a new ranging measurement from anchor anchor_id (0 or 1)."""
        if anchor_id not in (0, 1):
            return
        with self._lock:
            s = self._anchors[anchor_id]
            s.distance_m = distance_m
            s.timestamp  = time.monotonic()
            s.valid      = True

    def compute(self) -> 'UwbResult':
        """
        Derive bearing, distance, confidence, and fix quality from latest ranges.

        Called at the publish rate (≥10 Hz).  Applies Kalman smoothing to bearing.
        """
        now = time.monotonic()
        with self._lock:
            a0 = self._anchors[0]
            a1 = self._anchors[1]
            v0 = a0.valid and (now - a0.timestamp) < self.stale_timeout
            v1 = a1.valid and (now - a1.timestamp) < self.stale_timeout
            d0 = a0.distance_m
            d1 = a1.distance_m

            if not v0 and not v1:
                return UwbResult(valid=False)

            if v0 and v1:
                bearing_raw, conf = bearing_from_ranges(d0, d1, self.baseline_m)
                fix_quality = FIX_DUAL
                distance_m  = (d0 + d1) / 2.0
            elif v0:
                bearing_raw, conf = bearing_single_anchor(d0, self.baseline_m)
                fix_quality = FIX_SINGLE
                distance_m  = d0
                d1 = 0.0
            else:
                bearing_raw, conf = bearing_single_anchor(d1, self.baseline_m)
                fix_quality = FIX_SINGLE
                distance_m  = d1
                d0 = 0.0

            bearing_smooth = self._kalman.update(bearing_raw)

        return UwbResult(
            valid         = True,
            bearing_deg   = bearing_smooth,
            distance_m    = distance_m,
            confidence    = conf,
            anchor0_dist  = d0,
            anchor1_dist  = d1,
            baseline_m    = self.baseline_m,
            fix_quality   = fix_quality,
        )


@dataclass
class UwbResult:
    """Computed UWB fix — mirrors UwbTarget.msg fields."""
    valid:        bool  = False
    bearing_deg:  float = 0.0
    distance_m:   float = 0.0
    confidence:   float = 0.0
    anchor0_dist: float = 0.0
    anchor1_dist: float = 0.0
    baseline_m:   float = 0.25
    fix_quality:  int   = FIX_NONE


# ── Serial reader (runs in background thread) ──────────────────────────────────

class AnchorSerialReader:
    """
    Background thread that reads from one anchor's serial port and calls
    state.update_anchor() on each valid RANGE frame.

    Designed to be used with a real serial.Serial object in production, or
    any file-like object with a readline() method for testing.
    """

    def __init__(
        self,
        anchor_id: int,
        port,           # serial.Serial or file-like (readline() interface)
        state:     UwbRangingState,
        logger=None,
    ) -> None:
        self._anchor_id = anchor_id
        self._port      = port
        self._state     = state
        self._log       = logger
        self._running   = False
        self._thread    = threading.Thread(target=self._run, daemon=True)

    def start(self) -> None:
        self._running = True
        self._thread.start()

    def stop(self) -> None:
        self._running = False

    def _run(self) -> None:
        while self._running:
            try:
                raw = self._port.readline()
                if isinstance(raw, bytes):
                    raw = raw.decode('ascii', errors='replace')
                frame = parse_frame(raw)
                if frame is not None:
                    self._state.update_anchor(frame.anchor_id, frame.distance_m)
            except Exception as e:
                if self._log:
                    self._log.warn(f'UWB anchor {self._anchor_id} read error: {e}')
                time.sleep(0.05)
