"""
test_obstacle_size.py — pytest tests for _obstacle_size.py (no ROS2 required).

Tests cover:
  CameraParams  — defaults and construction
  lidar_to_camera — coordinate transform
  project_to_pixel — pinhole projection + bounds checks
  sample_depth_median — empty/uniform/sparse depth images
  estimate_height — flat/obstacle/edge cases
  estimate_cluster_size — full pipeline with synthetic cluster + depth
"""

from __future__ import annotations

import math
from typing import List

import numpy as np
import pytest

from saltybot_bringup._obstacle_size import (
    CameraParams,
    ObstacleSizeEstimate,
    estimate_cluster_size,
    estimate_height,
    lidar_to_camera,
    project_to_pixel,
    sample_depth_median,
)
# Import Cluster helper from lidar_clustering for building test fixtures
from saltybot_bringup._lidar_clustering import Cluster


# ── Helpers ───────────────────────────────────────────────────────────────────

def _default_params(**kw) -> CameraParams:
    p = CameraParams()
    for k, v in kw.items():
        object.__setattr__(p, k, v)
    return p


def _blank_depth(h: int = 480, w: int = 640, val: int = 0) -> np.ndarray:
    return np.full((h, w), val, dtype=np.uint16)


def _obstacle_depth(
    h: int = 480, w: int = 640,
    depth_mm: int = 2000,
    u_c: int = 320, v_c: int = 240,
    half_w: int = 30, half_h: int = 60,
) -> np.ndarray:
    """Create a synthetic depth image with a rectangular obstacle."""
    img = _blank_depth(h, w)
    r0, r1 = max(0, v_c - half_h), min(h, v_c + half_h + 1)
    c0, c1 = max(0, u_c - half_w), min(w, u_c + half_w + 1)
    img[r0:r1, c0:c1] = depth_mm
    return img


def _make_cluster(
    cx: float = 2.0, cy: float = 0.0,
    width_m: float = 0.4, depth_m: float = 0.3,
    n_pts: int = 10,
) -> Cluster:
    """Construct a minimal Cluster NamedTuple for testing."""
    pts = np.column_stack([
        np.full(n_pts, cx),
        np.full(n_pts, cy),
    ]).astype(np.float64)
    centroid = np.array([cx, cy], dtype=np.float64)
    bbox_min = centroid - np.array([width_m / 2, depth_m / 2])
    bbox_max = centroid + np.array([width_m / 2, depth_m / 2])
    return Cluster(
        points=pts, centroid=centroid,
        bbox_min=bbox_min, bbox_max=bbox_max,
        width_m=width_m, depth_m=depth_m,
    )


# ── CameraParams ──────────────────────────────────────────────────────────────

class TestCameraParams:
    def test_defaults(self):
        p = CameraParams()
        assert p.fx == pytest.approx(383.0)
        assert p.fy == pytest.approx(383.0)
        assert p.cx == pytest.approx(320.0)
        assert p.cy == pytest.approx(240.0)
        assert p.width  == 640
        assert p.height == 480
        assert p.depth_scale == pytest.approx(0.001)
        assert p.ey == pytest.approx(0.05)

    def test_custom(self):
        p = CameraParams(fx=400.0, fy=400.0, cx=330.0, cy=250.0)
        assert p.fx == pytest.approx(400.0)


# ── lidar_to_camera ───────────────────────────────────────────────────────────

