"""
optical_flow_tracker.py — Sparse optical flow feature tracker.

Primary path: OpenCV CUDA (SparsePyrLKOpticalFlow + FAST detector).
Fallback:      OpenCV CPU (calcOpticalFlowPyrLK + goodFeaturesToTrack).

Both paths output the same data:  matched pixel pairs (prev_pts, curr_pts)
suitable for Essential-matrix estimation.

GPU path targets 30 Hz on Jetson Orin Nano Super.  The FAST detector on the
Ampere GPU is ~4x faster than the CPU equivalent at the same quality level.

Design notes:
  - Features are re-detected whenever tracking drops below `min_features`.
  - Forward-backward consistency check discards unreliable tracks.
  - Points in the centre 80% of the image (avoids fisheye edge distortion).
  - Max `max_features` points per frame — higher = more robust, more cost.
"""

from __future__ import annotations
from typing import Tuple

import numpy as np
import cv2


_USE_CUDA = hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0

# LK optical flow parameters
_LK_WIN   = (21, 21)
_LK_ITERS = 30
_LK_EPS   = 0.01
_LK_LEVELS = 4

# FAST detector threshold
_FAST_THR = 20


class OpticalFlowTracker:
    """
    Maintains a set of tracked feature points across consecutive frames.

    Usage:
        tracker = OpticalFlowTracker()
        prev_pts, curr_pts = tracker.update(prev_gray, curr_gray)
        # prev_pts / curr_pts: (N, 2) float32 arrays of matched pixel coords.
    """

    def __init__(self, max_features: int = 300, min_features: int = 80):
        self._max_features = max_features
        self._min_features = min_features
        self._prev_pts: np.ndarray | None = None
        self._prev_gray: np.ndarray | None = None
        self._use_cuda = _USE_CUDA

        if self._use_cuda:
            self._lk = cv2.cuda.SparsePyrLKOpticalFlow_create(
                winSize=_LK_WIN,
                maxLevel=_LK_LEVELS,
                iters=_LK_ITERS,
                useInitialFlow=False,
            )
            self._fast = cv2.cuda.FastFeatureDetector_create(
                threshold=_FAST_THR, nonmaxSuppression=True
            )

    def update(
        self,
        prev_gray: np.ndarray,
        curr_gray: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Track features from prev_gray to curr_gray.

        Returns:
            (prev_pts, curr_pts) — matched point pairs, shape (N, 2) float32.
            Returns empty arrays if tracking fails or there is no previous frame.
        """
        h, w = curr_gray.shape

        # Re-detect if we have too few features or no previous frame
        if self._prev_pts is None or len(self._prev_pts) < self._min_features:
            self._prev_pts = self._detect(prev_gray, h, w)
            self._prev_gray = prev_gray
            if self._prev_pts is None or len(self._prev_pts) < 8:
                return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32)

        # Track features
        try:
            if self._use_cuda:
                prev_pts_c, curr_pts, good = self._track_cuda(
                    self._prev_gray, curr_gray, self._prev_pts
                )
            else:
                prev_pts_c, curr_pts, good = self._track_cpu(
                    self._prev_gray, curr_gray, self._prev_pts
                )
        except Exception:
            self._prev_pts = None
            return np.zeros((0, 2), np.float32), np.zeros((0, 2), np.float32)

        prev_pts_c = prev_pts_c[good]
        curr_pts   = curr_pts[good]

        # Update state
        self._prev_pts  = curr_pts
        self._prev_gray = curr_gray

        return prev_pts_c.reshape(-1, 2), curr_pts.reshape(-1, 2)

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

    # ── Detection ─────────────────────────────────────────────────────────────

    def _detect(self, gray: np.ndarray, h: int, w: int) -> np.ndarray | None:
        """Detect FAST features in the inner 80% of the image."""
        margin_x = int(w * 0.10)
        margin_y = int(h * 0.10)
        roi = gray[margin_y:h - margin_y, margin_x:w - margin_x]

        if self._use_cuda:
            pts = self._detect_cuda(roi)
        else:
            pts = self._detect_cpu(roi)

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

        # Offset back to full-image coordinates
        pts[:, 0] += margin_x
        pts[:, 1] += margin_y

        # Subsample if too many
        if len(pts) > self._max_features:
            idx = np.random.choice(len(pts), self._max_features, replace=False)
            pts = pts[idx]

        return pts.astype(np.float32)

    def _detect_cuda(self, roi: np.ndarray) -> np.ndarray | None:
        gpu_img = cv2.cuda_GpuMat()
        gpu_img.upload(roi)
        kp = self._fast.detect(gpu_img, None)
        if not kp:
            return None
        pts = cv2.KeyPoint_convert(kp)
        return pts.reshape(-1, 2)

    def _detect_cpu(self, roi: np.ndarray) -> np.ndarray | None:
        pts = cv2.goodFeaturesToTrack(
            roi,
            maxCorners=self._max_features,
            qualityLevel=0.01,
            minDistance=10,
            blockSize=7,
        )
        if pts is None:
            return None
        return pts.reshape(-1, 2)

    # ── Tracking ──────────────────────────────────────────────────────────────

    def _track_cuda(
        self, prev: np.ndarray, curr: np.ndarray, pts: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        pts_gpu   = cv2.cuda_GpuMat(pts.reshape(-1, 1, 2))
        prev_gpu  = cv2.cuda_GpuMat(prev)
        curr_gpu  = cv2.cuda_GpuMat(curr)
        curr_pts_gpu, status_gpu, _ = self._lk.calc(prev_gpu, curr_gpu, pts_gpu, None)
        curr_pts = curr_pts_gpu.download().reshape(-1, 2)
        status   = status_gpu.download().ravel()

        # Forward-backward check
        bw_pts_gpu, bw_status_gpu, _ = self._lk.calc(curr_gpu, prev_gpu, curr_pts_gpu, None)
        bw_pts   = bw_pts_gpu.download().reshape(-1, 2)
        bw_status = bw_status_gpu.download().ravel()

        fb_err = np.linalg.norm(pts.reshape(-1, 2) - bw_pts, axis=1)
        good   = (status > 0) & (bw_status > 0) & (fb_err < 1.0)
        return pts, curr_pts, good

    def _track_cpu(
        self, prev: np.ndarray, curr: np.ndarray, pts: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        lk_params = dict(
            winSize=_LK_WIN,
            maxLevel=_LK_LEVELS,
            criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
                      _LK_ITERS, _LK_EPS),
        )
        curr_pts, st_fwd, _ = cv2.calcOpticalFlowPyrLK(
            prev, curr, pts.reshape(-1, 1, 2), None, **lk_params
        )
        bw_pts, st_bwd, _   = cv2.calcOpticalFlowPyrLK(
            curr, prev, curr_pts, None, **lk_params
        )
        fb_err = np.linalg.norm(
            pts.reshape(-1, 2) - bw_pts.reshape(-1, 2), axis=1
        )
        good = (
            (st_fwd.ravel() > 0) &
            (st_bwd.ravel() > 0) &
            (fb_err < 1.0)
        )
        return pts.reshape(-1, 2), curr_pts.reshape(-1, 2), good
