#!/usr/bin/env python3
"""
uwb_position_node.py — ROS2 UWB position node for Jetson (Issue #546).

Reads JSON range measurements from an ESP32 UWB tag (DW3000) over USB
serial, trilaterates position from 3+ fixed infrastructure anchors, and
publishes the robot's absolute position in the map frame.

Serial protocol (one JSON object per line, UTF-8):
────────────────────────────────────────────────────
  Tag → Jetson (full frame, all anchors at once):
    {"ts": 123456, "ranges": [
        {"id": 0, "d_mm": 1500, "rssi": -65.0},
        {"id": 1, "d_mm": 2100, "rssi": -68.0},
        {"id": 2, "d_mm": 1800, "rssi": -70.0}
    ]}

  Alternate (per-anchor, one line per measurement):
    {"id": 0, "d_mm": 1500, "rssi": -65.0}

  Both "d_mm" and "range_mm" are accepted.

Anchor positions:
─────────────────
  Fixed anchors are configured in uwb_position_params.yaml as flat arrays:
    anchor_ids:       [0, 1, 2, 3]
    anchor_positions: [x0, y0, z0, x1, y1, z1, ...]   # metres in map frame

Publishes:
──────────
  /saltybot/uwb/pose          (geometry_msgs/PoseStamped)    10 Hz
  /saltybot/uwb/range/<id>    (saltybot_uwb_msgs/UwbRange)   per anchor, on arrival
  /saltybot/uwb/status        (std_msgs/String)              10 Hz JSON diagnostics

TF2:
────
  Broadcasts  uwb_link → map  from /saltybot/uwb/pose position.

Outlier rejection:
──────────────────
  After initial trilateration, anchors with range residual > outlier_threshold_m
  are dropped.  If ≥ 3 inliers remain (2D mode) the position is re-solved.
  Anchors rejected in N_strikes_max consecutive frames are flagged in status.
"""

import json
import math
import threading
import time
from typing import Dict, List, Optional, Tuple

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from geometry_msgs.msg import (
    PoseStamped, TransformStamped,
    Quaternion,
)
from std_msgs.msg import Header, String
from tf2_ros import TransformBroadcaster

from saltybot_uwb_msgs.msg import UwbRange

from saltybot_uwb_position.trilateration import (
    trilaterate,
    reject_outliers,
    residuals,
    KalmanFilter3D,
)

try:
    import serial
    _SERIAL_OK = True
except ImportError:
    _SERIAL_OK = False


# ── Serial reader thread ──────────────────────────────────────────────────────

class SerialJsonReader(threading.Thread):
    """Background thread: reads newline-delimited JSON from a serial port."""

    def __init__(self, port: str, baudrate: int, callback, logger) -> None:
        super().__init__(daemon=True)
        self._port     = port
        self._baudrate = baudrate
        self._callback = callback
        self._logger   = logger
        self._running  = False
        self._ser      = None

    def run(self) -> None:
        self._running = True
        while self._running:
            try:
                self._ser = serial.Serial(
                    self._port, self._baudrate, timeout=1.0
                )
                self._logger.info(f"UWB serial: opened {self._port} @ {self._baudrate}")
                self._read_loop()
            except Exception as exc:
                self._logger.warning(
                    f"UWB serial error ({self._port}): {exc} — retry in 2 s"
                )
                if self._ser and self._ser.is_open:
                    self._ser.close()
                time.sleep(2.0)

    def _read_loop(self) -> None:
        while self._running:
            try:
                raw = self._ser.readline()
                if not raw:
                    continue
                line = raw.decode("utf-8", errors="replace").strip()
                if not line:
                    continue
                try:
                    obj = json.loads(line)
                    self._callback(obj)
                except json.JSONDecodeError:
                    pass  # ignore malformed lines silently
            except Exception as exc:
                self._logger.warning(f"UWB read error: {exc}")
                break

    def stop(self) -> None:
        self._running = False
        if self._ser and self._ser.is_open:
            self._ser.close()


# ── Node ──────────────────────────────────────────────────────────────────────

