"""
Tests for the core functionality of the ExactCIs package.
"""

import pytest
import math
import time
from exactcis.core import (
    validate_counts, support, pmf, pmf_weights, find_root, find_smallest_theta,
    logsumexp, apply_haldane_correction, log_binom_coeff, log_nchg_pmf, 
    log_nchg_cdf, log_nchg_sf, find_root_log, find_plateau_edge
)
import numpy as np


def test_validate_counts_valid():
    """Test that validate_counts accepts valid inputs."""
    validate_counts(1, 2, 3, 4)  # Should not raise


def test_validate_counts_negative():
    """Test that validate_counts rejects negative counts."""
    with pytest.raises(ValueError):
        validate_counts(-1, 2, 3, 4)


def test_validate_counts_empty_margin():
    """Test that validate_counts rejects tables with empty margins."""
    with pytest.raises(ValueError):
        validate_counts(0, 0, 3, 4)  # Empty row
    with pytest.raises(ValueError):
        validate_counts(1, 2, 0, 0)  # Empty row
    with pytest.raises(ValueError):
        validate_counts(0, 2, 0, 4)  # Empty column
    with pytest.raises(ValueError):
        validate_counts(1, 0, 3, 0)  # Empty column


def test_support():
    """Test the support function."""
    # Test that support returns the correct support array for different inputs
    assert np.array_equal(support(5, 5, 5).x, np.array([0, 1, 2, 3, 4, 5]))
    assert np.array_equal(support(3, 2, 2).x, np.array([0, 1, 2]))
    assert np.array_equal(support(2, 3, 2).x, np.array([0, 1, 2]))
    assert np.array_equal(support(2, 2, 3).x, np.array([1, 2]))


def test_pmf_sums_to_one():
    """Test that the PMF sums to 1."""
    n1, n2, m = 5, 5, 5
    theta = 2.0
    supp = support(n1, n2, m)
    total = sum(pmf(k, n1, n2, m, theta) for k in supp.x)
    assert abs(total - 1.0) < 1e-10


def test_pmf_weights():
    """Test the pmf_weights function."""
    # Test with simple parameters
    n1, n2, m = 5, 5, 5
    theta = 2.0
    supp, probs = pmf_weights(n1, n2, m, theta)

    # Check that support is correct
    assert np.array_equal(supp, np.array([0, 1, 2, 3, 4, 5]))

    # Check that probabilities sum to 1
    assert abs(sum(probs) - 1.0) < 1e-10

    # Check that probabilities are non-negative
    assert all(p >= 0 for p in probs)

    # Check that the length of support and probabilities match
    assert len(supp) == len(probs)

    # Test with theta = 0
    supp, probs = pmf_weights(5, 5, 5, 0.0)
    assert probs[0] == 1.0 and all(p == 0.0 for p in probs[1:])

    # Test that the pmf_weights match the pmf function
    n1, n2, m = 4, 6, 3
    theta = 1.5
    supp, probs = pmf_weights(n1, n2, m, theta)
    for i, k in enumerate(supp):
        assert abs(probs[i] - pmf(k, n1, n2, m, theta)) < 1e-10


def test_find_root():
    """Test the find_root function."""
    # Test with a simple function: f(x) = x^2 - 4
    f = lambda x: x**2 - 4
    root = find_root(f, 0, 3)
    assert abs(root - 2.0) < 1e-8


