"""
stereo_vo.py — Stereo visual odometry using Essential matrix decomposition.

Input per frame:
  - Tracked feature pairs (prev_pts, curr_pts) from OpticalFlowTracker
  - Depth image (float32, metres) aligned to the IR1/left camera
  - Camera intrinsic matrix K

Pipeline:
  1. Find Essential matrix E from point correspondences (5-point RANSAC).
  2. Decompose E → [R|t] (unit-vector translation).
  3. Recover metric scale from depth at tracked feature points.
  4. Compose the incremental SE(3) pose and integrate into global estimate.

Scale recovery:
  Median depth at both prev_pts and curr_pts yields consistent scale.
  Only depth samples in [0.3, 6.0m] range are used (D435i reliable range).

Output:
  StereoVO.update() returns a PoseEstimate dataclass with:
    x, y, z       — position (m)
    roll, pitch, yaw — orientation (rad)
    vx, vz, omega — linear + angular velocity (m/s, rad/s) for current step
    valid          — False if insufficient inliers or depth samples

Coordinate frame:
  x = forward, y = left, z = up  (ROS REP-105 base_link convention).
  The camera looks forward along its +Z optical axis; we rotate into base_link.
"""

from __future__ import annotations
import math
from dataclasses import dataclass, field
from typing import Tuple

import numpy as np
import cv2


@dataclass
class PoseEstimate:
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0
    roll:  float = 0.0
    pitch: float = 0.0
    yaw:   float = 0.0
    vx:    float = 0.0
    vy:    float = 0.0
    omega: float = 0.0
    valid: bool  = False


# Depth validity window
_D_MIN = 0.3
_D_MAX = 6.0

# Essential matrix RANSAC
_RANSAC_THRESHOLD = 1.0    # pixels
_RANSAC_PROB      = 0.999
_MIN_INLIERS      = 20

# Minimum depth samples for reliable scale
_MIN_DEPTH_SAMPLES = 10

# Camera→base_link rotation (camera looks forward along +Z, base_link uses +X)
# Rotation: camera_Z → base_X, camera_-Y → base_Z
_R_CAM_TO_BASE = np.array([
    [ 0,  0,  1],
    [-1,  0,  0],
    [ 0, -1,  0],
], dtype=np.float64)


def _median_depth(depth: np.ndarray, pts: np.ndarray) -> float:
    """Median depth at pixel locations pts (Nx2 float32).  Returns 0 on failure."""
    h, w = depth.shape
    u = np.clip(pts[:, 0].astype(int), 0, w - 1)
    v = np.clip(pts[:, 1].astype(int), 0, h - 1)
    samples = depth[v, u]
    valid = samples[(samples > _D_MIN) & (samples < _D_MAX)]
    return float(np.median(valid)) if len(valid) >= _MIN_DEPTH_SAMPLES else 0.0


def _rot_to_yaw(R: np.ndarray) -> float:
    """Extract yaw from rotation matrix (Z-axis rotation in XY plane)."""
    return math.atan2(float(R[1, 0]), float(R[0, 0]))


