"""
calibration_node.py — UWB anchor auto-calibration ROS2 service (Issue #602)

Service:  /saltybot/uwb/calibrate_anchors  (CalibrateAnchors)

Workflow
────────
  1. For every ordered pair (i, j) with i < j, send AT+PEER_RANGE=<j> over
     anchor i's serial port, collect `n_samples` measurements, average them.
  2. Build the symmetric N×N distance matrix D.
  3. Run classical MDS (mds_math.mds_2d) to recover (x, y) positions.
  4. Align the frame: anchor-0 at origin, anchor-1 on +X axis.
  5. Return positions + residual RMS + JSON config string.

Parameters (ROS2, see config/calibration_params.yaml)
──────────────────────────────────────────────────────
  anchor_serials   list of serial port paths, index = anchor ID
                   e.g. ["/dev/uwb-anchor0", "/dev/uwb-anchor1", ...]
  baud_rate        115200
  peer_range_timeout_s  2.0   (per AT+PEER_RANGE= attempt)
  default_n_samples     5
"""

from __future__ import annotations

import json
import re
import time
from typing import Optional

import numpy as np
import serial

import rclpy
from rclpy.node import Node

from saltybot_uwb_calibration_msgs.srv import CalibrateAnchors
from saltybot_uwb_calibration.mds_math import mds_2d, anchor_frame_align


_PEER_RANGE_RE = re.compile(
    r"\+PEER_RANGE:(\d+),(\d+),(\d+),([-\d.]+)"
)
_PEER_ERR_RE = re.compile(
    r"\+PEER_RANGE:ERR,(\d+),(\w+)"
)


