import pytest
import numpy as np
from intreg import MeIntReg
from scipy.optimize import OptimizeResult


@pytest.fixture
def simple_point_data():
    """Fixture for simple point data (uncensored) with covariates."""
    y_lower = np.array([1, 1, 2])
    y_upper = np.array([1, 1, 2])
    X = np.array([[1, 0], [1, 1], [1, 2]])  # Covariates with intercept and slope
    random_effects = np.array([0, 1, 1])
    return y_lower, y_upper, X, random_effects


@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])
    X = np.array([[1, 0], [1, 1], [1, 2]])
    random_effects = np.array([0, 1, 1])
    return y_lower, y_upper, X, random_effects


@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])
    X = np.array([[1, 0], [1, 1], [1, 2]])
    random_effects = np.array([0, 1, 1])
    return y_lower, y_upper, X, random_effects


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


@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])
    X = np.array([[1, 0], [1, 1], [1, 2], [1, 3]])
    random_effects = np.array([0, 1, 1, 2])
    return y_lower, y_upper, X, random_effects


def test_initial_params(
    simple_point_data,
    left_censored_data,
    right_censored_data,
    mixed_interval_and_point_data,
    clear_interval_censored_data,
):
    """Test _initial_params method for mixed-effects interval regression."""

    test_cases = [
        (simple_point_data, "simple point data"),
        (left_censored_data, "left censored data"),
        (right_censored_data, "right censored data"),
        (mixed_interval_and_point_data, "mixed inteval and point data"),
        (clear_interval_censored_data, "clear interval censored data"),
    ]

    for (y_lower, y_upper, X, random_effects), description in test_cases:
        model = MeIntReg(y_lower, y_upper, X, random_effects)

        initial_params = model._initial_params()
        beta_init = initial_params[: X.shape[1]]
        u_init = initial_params[X.shape[1] : X.shape[1] + model.n_random_effects]
        log_sigma_init = initial_params[-1]

        # Assert beta is computed via linear regression
        midpoints = (y_lower + y_upper) / 2.0
        valid_mask = np.isfinite(midpoints)
        X, midpoints = X[valid_mask], midpoints[valid_mask]
        expected_beta = np.linalg.lstsq(X, midpoints, rcond=None)[0]

        assert np.allclose(
            beta_init, expected_beta, atol=1e-4
        ), f"Failed for {description}: beta initialization incorrect"

        # Assert random effects are initialized as zeros
        assert np.allclose(
            u_init, 0, atol=1e-4
        ), f"Failed for {description}: random effects not initialized to zero"

        # Assert log_sigma is a log of the standard deviation of midpoints
        expected_log_sigma = np.log(np.nanstd(midpoints))
        assert np.isclose(
            log_sigma_init, expected_log_sigma, atol=1e-4
        ), f"Failed for {description}: log(sigma) initialization incorrect"


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, X, random_effects), description in test_cases:
        model = MeIntReg(y_lower, y_upper, X, random_effects)

        params = model._initial_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}"
        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, X, random_effects = simple_point_data
    model = MeIntReg(y_lower, y_upper, X, random_effects)
    result = model.fit()
    print(result)
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 0.6, atol=1e-1)  # Expect mu near 0.6
    assert np.isclose(result.x[1], 1, atol=1e-1)  # Expect sigma near 1


def test_left_censored_data_fit(left_censored_data):
    y_lower, y_upper, X, random_effects = left_censored_data
    model = MeIntReg(y_lower, y_upper, X, random_effects)
    result = model.fit()
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], -0.5, atol=1e-1)
    assert np.isclose(result.x[1], 1, atol=1e-1)

def test_right_censored_data_fit(right_censored_data):
    y_lower, y_upper, X, random_effects = right_censored_data
    model = MeIntReg(y_lower, y_upper, X, random_effects)
    result = model.fit()
    print (result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 1.2, atol=1e-1) 
    assert np.isclose(result.x[1], 1.1, atol=1e-1) 


def test_mixed_interval_and_point_data_fit(mixed_interval_and_point_data):
    y_lower, y_upper, X, random_effects = mixed_interval_and_point_data
    model = MeIntReg(y_lower, y_upper, X, random_effects)
    result = model.fit()
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 1.7, atol=1e-1)  
    assert np.isclose(result.x[1], 0.4, atol=1e-1) 

def test_clear_interval_censored_data_fit(clear_interval_censored_data):
    y_lower, y_upper, X, random_effects = clear_interval_censored_data
    model = MeIntReg(y_lower, y_upper, X, random_effects)
    result = model.fit()
    print(result.x[0], result.x[1])
    assert isinstance(result, OptimizeResult)
    assert np.isclose(result.x[0], 1.5, atol=1e-1)  
    assert np.isclose(result.x[1], 1, atol=1e-1)  