class TestLidarToCamera:
    def test_forward_no_offset(self):
        """A point directly in front of the LIDAR (y_lidar=0) → x_cam=0."""
        p = CameraParams(ex=0.0, ey=0.0, ez=0.0)
        x_cam, y_cam, z_cam = lidar_to_camera(3.0, 0.0, p)
        assert x_cam == pytest.approx(0.0)
        assert y_cam == pytest.approx(0.0)
        assert z_cam == pytest.approx(3.0)

    def test_lateral_maps_to_x(self):
        """y_lidar positive (left) → x_cam negative (right convention)."""
        p = CameraParams(ex=0.0, ey=0.0, ez=0.0)
        x_cam, y_cam, z_cam = lidar_to_camera(0.0, 1.0, p)
        assert x_cam == pytest.approx(-1.0)
        assert z_cam == pytest.approx(0.0)

    def test_extrinsic_offset_applied(self):
        """Extrinsic translation is added to camera-frame coords."""
        p = CameraParams(ex=0.1, ey=0.05, ez=-0.02)
        x_cam, y_cam, z_cam = lidar_to_camera(2.0, 0.0, p)
        assert x_cam == pytest.approx(0.1)      # -y_lidar + ex = 0 + 0.1
        assert y_cam == pytest.approx(0.05)     # 0 + ey
        assert z_cam == pytest.approx(1.98)     # 2.0 + (-0.02)

    def test_zero_point(self):
        p = CameraParams(ex=0.0, ey=0.0, ez=0.0)
        x_cam, y_cam, z_cam = lidar_to_camera(0.0, 0.0, p)
        assert z_cam == pytest.approx(0.0)


# ── project_to_pixel ──────────────────────────────────────────────────────────

class TestProjectToPixel:
    def test_centre_projects_to_principal(self):
        """A point exactly on the optical axis projects to (cx, cy)."""
        p = CameraParams(fx=400.0, fy=400.0, cx=320.0, cy=240.0)
        px = project_to_pixel(0.0, 0.0, 2.0, p)
        assert px is not None
        u, v = px
        assert u == 320
        assert v == 240

    def test_negative_z_returns_none(self):
        p = CameraParams()
        assert project_to_pixel(0.0, 0.0, -1.0, p) is None

    def test_zero_z_returns_none(self):
        p = CameraParams()
        assert project_to_pixel(0.0, 0.0, 0.0, p) is None

    def test_off_to_right(self):
        """A point to the right of axis should land to the right of cx."""
        p = CameraParams(fx=400.0, cx=320.0, cy=240.0)
        px = project_to_pixel(0.5, 0.0, 2.0, p)
        assert px is not None
        u, _ = px
        assert u > 320

    def test_out_of_image_returns_none(self):
        p = CameraParams(fx=100.0, fy=100.0, cx=50.0, cy=50.0,
                         width=100, height=100)
        # Very far off-axis → outside image
        px = project_to_pixel(10.0, 0.0, 1.0, p)
        assert px is None

    def test_known_projection(self):
        """Verify exact pixel for known 3D point."""
        p = CameraParams(fx=400.0, fy=400.0, cx=320.0, cy=240.0,
                         width=640, height=480)
        # x_cam=1.0, z_cam=4.0 → u = 400*(1/4)+320 = 420
        # y_cam=0.0            → v = 400*(0/4)+240 = 240
        px = project_to_pixel(1.0, 0.0, 4.0, p)
        assert px == (420, 240)


# ── sample_depth_median ───────────────────────────────────────────────────────

class TestSampleDepthMedian:
    def test_all_zeros_returns_zero(self):
        img = _blank_depth()
        d, n = sample_depth_median(img, 320, 240, window_px=5)
        assert d == pytest.approx(0.0)
        assert n == 0

    def test_uniform_image(self):
        """All pixels set to 2000 mm → median = 2.0 m."""
        img = _blank_depth(val=2000)
        d, n = sample_depth_median(img, 320, 240, window_px=5, depth_scale=0.001)
        assert d == pytest.approx(2.0)
        assert n > 0

    def test_scale_applied(self):
        img = _blank_depth(val=3000)
        d, _ = sample_depth_median(img, 320, 240, window_px=3, depth_scale=0.001)
        assert d == pytest.approx(3.0)

    def test_window_clips_at_image_edge(self):
        """Sampling near edge should not crash."""
        img = _blank_depth(val=1500)
        d, n = sample_depth_median(img, 0, 0, window_px=10)
        assert d > 0.0
        assert n > 0

    def test_sparse_window(self):
        """Only some pixels are valid — median should reflect valid ones."""
        img = _blank_depth()
        img[240, 320] = 1000   # single valid pixel
        d, n = sample_depth_median(img, 320, 240, window_px=5, depth_scale=0.001)
        assert d == pytest.approx(1.0)
        assert n == 1

    def test_mixed_window_median(self):
        """Median of [1000, 2000, 3000] mm = 2000 mm = 2.0 m."""
        img = _blank_depth()
        img[240, 319] = 1000
        img[240, 320] = 2000
        img[240, 321] = 3000
        d, n = sample_depth_median(img, 320, 240, window_px=1, depth_scale=0.001)
        assert d == pytest.approx(2.0)
        assert n == 3