def test_find_smallest_theta():
    """Test the find_smallest_theta function."""
    import logging
    logger = logging.getLogger(__name__)
    
    # Create a simple test function that returns p-value = alpha when theta = 2.0
    def test_func(theta):
        return 0.025 if theta >= 2.0 else 0.05
    
    # Test with two-sided=True (default)
    # Since test_func(1.0) = 0.05 = alpha, and increasing=False (default), 
    # the function should return lo=1.0 as the smallest theta where f(theta) >= alpha
    theta = find_smallest_theta(test_func, alpha=0.05, lo=1.0, hi=3.0)
    logger.info(f"Two-sided=True result: {theta:.6f}")
    # The function correctly returns 1.0 since that's the smallest theta where f(theta) >= alpha
    assert abs(theta - 1.0) < 1e-6, f"Expected theta = 1.0, got {theta}"
    
    # Test with two-sided=False
    theta = find_smallest_theta(test_func, alpha=0.05, lo=1.0, hi=3.0, two_sided=False)
    logger.info(f"Two-sided=False result: {theta:.6f}")
    # For two_sided=False, same logic applies - returns 1.0 as smallest theta where f(theta) >= alpha
    assert abs(theta - 1.0) < 1e-6, f"Expected theta = 1.0, got {theta}"
    
    # Test with a continuous function
    def continuous_func(theta):
        return 0.05 * (theta / 2.0)
    
    theta = find_smallest_theta(continuous_func, alpha=0.05, lo=1.0, hi=3.0)
    logger.info(f"Continuous function result: {theta:.6f}")
    # Match the actual implementation behavior
    assert 1.9 < theta < 2.1, f"Expected theta near 2.0, got {theta}"


# ============================================================================
# PHASE 1 CRITICAL TESTS - HIGH PRIORITY FUNCTIONS
# ============================================================================

@pytest.mark.core
@pytest.mark.fast
def test_logsumexp_basic():
    """Test basic logsumexp functionality."""
    # Test basic functionality
    assert abs(logsumexp([0, 0]) - math.log(2)) < 1e-10
    
    # Test single element
    assert abs(logsumexp([1.5]) - 1.5) < 1e-10
    
    # Test empty list
    assert logsumexp([]) == float('-inf')


@pytest.mark.core
@pytest.mark.fast
def test_logsumexp_numerical_stability():
    """Test logsumexp numerical stability with large values."""
    # Test with large values (numerical stability)
    large_vals = [700, 701, 702]
    result = logsumexp(large_vals)
    assert not math.isinf(result)
    assert not math.isnan(result)
    
    # Test with very negative values
    small_vals = [-700, -701, -702]
    result = logsumexp(small_vals)
    assert abs(result - (-700 + math.log(1 + math.exp(-1) + math.exp(-2)))) < 1e-10


@pytest.mark.core
@pytest.mark.fast
def test_logsumexp_infinity_handling():
    """Test logsumexp with infinity values."""
    # Test with -inf values
    assert logsumexp([float('-inf'), 0]) == 0
    
    # Test with all -inf values
    assert logsumexp([float('-inf'), float('-inf')]) == float('-inf')
    
    # Test with mixed -inf and finite values
    result = logsumexp([float('-inf'), 1.0, 2.0])
    expected = logsumexp([1.0, 2.0])
    assert abs(result - expected) < 1e-10


@pytest.mark.core
@pytest.mark.fast
def test_apply_haldane_correction():
    """Test Haldane correction application."""
    # Test with zeros - should add 0.5 to all cells
    result = apply_haldane_correction(0, 5, 8, 10)
    assert result == (0.5, 5.5, 8.5, 10.5)
    
    # Test without zeros - should return unchanged
    result = apply_haldane_correction(1, 5, 8, 10)
    assert result == (1.0, 5.0, 8.0, 10.0)
    
    # Test with multiple zeros
    result = apply_haldane_correction(0, 0, 8, 10)
    assert result == (0.5, 0.5, 8.5, 10.5)
    
    # Test with all zeros
    result = apply_haldane_correction(0, 0, 0, 0)
    assert result == (0.5, 0.5, 0.5, 0.5)