class CalibrationNode(Node):

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

        self.declare_parameter("anchor_serials", [
            "/dev/uwb-anchor0",
            "/dev/uwb-anchor1",
            "/dev/uwb-anchor2",
            "/dev/uwb-anchor3",
        ])
        self.declare_parameter("baud_rate", 115200)
        self.declare_parameter("peer_range_timeout_s", 2.0)
        self.declare_parameter("default_n_samples", 5)

        self._serials: list[str] = (
            self.get_parameter("anchor_serials").value
        )
        self._baud: int = self.get_parameter("baud_rate").value
        self._timeout: float = self.get_parameter("peer_range_timeout_s").value
        self._default_n: int = self.get_parameter("default_n_samples").value

        self._srv = self.create_service(
            CalibrateAnchors,
            "/saltybot/uwb/calibrate_anchors",
            self._handle_calibrate,
        )

        self.get_logger().info(
            f"UWB calibration service ready — "
            f"{len(self._serials)} configured anchors"
        )

    # ── Service handler ────────────────────────────────────────────────────

    def _handle_calibrate(
        self,
        request: CalibrateAnchors.Request,
        response: CalibrateAnchors.Response,
    ) -> CalibrateAnchors.Response:

        # Determine anchor IDs to calibrate
        if request.anchor_ids:
            ids = list(request.anchor_ids)
        else:
            ids = list(range(len(self._serials)))

        n_anchors = len(ids)
        if n_anchors < 4:
            response.success = False
            response.message = (
                f"Need at least 4 anchors, got {n_anchors}"
            )
            return response

        n_samples = int(request.n_samples) if request.n_samples > 0 \
            else self._default_n

        self.get_logger().info(
            f"Starting calibration: anchors={ids}, n_samples={n_samples}"
        )

        # Open serial ports
        ports: dict[int, serial.Serial] = {}
        try:
            for anchor_id in ids:
                if anchor_id >= len(self._serials):
                    raise RuntimeError(
                        f"No serial port configured for anchor {anchor_id}"
                    )
                port_path = self._serials[anchor_id]
                ports[anchor_id] = serial.Serial(
                    port_path,
                    baudrate=self._baud,
                    timeout=self._timeout,
                )
                # Flush any stale data
                ports[anchor_id].reset_input_buffer()
                time.sleep(0.05)
                self.get_logger().info(
                    f"  Opened anchor {anchor_id} → {port_path}"
                )
        except Exception as exc:
            _close_all(ports)
            response.success = False
            response.message = f"Serial open failed: {exc}"
            return response

        # Measure all unique pairs
        D = np.zeros((n_anchors, n_anchors), dtype=float)
        id_to_idx = {aid: idx for idx, aid in enumerate(ids)}

        try:
            for i_idx, i_id in enumerate(ids):
                for j_idx, j_id in enumerate(ids):
                    if j_idx <= i_idx:
                        continue

                    dist_m = self._measure_pair(
                        ports[i_id], i_id, j_id, n_samples
                    )
                    if dist_m is None:
                        raise RuntimeError(
                            f"No valid range between anchor {i_id} "
                            f"and anchor {j_id}"
                        )
                    D[i_idx, j_idx] = dist_m
                    D[j_idx, i_idx] = dist_m
                    self.get_logger().info(
                        f"  [{i_id}↔{j_id}] = {dist_m:.3f} m"
                    )

        except Exception as exc:
            _close_all(ports)
            response.success = False
            response.message = str(exc)
            return response

        _close_all(ports)

        # MDS
        try:
            positions, residual = mds_2d(D)
            positions = anchor_frame_align(positions)
        except Exception as exc:
            response.success = False
            response.message = f"MDS failed: {exc}"
            return response

        # Build JSON config
        pos_dict = {}
        for idx, aid in enumerate(ids):
            pos_dict[str(aid)] = [
                round(float(positions[idx, 0]), 4),
                round(float(positions[idx, 1]), 4),
                0.0,
            ]
        json_str = json.dumps(pos_dict, indent=2)

        self.get_logger().info(
            f"Calibration complete — residual RMS {residual:.4f} m\n{json_str}"
        )

        # Fill response
        response.success = True
        response.message = (
            f"OK — residual RMS {residual*1000:.1f} mm"
        )
        response.anchor_ids = [int(a) for a in ids]
        response.positions_x = [float(positions[i, 0]) for i in range(n_anchors)]
        response.positions_y = [float(positions[i, 1]) for i in range(n_anchors)]
        response.positions_z = [0.0] * n_anchors
        response.residual_rms_m = float(residual)
        response.anchor_positions_json = json_str
        return response

    # ── Per-pair measurement ───────────────────────────────────────────────

    def _measure_pair(
        self,
        port: serial.Serial,
        my_id: int,
        peer_id: int,
        n_samples: int,
    ) -> Optional[float]:
        """
        Send AT+PEER_RANGE=<peer_id> `n_samples` times over `port`.
        Returns the mean range in metres, or None if all attempts fail.
        """
        measurements: list[float] = []

        for _ in range(n_samples):
            cmd = f"AT+PEER_RANGE={peer_id}\r\n"
            port.write(cmd.encode())
            port.flush()

            deadline = time.monotonic() + self._timeout
            while time.monotonic() < deadline:
                line = port.readline().decode(errors="replace").strip()
                if not line:
                    continue

                m = _PEER_RANGE_RE.match(line)
                if m:
                    # +PEER_RANGE:<my_id>,<peer_id>,<mm>,<rssi>
                    if int(m.group(1)) == my_id and int(m.group(2)) == peer_id:
                        mm = float(m.group(3))
                        measurements.append(mm / 1000.0)  # mm → m
                        break

                if _PEER_ERR_RE.match(line):
                    break  # timeout from firmware — try next sample

            time.sleep(0.05)  # inter-command gap

        if not measurements:
            return None
        return float(np.mean(measurements))


# ── Helpers ────────────────────────────────────────────────────────────────

def _close_all(ports: dict) -> None:
    for p in ports.values():
        try:
            p.close()
        except Exception:
            pass


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


if __name__ == "__main__":
    main()
