"""Round-9 provenance: survival + remaining IV variants (66 → 76/925)."""
from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

import statspai as sp


@pytest.fixture
def surv_df():
    rng = np.random.default_rng(0)
    n = 200
    return pd.DataFrame({
        "t": rng.exponential(2.0, size=n),
        "e": rng.binomial(1, 0.7, size=n),
        "g": rng.binomial(1, 0.5, size=n),
        "x1": rng.normal(size=n),
        "cluster": rng.integers(0, 5, size=n),
    })


@pytest.fixture
def iv_df():
    rng = np.random.default_rng(1)
    n = 200
    z = rng.normal(size=n)
    x = z + rng.normal(size=n)
    return pd.DataFrame({
        "y": x + rng.normal(size=n), "x": x, "z": z,
        "z2": rng.normal(size=n),
    })


# --- Survival ---------------------------------------------------------

class TestKaplanMeierProvenance:
    def test_attached(self, surv_df):
        r = sp.kaplan_meier(surv_df, duration="t", event="e", group="g")
        assert sp.get_provenance(r).function == "sp.survival.kaplan_meier"


class TestCoxProvenance:
    def test_attached(self, surv_df):
        r = sp.cox(formula="t ~ x1", data=surv_df,
                    duration="t", event="e")
        assert sp.get_provenance(r).function == "sp.survival.cox"


class TestAftProvenance:
    def test_attached(self, surv_df):
        from statspai.survival.aft import aft
        r = aft("t + e ~ x1", data=surv_df, family="weibull")
        assert sp.get_provenance(r).function == "sp.survival.aft"


class TestCoxFrailtyProvenance:
    def test_attached(self, surv_df):
        from statspai.survival.frailty import cox_frailty
        r = cox_frailty("t + e ~ x1", data=surv_df,
                          cluster="cluster")
        assert sp.get_provenance(r).function == "sp.survival.cox_frailty"


class TestCausalSurvivalForestProvenance:
    def test_attached(self, surv_df):
        from statspai.survival.causal_forest import causal_survival_forest
        df = surv_df.assign(d=surv_df["g"])  # treat = g
        r = causal_survival_forest(
            df, time="t", event="e", treat="d",
            covariates=["x1"], n_trees=20,
        )
        assert sp.get_provenance(r).function == "sp.survival.causal_survival_forest"


# --- IV variants ------------------------------------------------------

class TestKernelIvProvenance:
    def test_attached(self, iv_df):
        from statspai.iv.kernel_iv import kernel_iv
        r = kernel_iv(data=iv_df, y="y", treat="x",
                        instrument="z", n_boot=20)
        assert sp.get_provenance(r).function == "sp.iv.kernel_iv"


class TestNpivProvenance:
    def test_attached(self, iv_df):
        from statspai.iv.npiv import npiv
        r = npiv(y="y", endog="x",
                  instruments=iv_df[["z"]], data=iv_df,
                  k_d=3, k_z=3)
        assert sp.get_provenance(r).function == "sp.iv.npiv"


class TestManyWeakJiveProvenance:
    def test_attached(self, iv_df):
        from statspai.iv.many_weak import jive as mw_jive
        r = mw_jive(data=iv_df, y="y", endog="x", instruments=["z", "z2"])
        assert sp.get_provenance(r).function == "sp.iv.many_weak_jive"


class TestManyWeakArProvenance:
    def test_attached(self, iv_df):
        from statspai.iv.many_weak import many_weak_ar
        r = many_weak_ar(
            data=iv_df, y="y", endog="x", instruments=["z", "z2"],
            beta_grid=np.linspace(-2, 2, 21),
        )
        assert sp.get_provenance(r).function == "sp.iv.many_weak_ar"


class TestContinuousIvLateProvenance:
    def test_attached(self, iv_df):
        from statspai.iv.continuous_late import continuous_iv_late
        r = continuous_iv_late(
            data=iv_df, y="y", treat="x", instrument="z",
            n_quantiles=4, n_boot=20,
        )
        assert sp.get_provenance(r).function == "sp.iv.continuous_iv_late"