@pytest.mark.core
@pytest.mark.fast
def test_log_binom_coeff():
    """Test log binomial coefficient calculation."""
    # Test basic cases
    assert abs(log_binom_coeff(5, 2) - math.log(10)) < 1e-10  # C(5,2) = 10
    assert abs(log_binom_coeff(10, 3) - math.log(120)) < 1e-10  # C(10,3) = 120
    
    # Test edge cases
    assert log_binom_coeff(5, 0) == 0  # C(5,0) = 1, log(1) = 0
    assert log_binom_coeff(5, 5) == 0  # C(5,5) = 1, log(1) = 0
    
    # Test symmetry: C(n,k) = C(n,n-k)
    assert abs(log_binom_coeff(10, 3) - log_binom_coeff(10, 7)) < 1e-10
    
    # Test invalid cases
    assert log_binom_coeff(5, 6) == float('-inf')  # k > n
    assert log_binom_coeff(5, -1) == float('-inf')  # k < 0


@pytest.mark.core
@pytest.mark.fast
def test_log_nchg_pmf_basic():
    """Test log noncentral hypergeometric PMF basic functionality."""
    n1, n2, m1, theta = 5, 5, 5, 2.0
    supp = support(n1, n2, m1)
    
    # Test against regular pmf
    for k in supp.x:
        log_prob = log_nchg_pmf(k, n1, n2, m1, theta)
        regular_prob = pmf(k, n1, n2, m1, theta)
        if regular_prob > 0:  # Only compare when regular_prob is positive
            assert abs(math.exp(log_prob) - regular_prob) < 1e-10


@pytest.mark.core
@pytest.mark.fast
def test_log_nchg_pmf_extreme_values():
    """Test log PMF with extreme theta values."""
    n1, n2, m1 = 5, 5, 5
    supp = support(n1, n2, m1)
    
    # Test with theta = 0
    log_prob_zero = log_nchg_pmf(supp.min_val, n1, n2, m1, 0.0)
    assert log_prob_zero == 0.0  # Should be log(1) for theta=0
    
    # Other values should be -inf for theta=0
    for k in supp.x:
        if k != supp.min_val:
            assert log_nchg_pmf(k, n1, n2, m1, 0.0) == float('-inf')


@pytest.mark.core
@pytest.mark.fast
def test_log_nchg_pmf_out_of_support():
    """Test log PMF outside support."""
    n1, n2, m1, theta = 5, 5, 5, 2.0
    
    # Test out of support values
    assert log_nchg_pmf(-1, n1, n2, m1, theta) == float('-inf')
    assert log_nchg_pmf(10, n1, n2, m1, theta) == float('-inf')


@pytest.mark.core
@pytest.mark.fast
def test_log_nchg_cdf_basic():
    """Test log CDF basic functionality."""
    n1, n2, m1, theta = 5, 5, 5, 2.0
    supp = support(n1, n2, m1)
    
    # Test boundary conditions
    assert log_nchg_cdf(supp.min_val - 1, n1, n2, m1, theta) == float('-inf')
    assert log_nchg_cdf(supp.max_val, n1, n2, m1, theta) == 0.0  # log(1)
    
    # Test monotonicity: CDF should be increasing
    prev_val = float('-inf')
    for k in supp.x:
        curr_val = log_nchg_cdf(k, n1, n2, m1, theta)
        assert curr_val >= prev_val
        prev_val = curr_val


@pytest.mark.core
@pytest.mark.fast
def test_log_nchg_sf_basic():
    """Test log survival function basic functionality."""
    n1, n2, m1, theta = 5, 5, 5, 2.0
    supp = support(n1, n2, m1)
    
    # Test boundary conditions
    assert log_nchg_sf(supp.max_val, n1, n2, m1, theta) == float('-inf')  # log(0)
    assert log_nchg_sf(supp.min_val - 1, n1, n2, m1, theta) == 0.0  # log(1)
    
    # Test complementarity: log(CDF(k-1)) + log(SF(k-1)) should be close to log(1) = 0
    # Note: This is approximate due to numerical precision
    for k in supp.x[1:]:  # Skip first element
        log_cdf = log_nchg_cdf(k-1, n1, n2, m1, theta)
        log_sf = log_nchg_sf(k-1, n1, n2, m1, theta)
        # Use logsumexp to compute log(exp(log_cdf) + exp(log_sf))
        log_sum = logsumexp([log_cdf, log_sf])
        assert abs(log_sum - 0.0) < 1e-8  # Should be close to log(1) = 0


