import numpy as np
import pytest

from irae_risk.logistic import LogisticPolyRegression


@pytest.fixture
def regression_data():
    x = np.linspace(-2.0, 2.0, 20)
    y = np.array([0, 1] * 10)
    theta = np.array([-0.2, 0.8])
    return x, y, theta


def test_goodness_of_fit_uses_unpenalized_likelihood_by_default(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1, lam=(0.0, 0.5))

    result = model.goodness_of_fit(x, y, theta, bootstrap_samples=20)
    expected_llf = -model.get_nllf(x, y, theta)
    k = len(theta)
    n = len(x)

    assert result["LLF"] == pytest.approx(expected_llf)
    assert result["AIC"] == pytest.approx(2 * k - 2 * expected_llf)
    assert result["BIC"] == pytest.approx(k * np.log(n) - 2 * expected_llf)


def test_goodness_of_fit_can_include_regularization(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1, lam=(0.0, 0.5))

    unregularized = model.goodness_of_fit(
        x, y, theta, bootstrap_samples=20
    )
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
        bootstrap_samples=20,
    )
    penalty = model.penalty(theta)

    assert regularized["LLF"] == pytest.approx(
        -model.get_cost(x, y, theta)
    )
    assert unregularized["LLF"] - regularized["LLF"] == pytest.approx(penalty)
    assert regularized["AIC"] - unregularized["AIC"] == pytest.approx(
        2 * penalty
    )
    assert regularized["BIC"] - unregularized["BIC"] == pytest.approx(
        2 * penalty
    )


def test_regularization_switch_has_no_effect_without_penalty(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1)

    default = model.goodness_of_fit(x, y, theta, bootstrap_samples=20)
    regularized = model.goodness_of_fit(
        x,
        y,
        theta,
        regularization=True,
        bootstrap_samples=20,
    )

    assert regularized == pytest.approx(default)


def test_goodness_of_fit_reports_reproducible_bootstrap_deviance(regression_data):
    x, y, theta = regression_data
    model = LogisticPolyRegression(degree=1)

    first = model.goodness_of_fit(
        x,
        y,
        theta,
        bootstrap_samples=25,
        bootstrap_seed=123,
    )
    second = model.goodness_of_fit(
        x,
        y,
        theta,
        bootstrap_samples=25,
        bootstrap_seed=123,
    )

    assert "chi2" not in first
    assert "p-value(chi2)" not in first
    assert first["deviance"] == pytest.approx(2 * model.get_nllf(x, y, theta))
    assert first["p-value(deviance_bootstrap)"] == second[
        "p-value(deviance_bootstrap)"
    ]
    assert 0 < first["p-value(deviance_bootstrap)"] <= 1
    assert first["deviance_bootstrap_samples"] == 25