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


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

    design = np.column_stack([np.ones_like(x), x])
    probabilities = model.model(x, theta)
    weights = probabilities * (1 - probabilities)
    information = (design.T * weights) @ design
    bread = information + 2 * ridge * np.eye(2)
    bread_inv = np.linalg.pinv(bread, hermitian=True)

    scores = (y - probabilities)[:, None] * design
    robust_meat = scores.T @ scores

    expected_model = bread_inv @ information @ bread_inv
    expected_robust = bread_inv @ robust_meat @ bread_inv

    np.testing.assert_allclose(
        model.get_cov(x, y, theta),
        expected_model,
    )
    np.testing.assert_allclose(
        model.get_cov(x, y, theta, method="robust_sandwich"),
        expected_robust,
    )
    np.testing.assert_allclose(
        model.get_cov(x, y, theta, method="inverse_hessian"),
        bread_inv,
    )


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

    with pytest.raises(ValueError, match="Unknown covariance method"):
        model.get_cov(x, y, theta, method="not-a-method")

    l1_model = LogisticPolyRegression(degree=1, lam=(0.1, 0.0))
    with pytest.raises(ValueError, match="nonzero L1"):
        l1_model.get_cov(x, y, theta)