import numpy as np
import pytest
from aruco.ground_plane import (
    unproject_depth_to_points,
    detect_floor_plane,
    compute_consensus_plane,
    compute_floor_correction,
    FloorPlane,
    FloorCorrection,
)


def test_unproject_depth_to_points_simple():
    # Simple 3x3 depth map
    # K = identity for simplicity (fx=1, fy=1, cx=1, cy=1)
    # Pixel (1, 1) is center.
    # At (1, 1), u=1, v=1. x = (1-1)/1 = 0, y = (1-1)/1 = 0.
    # If depth is Z, point is (0, 0, Z).

    width, height = 3, 3
    K = np.array([[1, 0, 1], [0, 1, 1], [0, 0, 1]], dtype=np.float64)
    depth_map = np.zeros((height, width), dtype=np.float32)

    # Center pixel
    depth_map[1, 1] = 2.0

    # Top-left pixel (0, 0)
    # u=0, v=0. x = (0-1)/1 = -1. y = (0-1)/1 = -1.
    # Point: (-1*Z, -1*Z, Z)
    depth_map[0, 0] = 1.0

    points = unproject_depth_to_points(depth_map, K, depth_min=0.1, depth_max=5.0)

    # Should have 2 points (others are 0.0 which is < depth_min)
    assert points.shape == (2, 3)

    # Check center point
    # We don't know order, so check if expected points exist
    expected_center = np.array([0.0, 0.0, 2.0])
    expected_tl = np.array([-1.0, -1.0, 1.0])

    # Find matches
    has_center = np.any(np.all(np.isclose(points, expected_center, atol=1e-5), axis=1))
    has_tl = np.any(np.all(np.isclose(points, expected_tl, atol=1e-5), axis=1))

    assert has_center
    assert has_tl


def test_unproject_depth_to_points_stride():
    width, height = 10, 10
    K = np.eye(3)
    depth_map = np.ones((height, width), dtype=np.float32)

    points = unproject_depth_to_points(depth_map, K, stride=2)

    # 10x10 -> 5x5 = 25 points
    assert points.shape == (25, 3)


def test_unproject_depth_to_points_bounds():
    width, height = 3, 3
    K = np.eye(3)
    depth_map = np.array(
        [[0.05, 1.0, 11.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], dtype=np.float32
    )

    # 0.05 < 0.1 (min) -> excluded
    # 11.0 > 10.0 (max) -> excluded
    # 7 valid points
    points = unproject_depth_to_points(depth_map, K, depth_min=0.1, depth_max=10.0)
    assert points.shape == (7, 3)


def test_detect_floor_plane_perfect():
    # Create points on a perfect plane: y = -1.5 (floor at -1.5m)
    # Normal should be [0, 1, 0] (pointing up)
    # Plane eq: 0*x + 1*y + 0*z + d = 0 => y + d = 0 => -1.5 + d = 0 => d = 1.5

    # Generate grid of points
    x = np.linspace(-1, 1, 10)
    z = np.linspace(0, 5, 10)
    xx, zz = np.meshgrid(x, z)
    yy = np.full_like(xx, -1.5)

    points = np.stack([xx.flatten(), yy.flatten(), zz.flatten()], axis=1)

    # Add some noise to make it realistic but within threshold
    rng = np.random.default_rng(42)
    points += rng.normal(0, 0.001, points.shape)

    result = detect_floor_plane(points, distance_threshold=0.01, seed=42)

    assert result is not None
    assert isinstance(result, FloorPlane)
    normal = result.normal
    d = result.d
    inliers = result.num_inliers

    # Normal could be [0, 1, 0] or [0, -1, 0] depending on RANSAC
    # But we usually want it pointing "up" relative to camera or just consistent
    # Open3D segment_plane doesn't guarantee orientation

    # Check if it's vertical (y-axis aligned)
    assert abs(normal[1]) > 0.9

    # Check distance
    # If normal is [0, 1, 0], d should be 1.5
    # If normal is [0, -1, 0], d should be -1.5
    if normal[1] > 0:
        assert abs(d - 1.5) < 0.01
    else:
        assert abs(d + 1.5) < 0.01

    assert inliers == 100


def test_detect_floor_plane_with_outliers():
    # 100 inliers on floor y=-1.0
    inliers = np.zeros((100, 3))
    inliers[:, 0] = np.random.uniform(-1, 1, 100)
    inliers[:, 1] = -1.0
    inliers[:, 2] = np.random.uniform(1, 5, 100)

    # 50 outliers (walls, noise)
    outliers = np.random.uniform(-2, 2, (50, 3))
    outliers[:, 1] = np.random.uniform(-0.5, 1.0, 50)  # Above floor

    points = np.vstack([inliers, outliers])

    result = detect_floor_plane(points, distance_threshold=0.02, seed=42)

    assert result is not None
    assert abs(result.normal[1]) > 0.9  # Vertical normal
    assert result.num_inliers >= 100  # Should find all inliers


def test_detect_floor_plane_insufficient_points():
    points = np.array([[0, 0, 0], [1, 0, 0]])  # Only 2 points
    result = detect_floor_plane(points)
    assert result is None


def test_detect_floor_plane_no_plane():
    # Random cloud
    points = np.random.uniform(-1, 1, (100, 3))
    # With high threshold it might find something, but with low threshold and random points...
    # Actually RANSAC almost always finds *something* in 3 points.
    # But let's test that it runs without crashing.
    result = detect_floor_plane(points, distance_threshold=0.001, seed=42)
    # It might return None if it can't find enough inliers for a model
    # Open3D segment_plane usually returns a model even if bad.
    # We'll check our wrapper behavior.
    pass


def test_compute_consensus_plane_simple():
    # Two identical planes
    planes = [
        FloorPlane(normal=np.array([0, 1, 0]), d=1.5),
        FloorPlane(normal=np.array([0, 1, 0]), d=1.5),
    ]

    result = compute_consensus_plane(planes)

    np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6)
    assert abs(result.d - 1.5) < 1e-6


