"""
logger_node.py — UWB position logger and accuracy analyzer (Issue #634)

Subscriptions
─────────────
  /saltybot/pose/fused   geometry_msgs/PoseStamped   — EKF-fused pose (200 Hz)
  /saltybot/uwb/pose     geometry_msgs/PoseStamped   — raw UWB pose (10 Hz)
  /uwb/ranges            saltybot_uwb_msgs/UwbRangeArray — per-anchor ranges

Service
───────
  /saltybot/uwb/start_accuracy_test  StartAccuracyTest
    Collects N fused-pose samples at known position, computes accuracy
    metrics, publishes AccuracyReport, writes JSON summary.

Publishes
─────────
  /saltybot/uwb/accuracy_report  AccuracyReport

CSV logging (continuous, to log_dir)
─────────────────────────────────────
  fused_pose_<DATE>.csv   — timestamp_ns, x_m, y_m, heading_rad
  uwb_pose_<DATE>.csv     — timestamp_ns, x_m, y_m
  uwb_ranges_<DATE>.csv   — timestamp_ns, anchor_id, range_m, raw_mm, rssi, tag_id
  accuracy_<test_id>.json — accuracy test results

Parameters
──────────
  log_dir              ~/uwb_logs
  enable_csv           true
  default_n_samples    200
  default_timeout_s    30.0
  flush_interval_s     5.0
"""

from __future__ import annotations

import csv
import json
import math
import os
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import numpy as np

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

from geometry_msgs.msg import PoseStamped
from saltybot_uwb_msgs.msg import UwbRangeArray
from saltybot_uwb_logger_msgs.msg import AccuracyReport
from saltybot_uwb_logger_msgs.srv import StartAccuracyTest

from saltybot_uwb_logger.accuracy_stats import compute_stats, RangeAccum


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=10,
)


