"""
dock_detector.py — ArUco marker + IR beacon detection for docking (Issue #158).

Detection sources
─────────────────
  ArUco marker  : cv2.aruco marker on the dock face; PnP solve gives full 3-D
                  pose relative to the robot camera.  Marker ID and physical
                  size are configurable.

  IR beacon     : single ADC channel (normalised 0–1) with envelope follower
                  and amplitude threshold.  Provides presence detection only
                  (no pose), used as a fallback or range-alert.

Outputs
───────
  DockPose(distance_m, yaw_rad, lateral_m, source)
    distance_m  — straight-line distance to dock marker centre (m)
    yaw_rad     — bearing error (rad); +ve = dock is right of camera axis (tx > 0)
    lateral_m   — lateral offset (m) in camera frame; +ve = dock is right
    source      — "aruco" | "ir"

  ArucoDetector.detect() returns Optional[DockPose]; None if marker absent or
  OpenCV / cv2.aruco are unavailable at runtime.

  IRBeaconDetector.update(sample) returns bool (True = beacon detected).

Pure module — no ROS2 dependency.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional


# ── DockPose ─────────────────────────────────────────────────────────────────

@dataclass
class DockPose:
    """Dock pose in the robot's camera frame."""
    distance_m: float   # straight-line distance to marker (m)
    yaw_rad:    float   # bearing error (rad); +ve = dock is right of camera axis (tx > 0)
    lateral_m:  float   # lateral offset (m) in camera X; +ve = dock is right
    source:     str     # "aruco" | "ir"


# ── ArucoDetector ─────────────────────────────────────────────────────────────

class ArucoDetector:
    """
    Detect a single ArUco marker and solve its 3-D pose with PnP.

    Parameters
    ----------
    marker_id       : ArUco marker ID placed on the dock face (default 42)
    marker_size_m   : physical side length of the printed marker (m)
    aruco_dict_id   : cv2.aruco dictionary constant (default DICT_4X4_50 = 0)
    """

    MARKER_ID_DEFAULT  = 42
    MARKER_SIZE_DEFAULT = 0.10   # 10 cm

    def __init__(
        self,
        marker_id:     int   = MARKER_ID_DEFAULT,
        marker_size_m: float = MARKER_SIZE_DEFAULT,
        aruco_dict_id: int   = 0,    # cv2.aruco.DICT_4X4_50
    ):
        self._marker_id   = marker_id
        self._marker_size = marker_size_m
        self._dict_id     = aruco_dict_id
        self._ready: Optional[bool] = None   # None = not yet attempted

        # Set by _init(); None when OpenCV unavailable
        self._cv2          = None
        self._detector     = None
        self._obj_pts      = None

    # ── Initialisation ────────────────────────────────────────────────────────

    def _init(self) -> bool:
        """Lazy-initialise cv2.aruco.  Returns False if unavailable."""
        if self._ready is not None:
            return self._ready
        try:
            import cv2
            import numpy as np
            aruco_dict   = cv2.aruco.getPredefinedDictionary(self._dict_id)
            aruco_params = cv2.aruco.DetectorParameters()
            self._detector = cv2.aruco.ArucoDetector(aruco_dict, aruco_params)
            half = self._marker_size / 2.0
            self._obj_pts = np.array([
                [-half,  half, 0.0],
                [ half,  half, 0.0],
                [ half, -half, 0.0],
                [-half, -half, 0.0],
            ], dtype=np.float64)
            self._cv2   = cv2
            self._ready = True
        except Exception:
            self._ready = False
        return self._ready

    # ── Public API ────────────────────────────────────────────────────────────

    @property
    def available(self) -> bool:
        """True if cv2.aruco is importable and initialised."""
        return self._init()

    def detect(
        self,
        image,           # numpy.ndarray  H×W×3 BGR (or greyscale)
        camera_matrix,   # numpy.ndarray  3×3
        dist_coeffs,     # numpy.ndarray  1×4 or 1×5
    ) -> Optional[DockPose]:
        """
        Detect the dock marker in *image* and return its pose.

        Returns None if the target marker is absent, PnP fails, or OpenCV is
        not available.
        """
        if not self._init():
            return None

        try:
            corners, ids, _ = self._detector.detectMarkers(image)
        except Exception:
            return None

        if ids is None or len(ids) == 0:
            return None

        import numpy as np
        for i, mid in enumerate(ids.flatten()):
            if int(mid) != self._marker_id:
                continue
            img_pts = corners[i].reshape(4, 2).astype(np.float64)
            ok, rvec, tvec = self._cv2.solvePnP(
                self._obj_pts, img_pts,
                camera_matrix, dist_coeffs,
                flags=self._cv2.SOLVEPNP_IPPE_SQUARE,
            )
            if not ok:
                continue
            tx = float(tvec[0])
            ty = float(tvec[1])
            tz = float(tvec[2])
            distance = math.sqrt(tx * tx + ty * ty + tz * tz)
            yaw      = math.atan2(tx, tz)   # camera-x deviation from optical axis
            return DockPose(
                distance_m=distance,
                yaw_rad=yaw,
                lateral_m=tx,
                source="aruco",
            )
        return None


# ── IRBeaconDetector ──────────────────────────────────────────────────────────

class IRBeaconDetector:
    """
    IR beacon presence detector from a single normalised ADC channel.

    Maintains an exponential envelope follower.  When the smoothed amplitude
    exceeds *threshold* the beacon is considered detected.

    Parameters
    ----------
    threshold       : amplitude above which beacon is 'detected' (default 0.5)
    smoothing_alpha : EMA coefficient ∈ (0, 1]; 1 = no smoothing (default 0.3)
    """

    def __init__(
        self,
        threshold:       float = 0.5,
        smoothing_alpha: float = 0.3,
    ):
        self._threshold = max(1e-9, float(threshold))
        self._alpha     = max(1e-9, min(1.0, smoothing_alpha))
        self._envelope  = 0.0

    # ── Properties ────────────────────────────────────────────────────────────

    @property
    def envelope(self) -> float:
        """Current EMA envelope value."""
        return self._envelope

    @property
    def detected(self) -> bool:
        """True when smoothed amplitude ≥ threshold."""
        return self._envelope >= self._threshold

    # ── Update ────────────────────────────────────────────────────────────────

    def update(self, sample: float) -> bool:
        """
        Submit one ADC sample (normalised 0–1 or raw voltage).
        Returns True if the beacon is currently detected.
        """
        amp = abs(sample)
        self._envelope = (1.0 - self._alpha) * self._envelope + self._alpha * amp
        return self.detected

    def reset(self) -> None:
        """Clear envelope to zero (e.g. when node restarts)."""
        self._envelope = 0.0
