"""
_pointing_ray.py — 3-D pointing ray computation from hand landmarks + depth (no ROS2 deps).

Given MediaPipe Hands 21-landmark output, a depth image (float32, metres),
and camera intrinsic parameters, computes:
  - 3-D position of the INDEX_MCP (knuckle) — ray origin
  - 3-D direction (unit vector) INDEX_MCP → INDEX_TIP
  - Estimated pointing target (origin + direction * range)

Coordinate frame: camera_color_optical_frame (standard ROS optical frame,
+X right, +Y down, +Z into the scene).
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple


# MediaPipe Hands landmark indices (matches gesture_classifier.py)
_WRIST     = 0
_INDEX_MCP = 5
_INDEX_TIP = 8


def unproject(
    u: float, v: float, depth_m: float,
    fx: float, fy: float, cx: float, cy: float,
) -> Tuple[float, float, float]:
    """
    Back-project pixel (u, v) at depth depth_m to a 3-D point in the camera frame.

    Parameters
    ----------
    u, v    : pixel coordinates (float, origin top-left)
    depth_m : depth in metres
    fx, fy  : focal lengths (pixels)
    cx, cy  : principal point (pixels)

    Returns
    -------
    (X, Y, Z) in metres, camera_color_optical_frame
    """
    X = (u - cx) * depth_m / fx
    Y = (v - cy) * depth_m / fy
    Z = depth_m
    return X, Y, Z


def sample_depth(
    depth_img,
    u: float,
    v: float,
    window: int = 5,
    d_min: float = 0.1,
    d_max: float = 8.0,
) -> float:
    """
    Return the median of valid depth values in a (window × window) patch centred at (u, v).

    Parameters
    ----------
    depth_img : np.ndarray (H, W) float32, depth in metres
    u, v      : centre pixel (float)
    window    : patch side length (pixels)
    d_min     : minimum valid depth (m)
    d_max     : maximum valid depth (m)

    Returns
    -------
    float : median depth, or NaN if no valid samples in the patch
    """
    import numpy as np

    h, w = depth_img.shape[:2]
    half = window // 2
    u0 = max(0,     int(round(u)) - half)
    u1 = min(w,     int(round(u)) + half + 1)
    v0 = max(0,     int(round(v)) - half)
    v1 = min(h,     int(round(v)) + half + 1)

    patch = depth_img[v0:v1, u0:u1]
    valid = patch[(patch > d_min) & (patch < d_max)]
    if valid.size == 0:
        return float('nan')
    return float(np.median(valid))


def compute_pointing_ray(
    landmarks: List[Tuple[float, float, float]],
    depth_img,
    fx: float,
    fy: float,
    cx: float,
    cy: float,
    img_w: int,
    img_h: int,
    ref_distance: float = 2.0,
    depth_window: int = 5,
) -> Optional[dict]:
    """
    Compute a 3-D pointing ray from MediaPipe Hands landmarks and a depth image.

    Parameters
    ----------
    landmarks     : list of 21 (nx, ny, nz) normalised MediaPipe landmarks
                    (nx ∈ [0,1] left→right, ny ∈ [0,1] top→bottom)
    depth_img     : np.ndarray (H, W) float32, depth in metres
    fx, fy        : focal lengths (pixels)
    cx, cy        : principal point (pixels)
    img_w, img_h  : image dimensions (pixels)
    ref_distance  : fallback range when depth is unavailable (m)
    depth_window  : window size for depth median sampling (pixels)

    Returns
    -------
    dict with keys:
        origin_3d    : (X, Y, Z)  INDEX_MCP in camera frame (m)
        direction_3d : (dx, dy, dz)  unit vector MCP → TIP
        target_3d    : (X, Y, Z)  estimated pointing target (m)
        range_m      : distance from origin to target (m)
        mcp_uv       : (u, v) pixel coords of INDEX_MCP
        tip_uv       : (u, v) pixel coords of INDEX_TIP
    or None if the ray cannot be computed (both depths invalid).
    """
    mcp_lm = landmarks[_INDEX_MCP]
    tip_lm = landmarks[_INDEX_TIP]

    mcp_u, mcp_v = mcp_lm[0] * img_w, mcp_lm[1] * img_h
    tip_u, tip_v = tip_lm[0] * img_w, tip_lm[1] * img_h

    d_mcp = sample_depth(depth_img, mcp_u, mcp_v, depth_window)
    d_tip = sample_depth(depth_img, tip_u, tip_v, depth_window)

    # Need at least one valid depth
    if math.isnan(d_mcp) and math.isnan(d_tip):
        return None

    # Fill missing depth with the other measurement
    if math.isnan(d_mcp):
        d_mcp = d_tip
    if math.isnan(d_tip):
        d_tip = d_mcp

    mcp_3d = unproject(mcp_u, mcp_v, d_mcp, fx, fy, cx, cy)
    tip_3d = unproject(tip_u, tip_v, d_tip, fx, fy, cx, cy)

    # Ray direction: MCP → TIP
    dx = tip_3d[0] - mcp_3d[0]
    dy = tip_3d[1] - mcp_3d[1]
    dz = tip_3d[2] - mcp_3d[2]
    length = math.sqrt(dx * dx + dy * dy + dz * dz)

    if length < 1e-4:
        # Degenerate: MCP and TIP project to the same 3-D point
        # Fall back to the image-plane direction extended at ref_distance
        px = (tip_u - mcp_u) / (img_w or 1)
        py = (tip_v - mcp_v) / (img_h or 1)
        dx, dy, dz = px, py, 1.0
        length = math.sqrt(dx * dx + dy * dy + dz * dz) or 1.0

    dx /= length
    dy /= length
    dz /= length

    range_m = ref_distance
    target_3d = (
        mcp_3d[0] + dx * range_m,
        mcp_3d[1] + dy * range_m,
        mcp_3d[2] + dz * range_m,
    )

    return {
        'origin_3d':    mcp_3d,
        'direction_3d': (dx, dy, dz),
        'target_3d':    target_3d,
        'range_m':      range_m,
        'mcp_uv':       (mcp_u, mcp_v),
        'tip_uv':       (tip_u, tip_v),
    }