class UwbLoggerNode(Node):

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

        # ── Parameters ────────────────────────────────────────────────────
        self.declare_parameter("log_dir",            os.path.expanduser("~/uwb_logs"))
        self.declare_parameter("enable_csv",         True)
        self.declare_parameter("default_n_samples",  200)
        self.declare_parameter("default_timeout_s",  30.0)
        self.declare_parameter("flush_interval_s",   5.0)

        self._log_dir     = Path(self.get_parameter("log_dir").value)
        self._enable_csv  = self.get_parameter("enable_csv").value
        self._default_n   = self.get_parameter("default_n_samples").value
        self._default_to  = self.get_parameter("default_timeout_s").value
        self._flush_iv    = self.get_parameter("flush_interval_s").value

        if self._enable_csv:
            self._log_dir.mkdir(parents=True, exist_ok=True)

        # ── CSV writers ───────────────────────────────────────────────────
        date_tag = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        self._csv_lock   = threading.Lock()
        self._csv_files  = []
        self._fused_csv  = self._open_csv(
            f"fused_pose_{date_tag}.csv",
            ["timestamp_ns", "x_m", "y_m", "heading_rad"])
        self._uwb_csv    = self._open_csv(
            f"uwb_pose_{date_tag}.csv",
            ["timestamp_ns", "x_m", "y_m"])
        self._ranges_csv = self._open_csv(
            f"uwb_ranges_{date_tag}.csv",
            ["timestamp_ns", "anchor_id", "range_m", "raw_mm", "rssi", "tag_id"])

        # ── Accuracy test state ────────────────────────────────────────────
        self._test_lock       = threading.Lock()
        self._test_active     = False
        self._test_id         = ""
        self._test_truth_x    = 0.0
        self._test_truth_y    = 0.0
        self._test_n_target   = 0
        self._test_xs: list   = []
        self._test_ys: list   = []
        self._test_ranges: dict[int, RangeAccum] = {}
        self._test_t0         = 0.0
        self._test_timeout    = 0.0
        self._test_done_event = threading.Event()

        # ── Publisher ─────────────────────────────────────────────────────
        self._report_pub = self.create_publisher(
            AccuracyReport, "/saltybot/uwb/accuracy_report", 10
        )

        # ── Subscriptions ─────────────────────────────────────────────────
        self.create_subscription(
            PoseStamped, "/saltybot/pose/fused",
            self._fused_cb, _SENSOR_QOS)
        self.create_subscription(
            PoseStamped, "/saltybot/uwb/pose",
            self._uwb_pose_cb, _SENSOR_QOS)
        self.create_subscription(
            UwbRangeArray, "/uwb/ranges",
            self._ranges_cb, _SENSOR_QOS)

        # ── Service ───────────────────────────────────────────────────────
        self._srv = self.create_service(
            StartAccuracyTest,
            "/saltybot/uwb/start_accuracy_test",
            self._handle_start_test)

        # ── Periodic flush ────────────────────────────────────────────────
        self.create_timer(self._flush_iv, self._flush_csv)

        self.get_logger().info(
            f"UWB logger ready — log_dir={self._log_dir} csv={self._enable_csv}"
        )

    def destroy_node(self) -> None:
        self._flush_csv()
        for fh in self._csv_files:
            try:
                fh.close()
            except Exception:
                pass
        super().destroy_node()

    # ── Callbacks ─────────────────────────────────────────────────────────

    def _fused_cb(self, msg: PoseStamped) -> None:
        ts = _stamp_ns(msg.header.stamp)
        x  = msg.pose.position.x
        y  = msg.pose.position.y
        q  = msg.pose.orientation
        heading = 2.0 * math.atan2(float(q.z), float(q.w))

        if self._enable_csv:
            with self._csv_lock:
                if self._fused_csv:
                    self._fused_csv.writerow([ts, x, y, heading])

        with self._test_lock:
            if self._test_active:
                self._test_xs.append(x)
                self._test_ys.append(y)
                if len(self._test_xs) >= self._test_n_target:
                    self._test_active = False
                    self._test_done_event.set()

    def _uwb_pose_cb(self, msg: PoseStamped) -> None:
        ts = _stamp_ns(msg.header.stamp)
        x  = msg.pose.position.x
        y  = msg.pose.position.y
        if self._enable_csv:
            with self._csv_lock:
                if self._uwb_csv:
                    self._uwb_csv.writerow([ts, x, y])

    def _ranges_cb(self, msg: UwbRangeArray) -> None:
        ts = _stamp_ns(msg.header.stamp)
        for r in msg.ranges:
            if self._enable_csv:
                with self._csv_lock:
                    if self._ranges_csv:
                        self._ranges_csv.writerow([
                            ts, r.anchor_id, r.range_m,
                            r.raw_mm, r.rssi, r.tag_id,
                        ])
            with self._test_lock:
                if self._test_active:
                    aid = int(r.anchor_id)
                    if aid not in self._test_ranges:
                        self._test_ranges[aid] = RangeAccum(anchor_id=aid)
                    self._test_ranges[aid].add(float(r.range_m))

    # ── Service ────────────────────────────────────────────────────────────

    def _handle_start_test(
        self,
        request: StartAccuracyTest.Request,
        response: StartAccuracyTest.Response,
    ) -> StartAccuracyTest.Response:

        with self._test_lock:
            if self._test_active:
                response.success = False
                response.message = "Another test is already running"
                response.test_id = self._test_id
                return response

        n = int(request.n_samples) if request.n_samples > 0 else self._default_n
        timeout = float(request.timeout_s) if request.timeout_s > 0.0 else self._default_to
        test_id = (
            request.test_id.strip()
            if request.test_id.strip()
            else datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        )

        self.get_logger().info(
            f"Accuracy test '{test_id}' — "
            f"truth=({request.truth_x_m:.3f}, {request.truth_y_m:.3f}) "
            f"n={n} timeout={timeout}s"
        )

        with self._test_lock:
            self._test_id       = test_id
            self._test_truth_x  = float(request.truth_x_m)
            self._test_truth_y  = float(request.truth_y_m)
            self._test_n_target = n
            self._test_xs       = []
            self._test_ys       = []
            self._test_ranges   = {}
            self._test_t0       = time.monotonic()
            self._test_timeout  = timeout
            self._test_done_event.clear()
            self._test_active   = True

        threading.Thread(
            target=self._run_test,
            args=(test_id, request.truth_x_m, request.truth_y_m, n, timeout),
            daemon=True,
        ).start()

        response.success = True
        response.message = f"Test '{test_id}' started, collecting {n} samples"
        response.test_id = test_id
        return response

    def _run_test(
        self,
        test_id: str,
        truth_x: float,
        truth_y: float,
        n_target: int,
        timeout: float,
    ) -> None:
        finished = self._test_done_event.wait(timeout=timeout)

        with self._test_lock:
            self._test_active = False
            xs = list(self._test_xs)
            ys = list(self._test_ys)
            ranges_snap = dict(self._test_ranges)
            t0 = self._test_t0

        duration = time.monotonic() - t0
        n_got = len(xs)

        if n_got < 2:
            self.get_logger().error(
                f"Test '{test_id}' failed: only {n_got} samples in {duration:.1f}s"
            )
            return

        if not finished:
            self.get_logger().warn(
                f"Test '{test_id}' timed out — {n_got}/{n_target} samples"
            )

        try:
            stats = compute_stats(np.array(xs), np.array(ys), truth_x, truth_y)
        except ValueError as exc:
            self.get_logger().error(f"Stats failed: {exc}")
            return

        self.get_logger().info(
            f"Test '{test_id}' — n={stats.n_samples} "
            f"CEP50={stats.cep50*100:.1f}cm CEP95={stats.cep95*100:.1f}cm "
            f"RMSE={stats.rmse*100:.1f}cm max={stats.max_error*100:.1f}cm"
        )

        # Publish report
        report = AccuracyReport()
        report.header.stamp    = self.get_clock().now().to_msg()
        report.header.frame_id = "map"
        report.truth_x_m   = truth_x
        report.truth_y_m   = truth_y
        report.n_samples   = stats.n_samples
        report.mean_x_m    = stats.mean_x
        report.mean_y_m    = stats.mean_y
        report.std_x_m     = stats.std_x
        report.std_y_m     = stats.std_y
        report.std_2d_m    = stats.std_2d
        report.bias_x_m    = stats.bias_x
        report.bias_y_m    = stats.bias_y
        report.cep50_m     = stats.cep50
        report.cep95_m     = stats.cep95
        report.max_error_m = stats.max_error
        report.rmse_m      = stats.rmse
        report.test_id     = test_id
        report.duration_s  = duration

        anchor_ids = sorted(ranges_snap.keys())
        report.anchor_ids          = [int(a) for a in anchor_ids]
        report.anchor_range_mean_m = [ranges_snap[a].mean for a in anchor_ids]
        report.anchor_range_std_m  = [ranges_snap[a].std  for a in anchor_ids]

        self._report_pub.publish(report)

        # Write JSON
        if self._enable_csv:
            self._write_json(test_id, truth_x, truth_y, stats,
                             anchor_ids, ranges_snap, duration)

    # ── CSV / JSON helpers ─────────────────────────────────────────────────

    def _open_csv(
        self, filename: str, headers: list
    ) -> Optional[csv.writer]:
        if not self._enable_csv:
            return None
        path = self._log_dir / filename
        fh = open(path, "w", newline="", buffering=1)
        writer = csv.writer(fh)
        writer.writerow(headers)
        self._csv_files.append(fh)
        return writer

    def _flush_csv(self) -> None:
        for fh in self._csv_files:
            try:
                fh.flush()
            except Exception:
                pass

    def _write_json(self, test_id, truth_x, truth_y, stats,
                    anchor_ids, ranges_snap, duration) -> None:
        path = self._log_dir / f"accuracy_{test_id}.json"
        data = {
            "test_id":     test_id,
            "truth":       {"x_m": truth_x, "y_m": truth_y},
            "n_samples":   stats.n_samples,
            "duration_s":  round(duration, 3),
            "mean":        {"x_m": round(stats.mean_x, 4),
                            "y_m": round(stats.mean_y, 4)},
            "bias":        {"x_m": round(stats.bias_x, 4),
                            "y_m": round(stats.bias_y, 4)},
            "std":         {"x_m": round(stats.std_x, 4),
                            "y_m": round(stats.std_y, 4),
                            "2d_m": round(stats.std_2d, 4)},
            "cep50_m":     round(stats.cep50, 4),
            "cep95_m":     round(stats.cep95, 4),
            "rmse_m":      round(stats.rmse, 4),
            "max_error_m": round(stats.max_error, 4),
            "anchors": {
                str(aid): {
                    "range_mean_m": round(ranges_snap[aid].mean, 4),
                    "range_std_m":  round(ranges_snap[aid].std,  4),
                    "n":            ranges_snap[aid].count,
                }
                for aid in anchor_ids
            },
        }
        try:
            with open(path, "w") as f:
                json.dump(data, f, indent=2)
            self.get_logger().info(f"Report written: {path}")
        except OSError as exc:
            self.get_logger().error(f"JSON write failed: {exc}")


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

def _stamp_ns(stamp) -> int:
    return stamp.sec * 1_000_000_000 + stamp.nanosec


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


if __name__ == "__main__":
    main()