def test_compute_consensus_plane_weighted():
    # Two planes, one with more weight
    # Plane 1: normal [0, 1, 0], d=1.0
    # Plane 2: normal [0, 1, 0], d=2.0
    # Weights: [1, 3] -> weighted avg d should be (1*1 + 3*2)/4 = 7/4 = 1.75

    planes = [
        FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
        FloorPlane(normal=np.array([0, 1, 0]), d=2.0),
    ]
    weights = [1.0, 3.0]

    result = compute_consensus_plane(planes, weights)

    np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6)
    assert abs(result.d - 1.75) < 1e-6


def test_compute_consensus_plane_averaging_normals():
    # Two planes with slightly different normals
    # n1 = [0, 1, 0]
    # n2 = [0.1, 0.995, 0] (approx)

    n1 = np.array([0, 1, 0], dtype=np.float64)
    n2 = np.array([0.1, 1.0, 0], dtype=np.float64)
    n2 /= np.linalg.norm(n2)

    planes = [FloorPlane(normal=n1, d=1.0), FloorPlane(normal=n2, d=1.0)]

    result = compute_consensus_plane(planes)

    # Expected normal is roughly average (normalized)
    avg_n = (n1 + n2) / 2.0
    avg_d = 1.0  # (1.0 + 1.0) / 2.0
    norm = np.linalg.norm(avg_n)
    expected_n = avg_n / norm
    expected_d = avg_d / norm

    np.testing.assert_allclose(result.normal, expected_n, atol=1e-6)
    assert abs(result.d - expected_d) < 1e-6


def test_compute_consensus_plane_empty():
    with pytest.raises(ValueError):
        compute_consensus_plane([])


def test_compute_consensus_plane_flip_normals():
    # If one normal is flipped, it should be flipped back to align with the majority/first
    # n1 = [0, 1, 0]
    # n2 = [0, -1, 0]
    # d1 = 1.0
    # d2 = -1.0 (same plane, just flipped normal)

    planes = [
        FloorPlane(normal=np.array([0, 1, 0]), d=1.0),
        FloorPlane(normal=np.array([0, -1, 0]), d=-1.0),
    ]

    result = compute_consensus_plane(planes)

    # Should align to first one (arbitrary choice, but consistent)

    np.testing.assert_allclose(result.normal, np.array([0, 1, 0]), atol=1e-6)
    assert abs(result.d - 1.0) < 1e-6


def test_compute_floor_correction_identity():
    # Current floor is already at target
    # Target y = 0.0
    # Current plane: normal [0, 1, 0], d = 0.0 (y = 0)

    current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=0.0)

    result = compute_floor_correction(current_plane, target_floor_y=0.0)

    assert result.valid
    np.testing.assert_allclose(result.transform, np.eye(4), atol=1e-6)


def test_compute_floor_correction_translation_only():
    # Current floor is at y = -1.0
    # Plane eq: y + d = 0 => -1 + d = 0 => d = 1.0
    # Target y = 0.0
    # We need to move everything UP by 1.0 (Ty = 1.0)

    current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=1.0)

    result = compute_floor_correction(
        current_plane, target_floor_y=0.0, max_translation_m=2.0
    )

    assert result.valid
    expected = np.eye(4)
    expected[1, 3] = 1.0

    np.testing.assert_allclose(result.transform, expected, atol=1e-6)


def test_compute_floor_correction_rotation_only():
    # Current floor is tilted 45 deg around Z
    # Normal is [-0.707, 0.707, 0]
    # Target normal is [0, 1, 0]
    # We need to rotate -45 deg around Z to align normals

    angle = np.deg2rad(45)
    c, s = np.cos(angle), np.sin(angle)
    # Normal rotated by 45 deg around Z from [0, 1, 0]
    # Rz(45) @ [0, 1, 0] = [-s, c, 0] = [-0.707, 0.707, 0]
    normal = np.array([-s, c, 0])
    d = 0.0  # Passes through origin

    current_plane = FloorPlane(normal=normal, d=d)

    result = compute_floor_correction(
        current_plane, target_floor_y=0.0, max_rotation_deg=90.0
    )

    assert result.valid
    T_corr = result.transform

    # Check rotation part
    # Should be Rz(-45)
    angle_corr = np.deg2rad(-45)
    cc, ss = np.cos(angle_corr), np.sin(angle_corr)
    expected_R = np.array([[cc, -ss, 0], [ss, cc, 0], [0, 0, 1]])

    np.testing.assert_allclose(T_corr[:3, :3], expected_R, atol=1e-6)

    # Translation should be 0 since d=0 and we rotate around origin (roughly)
    assert np.linalg.norm(T_corr[:3, 3]) < 1e-6


def test_compute_floor_correction_bounds():
    # Request huge translation
    # Current floor y = -10.0 (d=10.0)
    # Target y = 0.0
    # Need Ty = 10.0
    # Max trans = 0.1

    current_plane = FloorPlane(normal=np.array([0, 1, 0]), d=10.0)

    result = compute_floor_correction(
        current_plane, target_floor_y=0.0, max_translation_m=0.1
    )

    assert not result.valid
    assert "exceeds limit" in result.reason