class StereoVO:
    """
    Incremental visual odometry with depth-aided scale recovery.

    Call update() once per frame pair.  Thread-safe for single caller.
    """

    def __init__(self, K: np.ndarray | None = None, dt: float = 1.0 / 30.0):
        """
        Args:
            K:   3×3 camera intrinsic matrix.  Can be set later via set_K().
            dt:  nominal frame interval (seconds) for velocity estimation.
        """
        self._K  = K
        self._dt = dt

        # Accumulated global pose as 4×4 SE(3)
        self._pose = np.eye(4, dtype=np.float64)

        # Previous-step pose for velocity estimation
        self._prev_pose = np.eye(4, dtype=np.float64)

        self._frame_count = 0

    def set_K(self, K: np.ndarray) -> None:
        self._K = K.reshape(3, 3).astype(np.float64)

    def reset(self) -> None:
        self._pose       = np.eye(4)
        self._prev_pose  = np.eye(4)
        self._frame_count = 0

    def update(
        self,
        prev_pts: np.ndarray,    # (N, 2) float32
        curr_pts: np.ndarray,    # (N, 2) float32
        depth:    np.ndarray,    # float32 H×W depth in metres
        dt:       float | None = None,
    ) -> PoseEstimate:
        """
        Estimate incremental motion from matched feature pairs.

        Returns a PoseEstimate.  On failure returns a copy of the last valid
        estimate with valid=False.
        """
        self._frame_count += 1
        dt = dt or self._dt

        if self._K is None or len(prev_pts) < 8:
            return PoseEstimate(
                x=self._pose[0, 3], y=self._pose[1, 3],
                yaw=_rot_to_yaw(self._pose[:3, :3]), valid=False
            )

        # ── Essential matrix via 5-point RANSAC ──────────────────────────────
        E, mask = cv2.findEssentialMat(
            prev_pts, curr_pts, self._K,
            method=cv2.RANSAC,
            prob=_RANSAC_PROB,
            threshold=_RANSAC_THRESHOLD,
        )
        if E is None or mask is None:
            return self._invalid_est()

        inlier_mask  = mask.ravel().astype(bool)
        n_inliers    = int(inlier_mask.sum())
        if n_inliers < _MIN_INLIERS:
            return self._invalid_est()

        p1_in = prev_pts[inlier_mask]
        p2_in = curr_pts[inlier_mask]

        # ── Decompose E → R, t ────────────────────────────────────────────────
        _, R, t_unit, _ = cv2.recoverPose(E, p1_in, p2_in, self._K)
        # t_unit is a unit vector; we need metric scale

        # ── Scale recovery from depth ─────────────────────────────────────────
        scale = _median_depth(depth, p1_in)
        if scale <= 0:
            # Fallback: try current frame depth
            scale = _median_depth(depth, p2_in)
        if scale <= 0:
            return self._invalid_est()

        # Metric translation in camera frame
        t_metric = t_unit.ravel() * scale

        # ── Build incremental SE(3) in camera frame ───────────────────────────
        T_inc_cam = np.eye(4, dtype=np.float64)
        T_inc_cam[:3, :3] = R
        T_inc_cam[:3,  3] = t_metric

        # ── Rotate into base_link frame ───────────────────────────────────────
        R_cb = _R_CAM_TO_BASE
        R_inc_base = R_cb @ R @ R_cb.T
        t_inc_base = R_cb @ t_metric

        T_inc = np.eye(4, dtype=np.float64)
        T_inc[:3, :3] = R_inc_base
        T_inc[:3,  3] = t_inc_base

        # ── Accumulate global pose ────────────────────────────────────────────
        self._prev_pose = self._pose.copy()
        self._pose = self._pose @ T_inc

        # ── Extract pose + velocity ───────────────────────────────────────────
        x   = float(self._pose[0, 3])
        y   = float(self._pose[1, 3])
        z   = float(self._pose[2, 3])
        yaw = _rot_to_yaw(self._pose[:3, :3])

        dx  = float(self._pose[0, 3] - self._prev_pose[0, 3])
        dy  = float(self._pose[1, 3] - self._prev_pose[1, 3])
        dyaw = yaw - _rot_to_yaw(self._prev_pose[:3, :3])
        # Wrap dyaw to [-π, π]
        dyaw = (dyaw + math.pi) % (2 * math.pi) - math.pi

        vx    = math.hypot(dx, dy) / max(dt, 1e-6)
        # Signed vx: positive if moving forward
        if abs(dx) + abs(dy) > 0:
            heading_now = yaw
            move_dir    = math.atan2(dy, dx)
            if abs(move_dir - heading_now) > math.pi / 2:
                vx = -vx

        omega = dyaw / max(dt, 1e-6)

        return PoseEstimate(
            x=x, y=y, z=z, yaw=yaw,
            vx=vx, vy=0.0, omega=omega,
            valid=True,
        )

    def correct_from_loop_closure(
        self,
        rtabmap_x: float,
        rtabmap_y: float,
        rtabmap_yaw: float,
        alpha: float = 0.3,
    ) -> None:
        """
        Soft-correct accumulated pose using RTAB-Map's graph-optimized position.

        Called when RTAB-Map reports a pose significantly different from ours
        (indicating a loop closure correction).  alpha controls how strongly
        we shift toward RTAB-Map's estimate (0=ignore, 1=fully accept).
        """
        self._pose[0, 3] += alpha * (rtabmap_x   - self._pose[0, 3])
        self._pose[1, 3] += alpha * (rtabmap_y   - self._pose[1, 3])
        curr_yaw = _rot_to_yaw(self._pose[:3, :3])
        dyaw     = rtabmap_yaw - curr_yaw
        dyaw     = (dyaw + math.pi) % (2 * math.pi) - math.pi
        correction_yaw = curr_yaw + alpha * dyaw
        c, s = math.cos(correction_yaw), math.sin(correction_yaw)
        self._pose[:3, :3] = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])

    def _invalid_est(self) -> PoseEstimate:
        return PoseEstimate(
            x=self._pose[0, 3],
            y=self._pose[1, 3],
            yaw=_rot_to_yaw(self._pose[:3, :3]),
            valid=False,
        )
