import numpy as np
import pytest
from aruco.pose_averaging import PoseAccumulator
from aruco.pose_math import rvec_tvec_to_matrix


def test_pose_accumulator_basic():
    acc = PoseAccumulator()
    T = np.eye(4)
    acc.add_pose(T, 0.1, 0)

    assert len(acc.poses) == 1
    assert acc.reproj_errors[0] == 0.1
    assert acc.frame_ids[0] == 0


def test_filter_by_reproj():
    acc = PoseAccumulator()
    acc.add_pose(np.eye(4), 0.1, 0)
    acc.add_pose(np.eye(4), 0.5, 1)
    acc.add_pose(np.eye(4), 1.0, 2)

    indices = acc.filter_by_reproj(0.6)
    assert indices == [0, 1]


def test_ransac_filter_translation_outlier():
    acc = PoseAccumulator()
    # Reference pose
    T_ref = np.eye(4)
    acc.add_pose(T_ref, 0.1, 0)

    # Inlier (small shift)
    T_inlier = np.eye(4)
    T_inlier[0, 3] = 0.01
    acc.add_pose(T_inlier, 0.1, 1)

    # Outlier (large shift)
    T_outlier = np.eye(4)
    T_outlier[0, 3] = 1.0
    acc.add_pose(T_outlier, 0.1, 2)

    inliers = acc.ransac_filter(trans_thresh_m=0.1)
    assert 0 in inliers
    assert 1 in inliers
    assert 2 not in inliers


def test_compute_robust_mean():
    acc = PoseAccumulator()
    # Two identical poses
    T = np.eye(4)
    T[0, 3] = 1.0
    acc.add_pose(T, 0.1, 0)
    acc.add_pose(T, 0.1, 1)

    T_mean, stats = acc.compute_robust_mean()

    np.testing.assert_allclose(T_mean, T, atol=1e-10)
    assert stats["n_inliers"] == 2
    assert stats["median_reproj_error"] == 0.1


def test_compute_robust_mean_with_outlier():
    acc = PoseAccumulator()
    T1 = np.eye(4)
    T1[0, 3] = 1.0

    T2 = np.eye(4)
    T2[0, 3] = 1.1

    T_outlier = np.eye(4)
    T_outlier[0, 3] = 10.0

    acc.add_pose(T1, 0.1, 0)
    acc.add_pose(T2, 0.2, 1)
    acc.add_pose(T_outlier, 0.5, 2)

    # Filter out the outlier manually or via RANSAC
    inliers = acc.ransac_filter(trans_thresh_m=0.5)
    T_mean, stats = acc.compute_robust_mean(inliers)

    assert stats["n_inliers"] == 2
    # Median of 1.0 and 1.1 is 1.05
    assert abs(T_mean[0, 3] - 1.05) < 1e-10