"""`RansacFilter` — shared 2D-2D RANSAC + median-residual helper (AZ-282 / E-CC-HELPERS).

Implements the `ransac_filter` contract v1.0.0
(`_docs/02_document/contracts/shared_helpers/ransac_filter.md`).

Stateless static-only design — `coderule.mdc` permits static methods for
pure self-contained computations. Determinism is guaranteed by setting
`cv2.setRNGSeed(0)` immediately before every `cv2.findHomography(...,
RANSAC)` call.

Public surface raises ONLY `RansacFilterError`; OpenCV's lower-level
exceptions are wrapped.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Final

import cv2
import numpy as np

from gps_denied_onboard.helpers.se3_utils import SE3, se3_to_matrix

__all__ = [
    "RansacFilter",
    "RansacFilterError",
    "RansacResult",
]


# RANSAC requires ≥4 points to fit a homography (4 pairs of (x,y) ↔ (x,y)).
_HOMOGRAPHY_MIN_POINTS: Final[int] = 4
_DETERMINISTIC_SEED: Final[int] = 0


class RansacFilterError(ValueError):
    """Raised when an input violates the public shape / dtype / threshold contract."""


@dataclass(frozen=True)
class RansacResult:
    """Frozen output of `RansacFilter.filter_correspondences`.

    The numpy arrays are not copied; consumers MUST treat them as
    read-only. `median_residual_px` is NaN when the inlier set is empty
    (matches `compute_reprojection_residual` semantics).
    """

    inlier_correspondences: np.ndarray
    inlier_count: int
    outlier_count: int
    median_residual_px: float


def _validate_correspondences(correspondences: np.ndarray, *, where: str) -> None:
    if not isinstance(correspondences, np.ndarray):
        raise RansacFilterError(
            f"{where}: expected np.ndarray; got {type(correspondences).__name__}"
        )
    if correspondences.ndim != 2 or correspondences.shape[1] != 4:
        raise RansacFilterError(
            f"{where}: correspondences must have shape (N, 4) [x_a, y_a, x_b, y_b]; "
            f"got {correspondences.shape}"
        )


def _validate_threshold(ransac_threshold_px: float) -> None:
    if not isinstance(ransac_threshold_px, (int, float)) or ransac_threshold_px <= 0:
        raise RansacFilterError(
            f"ransac_threshold_px must be a positive float; got {ransac_threshold_px!r}"
        )


def _validate_min_inliers(min_inliers: int) -> None:
    if not isinstance(min_inliers, int) or min_inliers < 0:
        raise RansacFilterError(f"min_inliers must be a non-negative int; got {min_inliers!r}")


def _median_residual(
    inliers: np.ndarray, K: np.ndarray, distortion: np.ndarray, pose_matrix: np.ndarray
) -> float:
    """Median pixel residual between (x_b, y_b) and reprojected (x_a, y_a)+pose.

    Treats each correspondence's image-a pixel as the ``observed`` point and
    its image-b pixel as the ``predicted`` point under the supplied pose.
    The 2D-to-3D back-projection assumes z=1 in camera-a frame — sufficient
    for the helper's purpose of giving consumers a deterministic, OpenCV-
    backed quality signal in pixels. The contract pins MEDIAN (NOT mean).
    """
    if inliers.shape[0] == 0:
        return float("nan")

    pts_a = inliers[:, :2].astype(np.float64, copy=False)
    pts_b = inliers[:, 2:].astype(np.float64, copy=False)
    # Back-project image-a pixels into 3D (z=1) in camera-a frame using K^{-1}.
    K_inv = np.linalg.inv(K)
    pixels_h = np.hstack([pts_a, np.ones((pts_a.shape[0], 1), dtype=np.float64)])
    rays = (K_inv @ pixels_h.T).T  # (N, 3)

    # Apply pose (R | t) from cam_a -> cam_b: p_b = R @ p_a + t.
    R = pose_matrix[:3, :3]
    t = pose_matrix[:3, 3]
    rvec, _ = cv2.Rodrigues(R)
    projected, _ = cv2.projectPoints(
        rays.reshape(-1, 1, 3), rvec=rvec, tvec=t, cameraMatrix=K, distCoeffs=distortion
    )
    projected_pts = projected.reshape(-1, 2)
    residuals = np.linalg.norm(projected_pts - pts_b, axis=1)
    return float(np.median(residuals))


class RansacFilter:
    """Shared 2D-2D RANSAC inlier filter + reprojection-residual helper.

    All methods are static; no module-level state. Calls into OpenCV pin
    the RANSAC seed for byte-equal determinism (AC-3).
    """

    @staticmethod
    def filter_correspondences(
        correspondences: np.ndarray,
        ransac_threshold_px: float,
        min_inliers: int,
    ) -> RansacResult:
        """Run `cv2.findHomography(..., RANSAC)` and return the inlier mask.

        ``min_inliers`` is informational only — see the contract's
        "Min-inliers semantics" invariant.
        """
        _validate_correspondences(correspondences, where="filter_correspondences")
        _validate_threshold(ransac_threshold_px)
        _validate_min_inliers(min_inliers)

        n_points = correspondences.shape[0]
        if n_points < _HOMOGRAPHY_MIN_POINTS:
            raise RansacFilterError(
                f"filter_correspondences: homography RANSAC requires ≥{_HOMOGRAPHY_MIN_POINTS} "
                f"correspondences; got {n_points}"
            )

        pts_a = correspondences[:, :2].astype(np.float64, copy=False)
        pts_b = correspondences[:, 2:].astype(np.float64, copy=False)

        cv2.setRNGSeed(_DETERMINISTIC_SEED)
        try:
            _H, mask = cv2.findHomography(
                pts_a,
                pts_b,
                method=cv2.RANSAC,
                ransacReprojThreshold=float(ransac_threshold_px),
            )
        except cv2.error as exc:
            raise RansacFilterError(f"filter_correspondences: OpenCV RANSAC failed: {exc}") from exc

        if mask is None:
            inlier_mask = np.zeros(n_points, dtype=bool)
        else:
            inlier_mask = mask.ravel().astype(bool)

        inliers = correspondences[inlier_mask]
        inlier_count = int(inliers.shape[0])
        outlier_count = n_points - inlier_count

        if inlier_count == 0:
            median_residual = float("nan")
        else:
            # Median residual from the homography fit itself — distance from
            # H @ pts_a to pts_b for the inlier subset. Reuse cv2.perspectiveTransform
            # to stay in OpenCV's reference frame; this matches the C3.5/C4 contract.
            warped = cv2.perspectiveTransform(
                inliers[:, :2].reshape(-1, 1, 2).astype(np.float64), _H
            ).reshape(-1, 2)
            residuals = np.linalg.norm(warped - inliers[:, 2:].astype(np.float64), axis=1)
            median_residual = float(np.median(residuals))

        return RansacResult(
            inlier_correspondences=inliers,
            inlier_count=inlier_count,
            outlier_count=outlier_count,
            median_residual_px=median_residual,
        )

    @staticmethod
    def compute_reprojection_residual(
        correspondences: np.ndarray,
        K: np.ndarray,
        distortion: np.ndarray,
        pose: SE3,
    ) -> float:
        """Median reprojection residual in pixels for the supplied inlier set.

        Empty inlier sets return ``NaN`` per AC-5. ``K`` MUST be (3, 3);
        ``distortion`` MUST be (5,) or (8,) — OpenCV's standard models.
        """
        _validate_correspondences(correspondences, where="compute_reprojection_residual")

        if not isinstance(K, np.ndarray):
            raise RansacFilterError(
                f"compute_reprojection_residual: K must be np.ndarray; got {type(K).__name__}"
            )
        if K.shape != (3, 3):
            raise RansacFilterError(
                f"compute_reprojection_residual: K must have shape (3, 3); got {K.shape}"
            )

        if not isinstance(distortion, np.ndarray):
            raise RansacFilterError(
                f"compute_reprojection_residual: distortion must be np.ndarray; "
                f"got {type(distortion).__name__}"
            )
        if distortion.ndim != 1 or distortion.shape[0] not in (5, 8):
            raise RansacFilterError(
                f"compute_reprojection_residual: distortion must have shape (5,) or (8,); "
                f"got {distortion.shape}"
            )

        K_f64 = K.astype(np.float64, copy=False)
        dist_f64 = distortion.astype(np.float64, copy=False)
        pose_matrix = se3_to_matrix(pose)

        try:
            return _median_residual(correspondences, K_f64, dist_f64, pose_matrix)
        except cv2.error as exc:
            raise RansacFilterError(
                f"compute_reprojection_residual: OpenCV projection failed: {exc}"
            ) from exc