class UwbPositionNode(Node):
    """ROS2 UWB position node — trilateration from 3+ fixed anchors."""

    _MIN_ANCHORS_2D = 3
    _MIN_ANCHORS_3D = 4

    def __init__(self) -> None:
        super().__init__("uwb_position")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("serial_port",          "/dev/ttyUSB0")
        self.declare_parameter("baudrate",             115200)

        # Anchor configuration: flat arrays, 3 floats per anchor
        self.declare_parameter("anchor_ids",       [0, 1, 2])
        self.declare_parameter("anchor_positions", [
            0.0, 0.0, 2.0,   # anchor 0: x, y, z (metres in map frame)
            5.0, 0.0, 2.0,   # anchor 1
            5.0, 5.0, 2.0,   # anchor 2
        ])

        self.declare_parameter("solve_z",              False)   # 2D mode by default
        self.declare_parameter("robot_z",              0.0)     # m — robot floor height
        self.declare_parameter("publish_rate",         10.0)    # Hz — pose + status
        self.declare_parameter("range_timeout_s",      1.5)     # s — stale range threshold
        self.declare_parameter("max_range_m",          30.0)    # m — validity max
        self.declare_parameter("min_range_m",          0.05)    # m — validity min
        self.declare_parameter("outlier_threshold_m",  0.4)     # m — RANSAC threshold
        self.declare_parameter("outlier_strikes_max",  5)       # frames before flagging
        self.declare_parameter("kf_process_noise",     0.05)    # Kalman Q
        self.declare_parameter("kf_meas_noise",        0.10)    # Kalman R
        self.declare_parameter("map_frame",            "map")
        self.declare_parameter("uwb_frame",            "uwb_link")

        # Load params
        self._port         = self.get_parameter("serial_port").value
        self._baudrate     = self.get_parameter("baudrate").value
        self._solve_z      = self.get_parameter("solve_z").value
        self._robot_z      = self.get_parameter("robot_z").value
        self._publish_rate = self.get_parameter("publish_rate").value
        self._timeout      = self.get_parameter("range_timeout_s").value
        self._max_r        = self.get_parameter("max_range_m").value
        self._min_r        = self.get_parameter("min_range_m").value
        self._outlier_thr  = self.get_parameter("outlier_threshold_m").value
        self._strikes_max  = self.get_parameter("outlier_strikes_max").value
        self._map_frame    = self.get_parameter("map_frame").value
        self._uwb_frame    = self.get_parameter("uwb_frame").value

        # Parse anchor config
        anchor_ids_raw  = self.get_parameter("anchor_ids").value
        anchor_pos_flat = self.get_parameter("anchor_positions").value
        self._anchor_ids, self._anchor_positions = self._parse_anchors(
            anchor_ids_raw, anchor_pos_flat
        )
        self.get_logger().info(
            f"UWB anchors: {[f'#{i}@({p[0]:.1f},{p[1]:.1f},{p[2]:.1f})' for i,p in zip(self._anchor_ids, self._anchor_positions)]}"
        )

        # ── State ─────────────────────────────────────────────────────────────
        self._lock   = threading.Lock()
        # anchor_id → (range_m, rssi, timestamp_monotonic)
        self._ranges: Dict[int, Tuple[float, float, float]] = {}
        self._strikes: Dict[int, int] = {aid: 0 for aid in self._anchor_ids}
        self._last_fix: Optional[Tuple[float, float, float]] = None
        self._has_fix  = False
        self._fix_age  = 0.0

        # Kalman filter
        self._kf = KalmanFilter3D(
            process_noise=self.get_parameter("kf_process_noise").value,
            measurement_noise=self.get_parameter("kf_meas_noise").value,
            dt=1.0 / self._publish_rate,
        )

        # Per-anchor outlier publishers (pre-create for configured anchors)
        sensor_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=5,
        )
        self._range_pubs: Dict[int, object] = {}
        for aid in self._anchor_ids:
            self._range_pubs[aid] = self.create_publisher(
                UwbRange, f"/saltybot/uwb/range/{aid}", sensor_qos
            )

        # Main publishers
        self._pose_pub   = self.create_publisher(PoseStamped, "/saltybot/uwb/pose",   10)
        self._status_pub = self.create_publisher(String,      "/saltybot/uwb/status", 10)

        # TF2
        self._tf_broadcaster = TransformBroadcaster(self)

        # ── Serial reader ──────────────────────────────────────────────────────
        if _SERIAL_OK:
            self._reader = SerialJsonReader(
                self._port, self._baudrate,
                callback=self._on_serial_json,
                logger=self.get_logger(),
            )
            self._reader.start()
        else:
            self.get_logger().warning(
                "pyserial not installed — serial I/O disabled (simulation mode)"
            )

        # ── Publish timer ──────────────────────────────────────────────────────
        self._prev_publish_t = self.get_clock().now().nanoseconds * 1e-9
        self.create_timer(1.0 / self._publish_rate, self._publish_cb)

        self.get_logger().info(
            f"UwbPositionNode ready — port={self._port} "
            f"anchors={len(self._anchor_ids)} "
            f"mode={'3D' if self._solve_z else '2D'} "
            f"rate={self._publish_rate:.0f}Hz"
        )

    # ── Anchor config parsing ─────────────────────────────────────────────────

    def _parse_anchors(
        self,
        ids_raw,
        pos_flat,
    ):
        ids = list(ids_raw)
        n   = len(ids)
        if len(pos_flat) < n * 3:
            raise ValueError(
                f"anchor_positions must have 3×anchor_ids entries. "
                f"Got {len(pos_flat)} values for {n} anchors."
            )
        positions = []
        for i in range(n):
            base = i * 3
            positions.append((
                float(pos_flat[base]),
                float(pos_flat[base + 1]),
                float(pos_flat[base + 2]),
            ))
        return ids, positions

    # ── Serial JSON callback (called from reader thread) ──────────────────────

    def _on_serial_json(self, obj: dict) -> None:
        """Parse incoming JSON and update range table."""
        now = time.monotonic()

        if "ranges" in obj:
            # Full-frame format: {"ts": ..., "ranges": [{id, d_mm, rssi}, ...]}
            for entry in obj["ranges"]:
                self._ingest_range_entry(entry, now)
        else:
            # Per-anchor format: {"id": 0, "d_mm": 1500, "rssi": -65.0}
            self._ingest_range_entry(obj, now)

    def _ingest_range_entry(self, entry: dict, timestamp: float) -> None:
        """Store a single anchor range measurement after validity checks."""
        try:
            anchor_id = int(entry.get("id", entry.get("anchor_id", -1)))
            # Accept both "d_mm" and "range_mm"
            raw_mm    = int(entry.get("d_mm", entry.get("range_mm", 0)))
            rssi      = float(entry.get("rssi", 0.0))
        except (KeyError, TypeError, ValueError):
            return

        if anchor_id not in self._anchor_ids:
            return  # unknown / unconfigured anchor

        range_m = raw_mm / 1000.0
        if range_m < self._min_r or range_m > self._max_r:
            return  # outside validity window

        with self._lock:
            self._ranges[anchor_id] = (range_m, rssi, timestamp)

        # Publish per-anchor range topic immediately
        self._publish_anchor_range(anchor_id, range_m, raw_mm, rssi)

    # ── Per-anchor range publisher ────────────────────────────────────────────

    def _publish_anchor_range(
        self, anchor_id: int, range_m: float, raw_mm: int, rssi: float
    ) -> None:
        if anchor_id not in self._range_pubs:
            return
        msg = UwbRange()
        msg.header.stamp    = self.get_clock().now().to_msg()
        msg.header.frame_id = self._map_frame
        msg.anchor_id       = anchor_id
        msg.range_m         = float(range_m)
        msg.raw_mm          = int(raw_mm)
        msg.rssi            = float(rssi)
        msg.tag_id          = ""
        self._range_pubs[anchor_id].publish(msg)

    # ── Main publish callback (10 Hz) ─────────────────────────────────────────

    def _publish_cb(self) -> None:
        now_ros  = self.get_clock().now()
        now_mono = time.monotonic()

        # Collect fresh ranges
        with self._lock:
            fresh = {
                aid: entry
                for aid, entry in self._ranges.items()
                if (now_mono - entry[2]) <= self._timeout
            }

        # Check minimum anchor count for chosen mode
        min_anch = self._MIN_ANCHORS_3D if self._solve_z else self._MIN_ANCHORS_2D
        state    = "no_fix"
        pos      = None
        used_ids: List[int] = []
        res_map: Dict[int, float] = {}

        if len(fresh) >= min_anch:
            active_ids = list(fresh.keys())
            anch_pos   = [self._anchor_positions[self._anchor_ids.index(i)]
                          for i in active_ids]
            anch_r     = [fresh[i][0] for i in active_ids]

            try:
                fixed_z = None if self._solve_z else self._robot_z
                raw_pos = trilaterate(anch_pos, anch_r, fixed_z=fixed_z)

                # Outlier rejection
                inlier_idx = reject_outliers(
                    anch_pos, anch_r, raw_pos,
                    threshold_m=self._outlier_thr,
                )
                res_all = residuals(anch_pos, anch_r, raw_pos)
                for k, aid in enumerate(active_ids):
                    res_map[aid] = round(res_all[k], 3)

                # Update consecutive strike counters
                for k, aid in enumerate(active_ids):
                    if k in inlier_idx:
                        self._strikes[aid] = 0
                    else:
                        self._strikes[aid] = self._strikes.get(aid, 0) + 1

                # Re-solve with inliers if any were rejected
                if len(inlier_idx) < len(active_ids) and len(inlier_idx) >= min_anch:
                    inlier_anch = [anch_pos[k] for k in inlier_idx]
                    inlier_r    = [anch_r[k]   for k in inlier_idx]
                    raw_pos     = trilaterate(inlier_anch, inlier_r, fixed_z=fixed_z)

                used_ids = [active_ids[k] for k in inlier_idx]
                pos = raw_pos

                # Kalman predict + update
                dt = now_ros.nanoseconds * 1e-9 - self._prev_publish_t
                self._kf.predict(dt=dt)
                self._kf.update(pos[0], pos[1], pos[2])
                pos = self._kf.position()

                self._last_fix = pos
                self._has_fix  = True
                self._fix_age  = 0.0
                state = "ok" if len(used_ids) >= min_anch else "degraded"

            except (ValueError, RuntimeError) as exc:
                self.get_logger().warning(f"Trilateration failed: {exc}")
                state = "degraded"

        elif self._has_fix:
            # KF predict-only: coast on last known position
            dt = now_ros.nanoseconds * 1e-9 - self._prev_publish_t
            self._kf.predict(dt=dt)
            pos = self._kf.position()
            self._fix_age += 1.0 / self._publish_rate
            state = "degraded" if self._fix_age < 5.0 else "no_fix"

        self._prev_publish_t = now_ros.nanoseconds * 1e-9

        if pos is None:
            self._publish_status(state, 0, [], {})
            return

        x, y, z = pos

        # ── Publish PoseStamped ────────────────────────────────────────────
        pose_msg = PoseStamped()
        pose_msg.header.stamp    = now_ros.to_msg()
        pose_msg.header.frame_id = self._map_frame
        pose_msg.pose.position.x = x
        pose_msg.pose.position.y = y
        pose_msg.pose.position.z = z
        pose_msg.pose.orientation.w = 1.0  # no orientation from UWB alone
        self._pose_pub.publish(pose_msg)

        # ── TF2: uwb_link → map ───────────────────────────────────────────
        tf_msg = TransformStamped()
        tf_msg.header.stamp    = now_ros.to_msg()
        tf_msg.header.frame_id = self._map_frame
        tf_msg.child_frame_id  = self._uwb_frame
        tf_msg.transform.translation.x = x
        tf_msg.transform.translation.y = y
        tf_msg.transform.translation.z = z
        tf_msg.transform.rotation.w    = 1.0
        self._tf_broadcaster.sendTransform(tf_msg)

        # ── Diagnostics ───────────────────────────────────────────────────
        self._publish_status(state, len(used_ids), used_ids, res_map)

    # ── Status publisher ──────────────────────────────────────────────────────

    def _publish_status(
        self,
        state:     str,
        n_anchors: int,
        used_ids:  List[int],
        res_map:   Dict[int, float],
    ) -> None:
        # Flag anchors with too many consecutive outlier strikes
        flagged = [
            aid for aid, s in self._strikes.items() if s >= self._strikes_max
        ]
        pos = self._kf.position() if self._has_fix else None
        status = {
            "state":          state,
            "active_anchors": n_anchors,
            "used_anchor_ids": used_ids,
            "flagged_anchors": flagged,
            "position": {
                "x": round(pos[0], 3),
                "y": round(pos[1], 3),
                "z": round(pos[2], 3),
            } if pos else None,
            "residuals_m": res_map,
            "fix_age_s":   round(self._fix_age, 2),
        }
        self._status_pub.publish(String(data=json.dumps(status)))


# ── Entry point ───────────────────────────────────────────────────────────────

def main(args=None) -> None:
    rclpy.init(args=args)
    node = UwbPositionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
