"""
_obstacle_size.py — Depth-based obstacle size estimation by projecting LIDAR
clusters into the D435i depth image (no ROS2 deps).

Algorithm
---------
1. Transform each LIDAR cluster centroid from LIDAR frame to camera frame
   using a configurable translation-only extrinsic (typical co-mounted setup).
2. Project to depth-image pixel coordinates via the pinhole camera model.
3. Sample the D435i depth image in a window around the projected pixel to
   obtain a robust Z estimate (uint16 mm → metres).
4. Derive horizontal width (metres) directly from the LIDAR cluster bounding
   box, which is more reliable than image-based re-measurement.
5. Derive vertical height (metres) by scanning a vertical strip in the depth
   image: collect pixels whose depth is within `z_tol` of the sampled Z, then
   compute the row extent and back-project using the camera fy.
6. Return an ObstacleSizeEstimate per cluster.

Coordinate frames
-----------------
LIDAR frame (+X forward, +Y left, +Z up — 2-D scan at z=0):
  x_lidar = forward distance (metres)
  y_lidar = lateral distance (metres, positive = left)
  z_lidar = 0   (2-D LIDAR always lies in the horizontal plane)

Camera frame (+X right, +Y down, +Z forward — standard pinhole):
  x_cam =  −y_lidar + ex
  y_cam =  −z_lidar + ey  =  ey          (z_lidar=0 for 2-D LIDAR)
  z_cam =   x_lidar + ez

Extrinsic parameters (CameraParams.ex/ey/ez) give the LIDAR origin position
in the camera coordinate frame:
  ex   — lateral offset  (positive = LIDAR origin is to the right of camera)
  ey   — vertical offset (positive = LIDAR origin is below camera)
  ez   — forward offset  (positive = LIDAR origin is in front of camera)

Public API
----------
  CameraParams           dataclass — camera intrinsics + co-mount extrinsics
  ObstacleSizeEstimate   NamedTuple — per-cluster result

  lidar_to_camera(x_l, y_l, params)          → (x_cam, y_cam, z_cam)
  project_to_pixel(x_cam, y_cam, z_cam, p)   → (u, v) or None if out-of-bounds
  sample_depth_median(depth_u16, u, v, win, scale) → (depth_m, n_valid)
  estimate_height(depth_u16, u_c, v_c, z_ref, params,
                  search_rows, col_hw, z_tol) → height_m
  estimate_cluster_size(cluster, depth_u16, params,
                        depth_window, search_rows, col_hw, z_tol,
                        cluster_id)           → ObstacleSizeEstimate
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import NamedTuple, Optional, Tuple

import numpy as np


# ── Camera parameters ─────────────────────────────────────────────────────────

@dataclass
class CameraParams:
    """Pinhole intrinsics and LIDAR→camera extrinsics for the D435i.

    Intrinsics (D435i 640×480 depth stream defaults — override from CameraInfo):
      fx, fy  : focal lengths in pixels
      cx, cy  : principal point in pixels
      width, height : image dimensions in pixels
      depth_scale   : multiply raw uint16 value by this to get metres
                      D435i default = 0.001 (raw value is millimetres)

    Extrinsics (LIDAR origin position in camera frame):
      ex, ey, ez : translation (metres) — see module docstring
    """
    fx:           float = 383.0    # D435i 640×480 depth approx
    fy:           float = 383.0
    cx:           float = 320.0
    cy:           float = 240.0
    width:        int   = 640
    height:       int   = 480
    depth_scale:  float = 0.001    # mm → m

    # LIDAR→camera extrinsics: LIDAR origin position in camera frame
    ex: float = 0.0    # lateral (m); 0 = LIDAR directly in front of camera
    ey: float = 0.05   # vertical (m); +0.05 = LIDAR 5 cm below camera
    ez: float = 0.0    # forward  (m); 0 = LIDAR at same depth as camera


# ── Result type ───────────────────────────────────────────────────────────────

class ObstacleSizeEstimate(NamedTuple):
    """Size estimate for one LIDAR cluster projected into the depth image."""
    obstacle_id:  int    # user-supplied cluster index (0-based) or track ID
    centroid_x:   float  # LIDAR-frame forward distance (metres)
    centroid_y:   float  # LIDAR-frame lateral distance (metres, +left)
    depth_z:      float  # D435i sampled depth at projected centroid (metres)
    width_m:      float  # horizontal size, LIDAR-derived (metres)
    height_m:     float  # vertical size, depth-image-derived (metres)
    pixel_u:      int    # projected centroid column (-1 = out of image)
    pixel_v:      int    # projected centroid row    (-1 = out of image)
    lidar_range:  float  # range from LIDAR origin to centroid (metres)
    confidence:   float  # 0.0–1.0; based on depth sample quality


# ── Coordinate transform ──────────────────────────────────────────────────────

def lidar_to_camera(
    x_lidar: float,
    y_lidar: float,
    params:  CameraParams,
) -> Tuple[float, float, float]:
    """
    Transform a 2-D LIDAR point to the camera coordinate frame.

    Parameters
    ----------
    x_lidar : forward distance in LIDAR frame (metres)
    y_lidar : lateral distance in LIDAR frame (metres, +left)
    params  : CameraParams with extrinsic translation ex/ey/ez

    Returns
    -------
    (x_cam, y_cam, z_cam) in camera frame (metres)
      x_cam = right, y_cam = down, z_cam = forward
    """
    x_cam = -y_lidar + params.ex
    y_cam =  params.ey          # LIDAR z=0 for 2-D scan
    z_cam =  x_lidar + params.ez
    return float(x_cam), float(y_cam), float(z_cam)


# ── Pinhole projection ────────────────────────────────────────────────────────

def project_to_pixel(
    x_cam:  float,
    y_cam:  float,
    z_cam:  float,
    params: CameraParams,
) -> Optional[Tuple[int, int]]:
    """
    Project a 3-D camera-frame point to image pixel coordinates.

    Returns None if z_cam ≤ 0 or the projection falls outside the image.

    Returns
    -------
    (u, v) integer pixel coordinates (column, row), or None.
    """
    if z_cam <= 0.0:
        return None
    u = params.fx * x_cam / z_cam + params.cx
    v = params.fy * y_cam / z_cam + params.cy
    ui, vi = int(round(u)), int(round(v))
    if 0 <= ui < params.width and 0 <= vi < params.height:
        return ui, vi
    return None


# ── Depth sampling ────────────────────────────────────────────────────────────

def sample_depth_median(
    depth_u16:   np.ndarray,
    u:           int,
    v:           int,
    window_px:   int   = 5,
    depth_scale: float = 0.001,
) -> Tuple[float, int]:
    """
    Median depth in a square window centred on (u, v) in a uint16 depth image.

    Zeros (invalid readings) are excluded from the median.

    Parameters
    ----------
    depth_u16   : (H, W) uint16 depth image (raw units, e.g. mm for D435i)
    u, v        : centre pixel (column, row)
    window_px   : half-side of the sampling square (window of 2*window_px+1)
    depth_scale : scale factor to convert raw units to metres (0.001 for mm)

    Returns
    -------
    (depth_m, n_valid) — median depth in metres and count of valid pixels.
    depth_m = 0.0 when n_valid == 0.
    """
    h, w = depth_u16.shape
    r0 = max(0, v - window_px)
    r1 = min(h, v + window_px + 1)
    c0 = max(0, u - window_px)
    c1 = min(w, u + window_px + 1)

    patch = depth_u16[r0:r1, c0:c1].astype(np.float64)
    valid = patch[patch > 0]
    if len(valid) == 0:
        return 0.0, 0
    return float(np.median(valid) * depth_scale), int(len(valid))


# ── Height estimation ─────────────────────────────────────────────────────────

def estimate_height(
    depth_u16:   np.ndarray,
    u_c:         int,
    v_c:         int,
    z_ref:       float,
    params:      CameraParams,
    search_rows: int   = 120,
    col_hw:      int   = 10,
    z_tol:       float = 0.30,
) -> float:
    """
    Estimate the vertical height (metres) of an obstacle in the depth image.

    Scans a vertical strip centred on (u_c, v_c) and finds the row extent of
    pixels whose depth is within `z_tol` metres of `z_ref`.

    Parameters
    ----------
    depth_u16   : (H, W) uint16 depth image
    u_c, v_c    : projected centroid pixel
    z_ref       : reference depth (metres) — from sample_depth_median
    params      : CameraParams (needs fy, depth_scale)
    search_rows : half-height of the search strip (rows above/below v_c)
    col_hw      : half-width  of the search strip (columns left/right of u_c)
    z_tol       : depth tolerance — pixels within z_ref ± z_tol are included

    Returns
    -------
    height_m — vertical extent in metres; 0.0 if fewer than 2 valid rows found.
    """
    if z_ref <= 0.0:
        return 0.0

    h, w = depth_u16.shape
    r0 = max(0, v_c - search_rows)
    r1 = min(h, v_c + search_rows + 1)
    c0 = max(0, u_c - col_hw)
    c1 = min(w, u_c + col_hw + 1)

    strip = depth_u16[r0:r1, c0:c1].astype(np.float64) * params.depth_scale

    # For each row in the strip, check if any valid pixel is within z_tol of z_ref
    z_min = z_ref - z_tol
    z_max = z_ref + z_tol
    valid_mask = (strip > 0) & (strip >= z_min) & (strip <= z_max)
    valid_rows = np.any(valid_mask, axis=1)   # (n_rows,) bool

    if valid_rows.sum() < 2:
        return 0.0

    row_indices = np.where(valid_rows)[0]     # rows within the strip
    v_top_strip    = int(row_indices[0])
    v_bottom_strip = int(row_indices[-1])

    # Convert to absolute image rows
    v_top    = r0 + v_top_strip
    v_bottom = r0 + v_bottom_strip

    row_extent = v_bottom - v_top
    if row_extent <= 0:
        return 0.0

    # Back-project row extent to metres at depth z_ref
    height_m = float(row_extent) * z_ref / params.fy
    return float(height_m)


# ── Main entry point ──────────────────────────────────────────────────────────

def estimate_cluster_size(
    cluster,                        # _lidar_clustering.Cluster NamedTuple
    depth_u16:   np.ndarray,        # (H, W) uint16 depth image
    params:      CameraParams,
    depth_window: int   = 5,        # half-side for depth median sampling
    search_rows:  int   = 120,      # half-height of height search strip
    col_hw:       int   = 10,       # half-width of height search strip
    z_tol:        float = 0.30,     # depth tolerance for height estimation
    obstacle_id:  int   = 0,        # caller-assigned ID
) -> ObstacleSizeEstimate:
    """
    Estimate the 3-D size of one LIDAR cluster using the D435i depth image.

    Parameters
    ----------
    cluster      : Cluster from _lidar_clustering.cluster_points()
    depth_u16    : uint16 depth image (H × W), values in millimetres
    params       : CameraParams with intrinsics + extrinsics
    depth_window : half-side of median depth sampling window (pixels)
    search_rows  : half-height for vertical height search (rows)
    col_hw       : half-width  for vertical height search (columns)
    z_tol        : depth tolerance when collecting obstacle pixels (metres)
    obstacle_id  : arbitrary integer ID for this estimate

    Returns
    -------
    ObstacleSizeEstimate
    """
    cx, cy = float(cluster.centroid[0]), float(cluster.centroid[1])
    lidar_range = float(math.sqrt(cx * cx + cy * cy))

    # Transform centroid to camera frame
    x_cam, y_cam, z_cam = lidar_to_camera(cx, cy, params)

    # Project to depth image pixel
    px = project_to_pixel(x_cam, y_cam, z_cam, params)
    if px is None:
        return ObstacleSizeEstimate(
            obstacle_id  = obstacle_id,
            centroid_x   = cx,
            centroid_y   = cy,
            depth_z      = 0.0,
            width_m      = float(cluster.width_m),
            height_m     = 0.0,
            pixel_u      = -1,
            pixel_v      = -1,
            lidar_range  = lidar_range,
            confidence   = 0.0,
        )

    u, v = px

    # Sample depth
    depth_m, n_valid = sample_depth_median(
        depth_u16, u, v, window_px=depth_window,
        depth_scale=params.depth_scale,
    )

    # Use LIDAR range as fallback when depth image has no valid reading
    if depth_m <= 0.0:
        depth_m = z_cam
        depth_conf = 0.0
    else:
        # Confidence: ratio of valid pixels in window + how close to expected Z
        max_valid = (2 * depth_window + 1) ** 2
        fill_ratio = min(1.0, n_valid / max(max_valid * 0.25, 1.0))
        z_err = abs(depth_m - z_cam)
        z_conf = max(0.0, 1.0 - z_err / max(z_cam, 0.1))
        depth_conf = float(0.6 * fill_ratio + 0.4 * z_conf)

    # Width: LIDAR gives reliable horizontal measurement
    width_m = float(cluster.width_m)

    # Height: from depth image vertical strip
    height_m = estimate_height(
        depth_u16, u, v, depth_m, params,
        search_rows=search_rows, col_hw=col_hw, z_tol=z_tol,
    )

    return ObstacleSizeEstimate(
        obstacle_id  = obstacle_id,
        centroid_x   = cx,
        centroid_y   = cy,
        depth_z      = float(depth_m),
        width_m      = width_m,
        height_m     = float(height_m),
        pixel_u      = int(u),
        pixel_v      = int(v),
        lidar_range  = lidar_range,
        confidence   = float(depth_conf),
    )