# ── estimate_height ───────────────────────────────────────────────────────────

class TestEstimateHeight:
    def test_blank_image_returns_zero(self):
        p = CameraParams(fy=383.0)
        img = _blank_depth()
        h = estimate_height(img, 320, 240, z_ref=2.0, params=p)
        assert h == pytest.approx(0.0)

    def test_zero_z_ref_returns_zero(self):
        p = CameraParams()
        img = _blank_depth(val=2000)
        h = estimate_height(img, 320, 240, z_ref=0.0, params=p)
        assert h == pytest.approx(0.0)

    def test_known_height(self):
        """
        Obstacle occupies rows 180–300 (120 row-pixel span) at depth 2.0m.
        Expected height ≈ 120 * 2.0 / 383.0 ≈ 0.627 m.
        """
        p = CameraParams(fy=383.0, depth_scale=0.001)
        img = _blank_depth()
        # Obstacle at z=2000 mm, rows 180–300, cols 300–340
        img[180:301, 300:341] = 2000
        h = estimate_height(img, 320, 240, z_ref=2.0, params=p,
                             search_rows=200, col_hw=30, z_tol=0.3)
        expected = 120 * 2.0 / 383.0
        assert h == pytest.approx(expected, rel=0.05)

    def test_single_row_returns_zero(self):
        """Only one valid row → span=0 → height=0."""
        p = CameraParams(fy=383.0, depth_scale=0.001)
        img = _blank_depth()
        img[240, 315:326] = 2000   # single row
        h = estimate_height(img, 320, 240, z_ref=2.0, params=p,
                             search_rows=20, col_hw=10, z_tol=0.3)
        assert h == pytest.approx(0.0)

    def test_depth_outside_tolerance_excluded(self):
        """Pixels far from z_ref should not contribute to height."""
        p = CameraParams(fy=383.0, depth_scale=0.001)
        img = _blank_depth()
        # Two rows at very different depths — only one within z_tol
        img[220, 315:326] = 2000   # z = 2.0 m (within tolerance)
        img[260, 315:326] = 5000   # z = 5.0 m (outside tolerance)
        h = estimate_height(img, 320, 240, z_ref=2.0, params=p,
                             search_rows=100, col_hw=10, z_tol=0.3)
        assert h == pytest.approx(0.0)   # only 1 valid row → 0


# ── estimate_cluster_size ─────────────────────────────────────────────────────

