"""Tests for quantile confidence interval using normal approximation to binomial."""

import numpy as np

from gri_utils.stats import quantile_ci

# Create a seeded random generator for reproducible tests
RNG = np.random.default_rng(42)


def test_median_ci_basic() -> None:
    """Test CI for median of normal data."""
    data = RNG.standard_normal(100)
    lower, upper = quantile_ci(data, quantile=0.5)

    assert isinstance(lower, float)
    assert isinstance(upper, float)
    assert lower < upper


def test_median_ci_contains_median() -> None:
    """Test that CI typically contains the sample median."""
    data = RNG.standard_normal(100)
    lower, upper = quantile_ci(data, quantile=0.5)
    median = np.median(data)

    assert lower < median < upper


def test_different_quantiles() -> None:
    """Test CI for different quantiles."""
    data = RNG.standard_normal(100)

    lower_25, upper_25 = quantile_ci(data, quantile=0.25)
    lower_50, upper_50 = quantile_ci(data, quantile=0.50)
    lower_75, upper_75 = quantile_ci(data, quantile=0.75)

    # Each CI should be valid
    assert lower_25 < upper_25
    assert lower_50 < upper_50
    assert lower_75 < upper_75

    # Lower quantiles should have lower bounds
    assert lower_25 < lower_50 < lower_75


def test_confidence_levels() -> None:
    """Test different confidence levels - wider CI for higher confidence."""
    data = RNG.standard_normal(100)

    lower_90, upper_90 = quantile_ci(data, quantile=0.5, confidence=0.90)
    lower_95, upper_95 = quantile_ci(data, quantile=0.5, confidence=0.95)
    lower_99, upper_99 = quantile_ci(data, quantile=0.5, confidence=0.99)

    # Higher confidence should give wider intervals
    width_90 = upper_90 - lower_90
    width_95 = upper_95 - lower_95
    width_99 = upper_99 - lower_99

    assert width_90 < width_95 < width_99


def test_insufficient_samples() -> None:
    """Test that insufficient samples returns NaN."""
    data = np.array([1, 2, 3, 4, 5])  # Only 5 samples, default min is 10
    lower, upper = quantile_ci(data)

    assert np.isnan(lower)
    assert np.isnan(upper)


def test_custom_min_samples() -> None:
    """Test custom minimum samples threshold."""
    data = np.array([1, 2, 3, 4, 5])

    # With min_samples=5, should work
    lower, upper = quantile_ci(data, min_samples=5)
    assert not np.isnan(lower)
    assert not np.isnan(upper)


def test_larger_sample_narrower_ci() -> None:
    """Test that larger samples give narrower CI."""
    data_small = RNG.standard_normal(50)
    data_large = RNG.standard_normal(500)

    lower_small, upper_small = quantile_ci(data_small)
    lower_large, upper_large = quantile_ci(data_large)

    width_small = upper_small - lower_small
    width_large = upper_large - lower_large

    assert width_large < width_small


def test_1d_array_input() -> None:
    """Test that 1D input returns scalar floats."""
    data = RNG.standard_normal(100)

    lower, upper = quantile_ci(data)
    assert isinstance(lower, float)
    assert isinstance(upper, float)
    assert lower < upper


# Multi-dimensional tests


def test_2d_array_axis0() -> None:
    """Test 2D array with axis=0 (samples in rows, variables in columns)."""
    data = RNG.standard_normal((100, 3))  # 100 samples, 3 variables

    lower, upper = quantile_ci(data, axis=0)

    # Output should be 1D with shape (3,)
    assert lower.shape == (3,)
    assert upper.shape == (3,)
    assert np.all(lower < upper)


def test_2d_array_axis1() -> None:
    """Test 2D array with axis=1 (variables in rows, samples in columns)."""
    data = RNG.standard_normal((3, 100))  # 3 variables, 100 samples each

    lower, upper = quantile_ci(data, axis=1)

    # Output should be 1D with shape (3,)
    assert lower.shape == (3,)
    assert upper.shape == (3,)
    assert np.all(lower < upper)


def test_2d_default_axis() -> None:
    """Test 2D array default axis behavior (axis=0)."""
    data = RNG.standard_normal((100, 3))

    lower_default, upper_default = quantile_ci(data)
    lower_axis0, upper_axis0 = quantile_ci(data, axis=0)

    np.testing.assert_array_equal(lower_default, lower_axis0)
    np.testing.assert_array_equal(upper_default, upper_axis0)


def test_row_vector_auto_axis() -> None:
    """Test that row vectors (1, n) automatically use axis=1."""
    data = RNG.standard_normal((1, 100))

    # Explicitly pass axis=1 to match expected behavior and satisfy type checker
    lower, upper = quantile_ci(data, axis=1)

    # Should return array output with shape (1,)
    assert lower.shape == (1,)
    assert upper.shape == (1,)
    assert lower[0] < upper[0]


def test_3d_array() -> None:
    """Test 3D array with explicit axis."""
    data = RNG.standard_normal((50, 3, 4))  # 50 samples along axis 0

    lower, upper = quantile_ci(data, axis=0)

    # Output should have shape (3, 4)
    assert lower.shape == (3, 4)
    assert upper.shape == (3, 4)
    assert np.all(lower < upper)


def test_3d_array_axis1() -> None:
    """Test 3D array with middle axis."""
    data = RNG.standard_normal((3, 50, 4))  # 50 samples along axis 1

    lower, upper = quantile_ci(data, axis=1)

    # Output should have shape (3, 4)
    assert lower.shape == (3, 4)
    assert upper.shape == (3, 4)
    assert np.all(lower < upper)


def test_negative_axis() -> None:
    """Test negative axis indexing."""
    data = RNG.standard_normal((100, 3))

    lower_neg, upper_neg = quantile_ci(data, axis=-1)
    lower_pos, upper_pos = quantile_ci(data, axis=1)

    # axis=-1 and axis=1 should be equivalent for 2D
    np.testing.assert_array_almost_equal(lower_neg, lower_pos)
    np.testing.assert_array_almost_equal(upper_neg, upper_pos)


def test_2d_insufficient_samples() -> None:
    """Test 2D array with insufficient samples returns NaN array."""
    data = RNG.standard_normal((5, 3))  # Only 5 samples

    lower, upper = quantile_ci(data, axis=0)

    assert lower.shape == (3,)
    assert upper.shape == (3,)
    assert np.all(np.isnan(lower))
    assert np.all(np.isnan(upper))


def test_2d_ci_contains_medians() -> None:
    """Test that 2D CIs contain sample medians."""
    data = RNG.standard_normal((100, 3))

    lower, upper = quantile_ci(data, axis=0)
    medians = np.median(data, axis=0)

    # Each CI should contain its corresponding median
    assert np.all(lower < medians)
    assert np.all(medians < upper)
