import numpy as np
import pandas as pd

from irae_risk import reporting
from irae_risk.logistic import DEFAULT_RANDOM_SEED, LogisticPolyRegression


def _reporting_frames():
    df_data = pd.DataFrame(
        {
            "X": [-1.0, 0.0, 1.0, 2.0],
            "Y": [0, 0, 1, 1],
            "scale": ["plain"] * 4,
            "dataset": ["FULL"] * 4,
        }
    )
    df_fit = pd.DataFrame(
        {
            "scale": ["plain"],
            "dataset": ["FULL"],
            "theta": [np.array([-0.2, 0.8])],
            "cov": [np.eye(2) * 0.01],
        }
    )
    return df_data, df_fit


def test_reporting_builders_read_theta_column():
    df_data, df_fit = _reporting_frames()
    model = LogisticPolyRegression(degree=1)
    df_fit_index = df_fit.set_index(["scale", "dataset"])

    points, fits = reporting.build_plot_data_basic(
        df_data,
        df_fit_index,
        model,
        n_grid=5,
    )
    _, _, intervals = reporting.build_plot_data(
        df_data,
        df_fit_index,
        alpha=0.05,
        lg=model,
        n_grid=5,
    )

    assert len(points) == len(df_data)
    assert len(fits) == 5
    assert set(intervals["method"]) == {"normal", "delta"}


def test_full_ci_builder_accepts_bootstrap_theta_results():
    df_data, df_fit = _reporting_frames()
    model = LogisticPolyRegression(degree=1)
    df_fit_index = df_fit.set_index(["scale", "dataset"])
    bootstrap_theta = {
        ("plain", "FULL", "parametric_boots"): np.array(
            [[-0.3, 0.7], [-0.2, 0.8], [-0.1, 0.9]]
        )
    }

    _, _, intervals = reporting.build_plot_data_full_ci(
        df_data,
        df_fit_index,
        model,
        alpha=0.05,
        boots_theta_results=bootstrap_theta,
        n_grid=5,
        analytic_methods=(),
        bootstrap_methods=("parametric_boots",),
    )

    assert set(intervals["method"]) == {"parametric_boots"}
    assert set(intervals["ci_source"]) == {"bootstrap"}


def test_bootstrap_loader_calls_theta_api(tmp_path):
    df_data, _ = _reporting_frames()

    class ThetaSampler:
        def get_normal_theta(self, x, y, m, seed):
            assert len(x) == len(y) == 4
            assert seed == DEFAULT_RANDOM_SEED
            return np.full((m, 2), 1.5)

    result = reporting.load_or_build_bootstrap_theta(
        df_data,
        ThetaSampler(),
        tmp_path,
        n_boots=3,
        methods=("normal",),
    )

    samples = result[("plain", "FULL", "normal")]
    assert samples.shape == (3, 2)
    assert np.all(samples == 1.5)
    assert (tmp_path / "boots_theta_results.pkl").is_file()