"""Unit tests for ground_plane.py — RANSAC fitting and point classification."""

import math
import numpy as np
import pytest

from saltybot_obstacle_detect.ground_plane import (
    fit_ground_plane,
    height_above_plane,
    obstacle_mask,
    ground_mask,
)


def make_flat_ground(n=500, y_ground=1.0, noise=0.02, rng=None):
    """Flat ground plane at y = y_ground (camera frame: Y down)."""
    if rng is None:
        rng = np.random.default_rng(0)
    x = rng.uniform(-2.0, 2.0, n).astype(np.float64)
    z = rng.uniform(0.5, 4.0, n).astype(np.float64)
    y = np.full(n, y_ground) + rng.normal(0, noise, n)
    return np.stack([x, y, z], axis=1)


def make_obstacle_above_ground(cx=0.0, y_ground=1.0, height=0.3, n=50, rng=None):
    """Points forming a box above the ground plane."""
    if rng is None:
        rng = np.random.default_rng(1)
    x = rng.uniform(cx - 0.15, cx + 0.15, n)
    z = rng.uniform(1.5, 2.0, n)
    y = np.full(n, y_ground - height)   # Y_down: smaller y = higher in world
    return np.stack([x, y, z], axis=1)


# ── fit_ground_plane ──────────────────────────────────────────────────────────

def test_fit_flat_ground_returns_result():
    pts   = make_flat_ground(n=400)
    plane = fit_ground_plane(pts, n_iter=50)
    assert plane is not None


def test_fit_ground_normal_points_upward():
    """Normal should have negative Y component (upward in camera frame)."""
    pts   = make_flat_ground(n=400)
    plane = fit_ground_plane(pts, n_iter=50)
    assert plane is not None
    normal, d = plane
    assert normal[1] < 0.0, f"Expected normal.y < 0, got {normal[1]:.3f}"


def test_fit_ground_normal_unit_vector():
    pts   = make_flat_ground(n=300)
    plane = fit_ground_plane(pts, n_iter=50)
    assert plane is not None
    normal, _ = plane
    assert abs(np.linalg.norm(normal) - 1.0) < 1e-6


def test_fit_too_few_points_returns_none():
    pts   = make_flat_ground(n=5)
    plane = fit_ground_plane(pts, n_iter=10)
    assert plane is None


def test_fit_empty_returns_none():
    plane = fit_ground_plane(np.zeros((0, 3)))
    assert plane is None


def test_fit_tilted_ground():
    """Tilted ground plane (robot on slope) should still be found."""
    rng = np.random.default_rng(42)
    n   = 400
    x   = rng.uniform(-2.0, 2.0, n)
    z   = rng.uniform(0.5, 4.0, n)
    # Tilt: y = 0.5 + 0.1*z (plane tilted 5.7° forward)
    y   = 0.5 + 0.1 * z + rng.normal(0, 0.02, n)
    pts = np.stack([x, y, z], axis=1)
    plane = fit_ground_plane(pts, n_iter=80)
    assert plane is not None


# ── height_above_plane ────────────────────────────────────────────────────────

def test_ground_pts_near_zero_height():
    ground_pts = make_flat_ground(n=200, y_ground=1.0, noise=0.005)
    plane      = fit_ground_plane(ground_pts, n_iter=60)
    assert plane is not None
    heights    = height_above_plane(ground_pts, plane)
    assert np.abs(heights).mean() < 0.10   # mean residual < 10 cm


def test_obstacle_pts_positive_height():
    rng        = np.random.default_rng(7)
    ground_pts = make_flat_ground(n=300, y_ground=1.0, rng=rng)
    obs_pts    = make_obstacle_above_ground(y_ground=1.0, height=0.3, rng=rng)
    all_pts    = np.vstack([ground_pts, obs_pts])
    plane      = fit_ground_plane(all_pts, n_iter=60)
    assert plane is not None
    heights    = height_above_plane(obs_pts, plane)
    assert heights.mean() > 0.1, f"Expected obstacles above ground, got mean h={heights.mean():.3f}"


# ── obstacle_mask ─────────────────────────────────────────────────────────────

def test_obstacle_mask_selects_correct_points():
    rng        = np.random.default_rng(3)
    ground_pts = make_flat_ground(n=300, y_ground=1.0, rng=rng)
    obs_pts    = make_obstacle_above_ground(y_ground=1.0, height=0.4, n=60, rng=rng)
    all_pts    = np.vstack([ground_pts, obs_pts])
    plane      = fit_ground_plane(all_pts, n_iter=60)
    assert plane is not None
    mask = obstacle_mask(all_pts, plane, min_height_m=0.05, max_height_m=0.80)
    # Most of the obs_pts should be selected
    n_ground_selected = mask[:300].sum()
    n_obs_selected    = mask[300:].sum()
    assert n_obs_selected > 40, f"Expected most obs selected, got {n_obs_selected}/60"
    assert n_ground_selected < 30, f"Expected few ground points selected, got {n_ground_selected}/300"


# ── ground_mask ───────────────────────────────────────────────────────────────

def test_ground_mask_covers_inliers():
    pts   = make_flat_ground(n=300, noise=0.01)
    plane = fit_ground_plane(pts, n_iter=60)
    assert plane is not None
    gmask = ground_mask(pts, plane, inlier_thresh=0.08)
    assert gmask.mean() > 0.7   # at least 70% of flat ground should be ground inliers