class TestEstimateClusterSize:
    def test_returns_named_tuple(self):
        cluster = _make_cluster()
        img = _blank_depth(val=2000)
        est = estimate_cluster_size(cluster, img, CameraParams())
        assert isinstance(est, ObstacleSizeEstimate)

    def test_width_comes_from_lidar(self):
        """Width should equal the LIDAR cluster width_m."""
        cluster = _make_cluster(cx=3.0, cy=0.0, width_m=0.5)
        img = _blank_depth(val=3000)
        est = estimate_cluster_size(cluster, img, CameraParams())
        assert est.width_m == pytest.approx(0.5)

    def test_centroid_preserved(self):
        cluster = _make_cluster(cx=2.0, cy=0.3)
        img = _blank_depth(val=2000)
        est = estimate_cluster_size(cluster, img, CameraParams())
        assert est.centroid_x == pytest.approx(2.0)
        assert est.centroid_y == pytest.approx(0.3)

    def test_lidar_range_correct(self):
        cluster = _make_cluster(cx=3.0, cy=4.0)  # range = 5.0
        img = _blank_depth(val=5000)
        est = estimate_cluster_size(cluster, img, CameraParams())
        assert est.lidar_range == pytest.approx(5.0)

    def test_blank_depth_gives_zero_confidence(self):
        """No valid depth pixels → confidence=0, depth_z falls back to z_cam."""
        cluster = _make_cluster(cx=2.0, cy=0.0)
        img = _blank_depth()
        p   = CameraParams(ex=0.0, ey=0.0, ez=0.0)
        est = estimate_cluster_size(cluster, img, p)
        assert est.confidence == pytest.approx(0.0)
        assert est.depth_z > 0.0  # LIDAR-fallback z_cam

    def test_obstacle_in_depth_image(self):
        """With a real depth patch, confidence > 0 and pixel is in-image."""
        # Cluster at x=2m forward, y=0 → z_cam=2m, x_cam=0, y_cam=ey
        p = CameraParams(fy=383.0, fx=383.0, cx=320.0, cy=240.0,
                         depth_scale=0.001, ex=0.0, ey=0.05, ez=0.0)
        cluster = _make_cluster(cx=2.0, cy=0.0)
        # z_cam=2m, x_cam=0, y_cam=0.05
        # u = 383*0/2+320 = 320, v = 383*0.05/2+240 ≈ 249
        img = _obstacle_depth(depth_mm=2000, u_c=320, v_c=249,
                               half_w=30, half_h=80)
        est = estimate_cluster_size(cluster, img, p,
                                    depth_window=5, search_rows=120,
                                    col_hw=10, z_tol=0.3, obstacle_id=1)
        assert est.confidence > 0.0
        assert 0 <= est.pixel_u < 640
        assert 0 <= est.pixel_v < 480
        assert est.depth_z == pytest.approx(2.0, abs=0.1)
        assert est.obstacle_id == 1

    def test_behind_camera_returns_zero_confidence(self):
        """Cluster behind camera (x_lidar < 0 with no ez) → not projected."""
        p = CameraParams(ex=0.0, ey=0.0, ez=0.0)
        # x_lidar = -1.0 → z_cam = -1 + 0 = -1 → behind camera
        cluster = _make_cluster(cx=-1.0, cy=0.0)
        img = _blank_depth(val=1000)
        est = estimate_cluster_size(cluster, img, p)
        assert est.confidence == pytest.approx(0.0)
        assert est.pixel_u == -1
        assert est.pixel_v == -1

    def test_height_estimated_from_depth(self):
        """When obstacle spans rows, height_m > 0."""
        p = CameraParams(fy=383.0, fx=383.0, cx=320.0, cy=240.0,
                         depth_scale=0.001, ex=0.0, ey=0.05, ez=0.0)
        cluster = _make_cluster(cx=2.0, cy=0.0)
        # Expected pixel: u≈320, v≈249; give obstacle a 100-row span
        img = _obstacle_depth(depth_mm=2000, u_c=320, v_c=249,
                               half_w=20, half_h=50)
        est = estimate_cluster_size(cluster, img, p,
                                    search_rows=150, col_hw=15, z_tol=0.3)
        assert est.height_m > 0.0

    def test_obstacle_id_passed_through(self):
        cluster = _make_cluster()
        img = _blank_depth(val=2000)
        est = estimate_cluster_size(cluster, img, CameraParams(), obstacle_id=42)
        assert est.obstacle_id == 42

    def test_confidence_bounded(self):
        cluster = _make_cluster(cx=2.0)
        img = _blank_depth(val=2000)
        est = estimate_cluster_size(cluster, img, CameraParams())
        assert 0.0 <= est.confidence <= 1.0
