import pytest
import numpy as np
from intreg import IntReg
from scipy.optimize import OptimizeResult
from scipy.stats import norm


@pytest.fixture
def simple_point_data():
    """Fixture for simple point data (uncensored)."""
    y_lower = np.array([1, 1, 2])
    y_upper = np.array([1, 1, 2])
    return y_lower, y_upper


@pytest.fixture
def left_censored_data():
    """Fixture for left-censored data."""
    y_lower = np.array([-np.inf, 0, 1])
    y_upper = np.array([1, 1, 2])
    return y_lower, y_upper


@pytest.fixture
def right_censored_data():
    """Fixture for right-censored data."""
    y_lower = np.array([1, 1, 2,])
    y_upper = np.array([np.inf, 2, 3])
    return y_lower, y_upper


@pytest.fixture
def mixed_interval_and_point_data():
    """Fixture for mixed interval and point data."""
    y_lower = np.array([1, 2, 3])
    y_upper = np.array([1, 3, 4])
    return y_lower, y_upper


@pytest.fixture
def clear_interval_censored_data():
    """Fixture for clear interval-censored data with obvi ous midpoints."""
    y_lower = np.array([1, 2, 3, 4])
    y_upper = np.array([2, 3, 4, 5])
    return y_lower, y_upper


def test_initial_params(
    simple_point_data,
    left_censored_data,
    right_censored_data,
    clear_interval_censored_data,
):
    """Test _initial_params method with all types of mock data (point, left-censored, right-censored, and interval-censored)."""

    # Define test cases as tuples: (data, expected_mu, expected_log_sigma_func)
    # We define expected values based on the interpretation of each dataset.
    test_cases = [
        (
            simple_point_data,
            lambda midpoints: np.nanmean(midpoints),
            lambda midpoints: np.log(np.nanstd(midpoints)),
        ),
        (
            left_censored_data,
            lambda midpoints: np.nanmean(midpoints),
            lambda midpoints: np.log(np.nanstd(midpoints)),
        ),
        (
            right_censored_data,
            lambda midpoints: np.nanmean(midpoints),
            lambda midpoints: np.log(np.nanstd(midpoints)),
        ),
        (
            clear_interval_censored_data,
            lambda midpoints: np.nanmean(midpoints),
            lambda midpoints: np.log(np.nanstd(midpoints)),
        ),
    ]

    for (y_lower, y_upper), expected_mu_func, expected_log_sigma_func in test_cases:
        model = IntReg(y_lower, y_upper)

        # Calculate midpoints for uncensored data
        midpoints = (y_lower + y_upper) / 2.0
        midpoints = np.where(np.isfinite(midpoints), midpoints, np.nan)

        # Calculate expected values
        expected_mu = expected_mu_func(midpoints)
        expected_log_sigma = expected_log_sigma_func(midpoints)

        # Call _initial_params and get the output
        initial_mu, initial_log_sigma = model._initial_params()

        # Assertions to check that _initial_params matches expected mu and log(sigma)
        assert np.isclose(
            initial_mu, expected_mu, atol=1e-4
        ), f"Failed for data {y_lower, y_upper}: expected mu {expected_mu}, got {initial_mu}"
        assert np.isclose(
            initial_log_sigma, expected_log_sigma, atol=1e-4
        ), f"Failed for data {y_lower, y_upper}: expected log(sigma) {expected_log_sigma}, got {initial_log_sigma}"


def test_log_L(
    simple_point_data,
    left_censored_data,
    right_censored_data,
    clear_interval_censored_data,
):
    """Test log_L using parameters generated by _initial_params on all data types."""

    # Define test cases as tuples: (data, description)
    test_cases = [
        (simple_point_data, "point data"),
        (right_censored_data, "right-censored data"),
        (left_censored_data, "left-censored data"),
        (clear_interval_censored_data, "interval-censored data"),
    ]

    for (y_lower, y_upper), description in test_cases:
        model = IntReg(y_lower, y_upper)

        params = model._initial_params()
        print ('params', params)
        log_likelihood = model.log_L(params)

        assert isinstance(
            log_likelihood, float
        ), f"log_L output should be float for {description}"
        assert np.isfinite(
            log_likelihood
        ), f"log_L output should be finite for {description}"

        print(log_likelihood)
        assert (
            log_likelihood > 0
        ), f"log_L output should be positive for {description}, got {log_likelihood}"


def test_simple_point_data_fit(simple_point_data):
    y_lower, y_upper = simple_point_data
    model = IntReg(y_lower, y_upper)
    result = model.fit()
    print (result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 1.3, atol=1e-1)  # Expect mu near 1
    assert np.isclose(result.x[1], -0.8, atol=1e-1)  # Expect sigma near -0.75


def test_left_censored_data_fit(left_censored_data):
    y_lower, y_upper = left_censored_data
    model = IntReg(y_lower, y_upper)
    result = model.fit()
    print (result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 0.9, atol=1e-1)  # Expect mu near 0.9
    assert np.isclose(result.x[1], -1.8, atol=1e-1)  # Expect sigma near -1.8


def test_right_censored_data_fit(right_censored_data):
    y_lower, y_upper = right_censored_data
    model = IntReg(y_lower, y_upper)
    result = model.fit()
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 2, atol=1e-1)  # Expect mu near 1.5
    assert np.isclose(result.x[1], -1.9, atol=1e-1)  # Expect sigma near -1.9


def test_mixed_interval_and_point_data_fit(mixed_interval_and_point_data):
    y_lower, y_upper = mixed_interval_and_point_data
    model = IntReg(y_lower, y_upper)
    result = model.fit()
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 2.3, atol=1e-1)  # Expect mu around 2.5
    assert np.isclose(result.x[1], 0.02, atol=1e-1)  # Expect log(sigma) = 0.02


def test_clear_interval_censored_data_fit(clear_interval_censored_data):
    y_lower, y_upper = clear_interval_censored_data
    model = IntReg(y_lower, y_upper)
    result = model.fit()
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 3, atol=1e-1)  # Expect mu around 2.5
    assert np.isclose(result.x[1], 0.08, atol=1e-1)  # Expect log(sigma) = 0.08