@pytest.mark.core
@pytest.mark.fast
def test_find_root_log_basic():
    """Test find_root_log basic functionality."""
    # Test with a simple function: f(x) = x - 2.0
    f = lambda x: x - 2.0
    result = find_root_log(f, lo=1.0, hi=4.0)
    assert result is not None
    assert abs(math.exp(result) - 2.0) < 1e-6


@pytest.mark.core
@pytest.mark.fast
def test_find_root_log_timeout():
    """Test find_root_log timeout functionality."""
    def slow_func(x):
        time.sleep(0.01)  # Small delay
        return x - 2.0
    
    # Create a timeout checker that always returns True (timeout)
    timeout_checker = lambda: True
    result = find_root_log(slow_func, lo=1.0, hi=4.0, timeout_checker=timeout_checker)
    assert result is None


@pytest.mark.core
@pytest.mark.fast
def test_find_root_log_invalid_bounds():
    """Test find_root_log with invalid bounds."""
    f = lambda x: x - 2.0
    
    # Test with negative bounds
    with pytest.raises(ValueError):
        find_root_log(f, lo=-1.0, hi=4.0)
    
    with pytest.raises(ValueError):
        find_root_log(f, lo=1.0, hi=-4.0)


@pytest.mark.core
@pytest.mark.fast
def test_find_root_log_unbracketable():
    """Test find_root_log with unbracketable root."""
    # Function that doesn't change sign in the given interval
    f_no_root = lambda x: x + 1  # Always positive
    
    # This should attempt to expand the search interval and eventually raise RuntimeError
    with pytest.raises(RuntimeError):
        find_root_log(f_no_root, lo=1.0, hi=2.0)


@pytest.mark.core
@pytest.mark.fast
def test_find_plateau_edge_basic():
    """Test find_plateau_edge basic functionality."""
    # Test with plateau function where we need to search for the edge
    def plateau_func(x):
        if x < 2.0:
            return 0.01  # Below target
        elif x < 4.0:
            return 0.05  # At target (plateau)
        else:
            return 0.1   # Above target
    
    result = find_plateau_edge(plateau_func, 1.0, 5.0, target=0.05)
    assert result is not None
    theta, iterations = result
    assert 1.9 < theta < 2.1  # Should find left edge of plateau


@pytest.mark.core
@pytest.mark.fast
def test_find_plateau_edge_timeout():
    """Test find_plateau_edge timeout functionality."""
    def slow_func(x):
        time.sleep(0.01)
        return 0.05
    
    # Timeout immediately
    timeout_checker = lambda: True
    result = find_plateau_edge(slow_func, 1.0, 5.0, target=0.05, timeout_checker=timeout_checker)
    assert result is None


@pytest.mark.core
@pytest.mark.fast
def test_find_plateau_edge_boundary_conditions():
    """Test find_plateau_edge boundary conditions."""
    # Function where lo already meets condition
    def func_lo_good(x):
        return 0.06  # Always above target
    
    result = find_plateau_edge(func_lo_good, 1.0, 5.0, target=0.05, increasing=True)
    assert result is not None
    theta, iterations = result
    assert theta == 1.0
    assert iterations == 0
    
    # Function where hi doesn't meet condition
    def func_hi_bad(x):
        return 0.04  # Always below target
    
    result = find_plateau_edge(func_hi_bad, 1.0, 5.0, target=0.05, increasing=True)
    assert result is not None
    theta, iterations = result
    assert theta == 5.0
    assert iterations == 0
