"""
orb_stereo_matcher.py — ORB feature detection and descriptor matching for stereo VO.

Provides two match types
────────────────────────
  Temporal  : left_gray(t-1) → left_gray(t)   — matched pairs for Essential-matrix estimation
  Stereo    : left_gray(t)   ↔ right_gray(t)  — disparity-based metric depth (scale source)

Matching
────────
  BFMatcher(NORM_HAMMING) with Lowe ratio test (default ratio=0.75).
  Stereo matching additionally enforces the epipolar row constraint (|Δrow| ≤ row_tol px)
  and requires positive disparity (left.x > right.x, i.e. object in front of cameras).

ORB parameters
──────────────
  nfeatures  : max keypoints per frame (default 500)
  scale_factor: image pyramid scale (default 1.2)
  nlevels    : pyramid levels (default 8)

Scale recovery (stereo)
───────────────────────
  For matched infra1/infra2 pairs with disparity d > 0:
      depth_i = baseline_m * fx / d_i
  Returns median over valid [d_min_px, d_max_px] disparity range.
  Returns 0.0 when fewer than min_stereo_samples good pairs exist.
"""

from __future__ import annotations

from typing import Tuple

import numpy as np
import cv2


# ── Matching parameters ────────────────────────────────────────────────────────
_LOWE_RATIO      = 0.75    # Lowe ratio test threshold
_MIN_MATCHES     = 20      # minimum temporal matches for VO
_ROW_TOL         = 2       # ±px for stereo epipolar row constraint
_DISP_MIN        = 1.0     # minimum valid disparity (pixels)
_DISP_MAX        = 192.0   # maximum valid disparity (D435i 640-px baseline, ~0.15m min depth)
_MIN_STEREO_SAMP = 10      # minimum stereo samples for reliable scale


class OrbStereoMatcher:
    """
    ORB-based feature tracker for monocular temporal matching and stereo scale.

    Usage (per frame)::

        matcher = OrbStereoMatcher()
        # --- temporal ---
        prev_pts, curr_pts = matcher.update(curr_gray_left)
        # --- stereo scale (optional, called same frame) ---
        depth_m = matcher.stereo_scale(curr_gray_left, curr_gray_right, fx, baseline_m)
    """

    def __init__(
        self,
        nfeatures:    int   = 500,
        scale_factor: float = 1.2,
        nlevels:      int   = 8,
    ) -> None:
        self._orb = cv2.ORB_create(
            nfeatures=nfeatures,
            scaleFactor=scale_factor,
            nLevels=nlevels,
            edgeThreshold=15,
            patchSize=31,
        )
        self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)

        # Previous-frame state
        self._prev_gray:  np.ndarray | None = None
        self._prev_kp:    list | None        = None
        self._prev_desc:  np.ndarray | None  = None

    # ── Temporal matching ──────────────────────────────────────────────────────

    def update(
        self,
        curr_gray: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Detect ORB features in curr_gray, match against previous frame.

        Returns
        -------
        (prev_pts, curr_pts) : (N, 2) float32 arrays of matched pixel coords.
        Returns empty (0, 2) arrays on first call or when matching fails.
        """
        kp, desc = self._orb.detectAndCompute(curr_gray, None)

        if (
            desc is None
            or len(kp) < 8
            or self._prev_desc is None
            or self._prev_kp is None
        ):
            # First frame — just store state
            self._prev_gray = curr_gray.copy()
            self._prev_kp   = kp
            self._prev_desc = desc
            return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32)

        # Ratio test matching
        prev_pts, curr_pts = self._ratio_match(
            self._prev_kp, self._prev_desc,
            kp, desc,
        )

        # Update state
        self._prev_gray = curr_gray.copy()
        self._prev_kp   = kp
        self._prev_desc = desc

        return prev_pts, curr_pts

    # ── Stereo scale ───────────────────────────────────────────────────────────

    def stereo_scale(
        self,
        left_gray:  np.ndarray,
        right_gray: np.ndarray,
        fx:         float,
        baseline_m: float,
    ) -> float:
        """
        Compute median depth (metres) from stereo ORB disparity.

        Matches ORB features between left and right images under the epipolar
        row constraint, then: depth_i = baseline_m * fx / disparity_i.

        Returns 0.0 on failure (< min_stereo_samples valid pairs).
        """
        if baseline_m <= 0 or fx <= 0:
            return 0.0

        kp_l, desc_l = self._orb.detectAndCompute(left_gray,  None)
        kp_r, desc_r = self._orb.detectAndCompute(right_gray, None)

        if desc_l is None or desc_r is None:
            return 0.0

        pts_l, pts_r = self._ratio_match(kp_l, desc_l, kp_r, desc_r)

        if len(pts_l) == 0:
            return 0.0

        # Epipolar row constraint: |row_l - row_r| ≤ _ROW_TOL
        row_diff = np.abs(pts_l[:, 1] - pts_r[:, 1])
        mask_row = row_diff <= _ROW_TOL

        # Positive disparity (left camera has larger x for points in front)
        disp = pts_l[:, 0] - pts_r[:, 0]
        mask_disp = (disp >= _DISP_MIN) & (disp <= _DISP_MAX)

        valid = mask_row & mask_disp
        if valid.sum() < _MIN_STEREO_SAMP:
            return 0.0

        depths = baseline_m * fx / disp[valid]
        return float(np.median(depths))

    # ── Reset ──────────────────────────────────────────────────────────────────

    def reset(self) -> None:
        self._prev_gray  = None
        self._prev_kp    = None
        self._prev_desc  = None

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _ratio_match(
        self,
        kp1:   list,
        desc1: np.ndarray,
        kp2:   list,
        desc2: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """BFMatcher kNN (k=2) + Lowe ratio test → (pts1, pts2) float32."""
        try:
            matches = self._matcher.knnMatch(desc1, desc2, k=2)
        except cv2.error:
            return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32)

        good = []
        for pair in matches:
            if len(pair) == 2:
                m, n = pair
                if m.distance < _LOWE_RATIO * n.distance:
                    good.append(m)

        if len(good) < _MIN_MATCHES:
            return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32)

        pts1 = np.array(
            [kp1[m.queryIdx].pt for m in good], dtype=np.float32
        ).reshape(-1, 2)
        pts2 = np.array(
            [kp2[m.trainIdx].pt for m in good], dtype=np.float32
        ).reshape(-1, 2)
        return pts1, pts2
