import numpy as np
import pytest
from aruco.alignment import (
    compute_face_normal,
    rotation_align_vectors,
    apply_alignment_to_pose,
    get_face_normal_from_geometry,
    detect_ground_face,
)


def test_compute_face_normal_valid_quad():
    # Define a quad in the XY plane (normal should be Z)
    corners = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], dtype=np.float64)

    # v1 = corners[1] - corners[0] = [1, 0, 0]
    # v2 = corners[3] - corners[0] = [0, 1, 0]
    # normal = [1, 0, 0] x [0, 1, 0] = [0, 0, 1]

    normal = compute_face_normal(corners)
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)


def test_compute_face_normal_valid_triangle():
    corners = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=np.float64)

    normal = compute_face_normal(corners)
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)


def test_compute_face_normal_degenerate():
    # Collinear points
    corners = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype=np.float64)
    with pytest.raises(ValueError, match="collinear or degenerate"):
        compute_face_normal(corners)

    # Too few points
    corners_few = np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float64)
    with pytest.raises(ValueError, match="At least 3 corners"):
        compute_face_normal(corners_few)


def test_rotation_align_vectors_identity():
    v1 = np.array([0, 0, 1], dtype=np.float64)
    v2 = np.array([0, 0, 1], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)
    np.testing.assert_allclose(R, np.eye(3), atol=1e-10)


def test_rotation_align_vectors_90_deg():
    v1 = np.array([1, 0, 0], dtype=np.float64)
    v2 = np.array([0, 1, 0], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)

    # Check that R @ v1 == v2
    np.testing.assert_allclose(R @ v1, v2, atol=1e-10)
    # Check orthogonality
    np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_rotation_align_vectors_antiparallel():
    v1 = np.array([0, 0, 1], dtype=np.float64)
    v2 = np.array([0, 0, -1], dtype=np.float64)
    R = rotation_align_vectors(v1, v2)

    np.testing.assert_allclose(R @ v1, v2, atol=1e-10)
    np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_apply_alignment_to_pose():
    # Identity pose
    T = np.eye(4, dtype=np.float64)
    # 90 deg rotation around Z
    R_align = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=np.float64)

    T_aligned = apply_alignment_to_pose(T, R_align)

    expected = np.eye(4)
    expected[:3, :3] = R_align
    np.testing.assert_allclose(T_aligned, expected, atol=1e-10)

    # Non-identity pose
    T_pose = np.eye(4)
    T_pose[:3, 3] = [1, 2, 3]
    T_aligned_pose = apply_alignment_to_pose(T_pose, R_align)

    # Pre-multiplication: R_align @ T_pose
    # Rotation part: R_align @ I = R_align
    # Translation part: R_align @ [1, 2, 3] = [-2, 1, 3]
    expected_pose = np.eye(4)
    expected_pose[:3, :3] = R_align
    expected_pose[:3, 3] = R_align @ [1, 2, 3]

    np.testing.assert_allclose(T_aligned_pose, expected_pose, atol=1e-10)


def test_get_face_normal_from_geometry():
    face_marker_map = {"top": [1, 2]}
    # Marker 1: XY plane (normal Z)
    # Marker 2: XY plane (normal Z)
    marker_geometry = {
        1: np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], dtype=np.float64),
        2: np.array([[2, 2, 0], [3, 2, 0], [3, 3, 0], [2, 3, 0]], dtype=np.float64),
    }

    normal = get_face_normal_from_geometry("top", marker_geometry, face_marker_map)
    assert normal is not None
    np.testing.assert_allclose(normal, np.array([0, 0, 1]), atol=1e-10)

    # Missing face map
    assert get_face_normal_from_geometry("top", marker_geometry, None) is None

    # Missing face name
    assert (
        get_face_normal_from_geometry("bottom", marker_geometry, face_marker_map)
        is None
    )

    # No visible markers for face
    assert (
        get_face_normal_from_geometry("top", {3: marker_geometry[1]}, face_marker_map)
        is None
    )


def test_detect_ground_face():
    face_marker_map = {
        "bottom": [1],
        "top": [2],
    }
    # Marker 1: normal [0, -1, 0] (aligned with camera up [0, -1, 0])
    # Marker 2: normal [0, 1, 0] (opposite to camera up)
    marker_geometry = {
        1: np.array(
            [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], dtype=np.float64
        ),  # Normal is [0, 1, 0] or [0, -1, 0]?
        # v1 = [1, 0, 0], v2 = [0, 0, 1] -> cross = [0, -1, 0]
        2: np.array([[0, 1, 0], [1, 1, 0], [1, 1, -1], [0, 1, -1]], dtype=np.float64),
        # v1 = [1, 0, 0], v2 = [0, 0, -1] -> cross = [0, 1, 0]
    }

    camera_up = np.array([0, -1, 0], dtype=np.float64)

    # Both visible
    res = detect_ground_face({1, 2}, marker_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res
    assert face_name == "bottom"
    np.testing.assert_allclose(normal, np.array([0, -1, 0]), atol=1e-10)

    # Case 1: We know about bottom, but only top is visible. Should pick bottom (best alignment).
    res = detect_ground_face({2}, marker_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res
    assert face_name == "bottom"
    np.testing.assert_allclose(normal, np.array([0, -1, 0]), atol=1e-10)

    # Case 2: We don't know about bottom (e.g. partial map). Should pick top (best available).
    partial_geometry = {2: marker_geometry[2]}
    res = detect_ground_face({2}, partial_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res
    assert face_name == "top"
    np.testing.assert_allclose(normal, np.array([0, 1, 0]), atol=1e-10)

    # None visible
    assert (
        detect_ground_face(set(), marker_geometry, camera_up, face_marker_map) is None
    )

    # Missing map
    assert detect_ground_face({1, 2}, marker_geometry, camera_up, None) is None


def test_detect_ground_face_geometric_priority():
    # Test that geometric alignment is preferred over semantic names
    # Scenario: 'bottom' face is tilted 45 deg, 'side' face is perfectly aligned with camera up
    # This simulates a box placed on its side

    face_marker_map = {
        "bottom": [1],
        "side": [2],
    }

    # Camera up is [0, -1, 0] (Y-down convention common in CV, or Y-up depending on setup)
    # Let's assume we want to align with [0, -1, 0]
    camera_up = np.array([0, -1, 0], dtype=np.float64)

    # Marker 1 (bottom): Tilted 45 deg. Normal = [0.707, -0.707, 0]
    # Dot product with [0, -1, 0] = 0.707
    marker_geometry = {
        1: np.array([[0, 0, 0], [1, 1, 0], [1, 1, 1], [0, 0, 1]], dtype=np.float64),
        # v1=[1,1,0], v2=[0,0,1] -> cross=[1, -1, 0] -> norm=[0.707, -0.707, 0]
        # Marker 2 (side): Perfectly aligned. Normal = [0, -1, 0]
        # Dot product with [0, -1, 0] = 1.0
        2: np.array([[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], dtype=np.float64),
        # v1=[1,0,0], v2=[0,0,1] -> cross=[0, -1, 0]
    }

    # OLD BEHAVIOR: would pick 'bottom' because of name
    # NEW BEHAVIOR: should pick 'side' because of better alignment score

    res = detect_ground_face({1, 2}, marker_geometry, camera_up, face_marker_map)
    assert res is not None
    face_name, normal = res

    # This assertion will fail until we fix the code
    assert face_name == "side"
    np.testing.assert_allclose(normal, np.array([0, -1, 0]), atol=1e-10)