"""Round-7 estimator provenance instrumentation.

Layered on Phases 3+4+7+8+9+10 (54 estimators). This round adds 7 more
spanning spatial econometrics / quantile / distributional / conformal /
bootstrap inference. Coverage 54/925 → **61/925**.

Estimators (7):
- ``sp.spatial.spatial_did`` — spatial-lag DiD with spillover decomposition.
- ``sp.spatial.spatial_iv`` — spatial 2SLS.
- ``sp.qte.dist_iv`` — distributional IV / quantile LATE.
- ``sp.qte.beyond_average_late`` — quantile LATE with imperfect compliance.
- ``sp.qte.qte_hd_panel`` — high-dim panel QTE via LASSO controls.
- ``sp.bootstrap`` — general-purpose bootstrap inference.
- ``sp.conformal_cate`` — conformal prediction intervals for CATE.
"""
from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

import statspai as sp


@pytest.fixture
def panel_df():
    rng = np.random.default_rng(0)
    rows = []
    for u in range(20):
        for t in range(4):
            rows.append({
                "i": u, "year": t,
                "y": rng.normal(),
                "d": int(t >= 2 and u < 10),
                "x1": rng.normal(),
            })
    return pd.DataFrame(rows)


@pytest.fixture
def W_matrix():
    n_u = 20
    W = np.zeros((n_u, n_u))
    for i in range(n_u - 1):
        W[i, i + 1] = 0.5
        W[i + 1, i] = 0.5
    rs = W.sum(axis=1, keepdims=True)
    rs[rs == 0] = 1
    return W / rs


@pytest.fixture
def iv_df():
    rng = np.random.default_rng(1)
    n = 200
    z = rng.normal(size=n)
    d = (z + rng.normal(size=n) > 0).astype(int)
    return pd.DataFrame({
        "y": 0.5 * d + rng.normal(size=n),
        "d": d,
        "z": z,
        "x1": rng.normal(size=n),
    })


# ---------------------------------------------------------------------------
# Spatial
# ---------------------------------------------------------------------------

class TestSpatialDidProvenance:
    def test_attached(self, panel_df, W_matrix):
        r = sp.spatial_did(panel_df, y="y", treat="d",
                            unit="i", time="year", W=W_matrix)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.spatial.spatial_did"


class TestSpatialIvProvenance:
    def test_attached(self, panel_df, W_matrix):
        # spatial_iv expects cross-section data shaped to W. Use first
        # period only as a 20-region cross-section.
        cs = panel_df[panel_df["year"] == 0].reset_index(drop=True)
        from statspai.spatial.iv import spatial_iv
        r = spatial_iv(cs, y="y", endog=["d"], exog=["x1"],
                        W=W_matrix, instruments=None)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.spatial.spatial_iv"


# ---------------------------------------------------------------------------
# Quantile / distributional
# ---------------------------------------------------------------------------

class TestDistIvProvenance:
    def test_attached(self, iv_df):
        from statspai.qte.dist_iv import dist_iv
        r = dist_iv(data=iv_df, y="y", treat="d", instrument="z",
                     n_boot=20)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.qte.dist_iv"


class TestBeyondAverageLateProvenance:
    def test_attached(self):
        # beyond_average_late requires a BINARY instrument.
        rng = np.random.default_rng(11)
        n = 200
        z_bin = rng.binomial(1, 0.5, size=n)
        # Fuzzy compliance: P(D=1|Z) higher when Z=1.
        d = ((rng.uniform(size=n) < 0.3 + 0.5 * z_bin)).astype(int)
        df = pd.DataFrame({
            "y": 0.5 * d + rng.normal(size=n),
            "d": d,
            "z": z_bin,
        })
        from statspai.qte.beyond_average import beyond_average_late
        r = beyond_average_late(data=df, y="y", treat="d",
                                  instrument="z", n_boot=20)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.qte.beyond_average_late"


class TestQteHdPanelProvenance:
    def test_attached(self, panel_df):
        from statspai.qte.hd_panel import qte_hd_panel
        r = qte_hd_panel(data=panel_df, y="y", treat="d",
                          unit="i", time="year", covariates=["x1"])
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.qte.qte_hd_panel"


# ---------------------------------------------------------------------------
# Bootstrap inference
# ---------------------------------------------------------------------------

class TestBootstrapProvenance:
    def test_attached(self):
        rng = np.random.default_rng(2)
        df = pd.DataFrame({"x": rng.normal(size=200)})
        r = sp.bootstrap(df,
                          statistic=lambda d: float(d["x"].mean()),
                          n_boot=50)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.bootstrap"
        assert prov.params["n_boot"] == 50


# ---------------------------------------------------------------------------
# Conformal CATE
# ---------------------------------------------------------------------------

class TestConformalCateProvenance:
    def test_attached(self):
        rng = np.random.default_rng(3)
        n = 250
        df = pd.DataFrame({
            "y": rng.normal(size=n),
            "d": rng.binomial(1, 0.5, size=n),
            "x1": rng.normal(size=n),
            "x2": rng.normal(size=n),
        })
        from statspai.conformal_causal.conformal_ite import conformal_cate
        r = conformal_cate(data=df, y="y", treat="d",
                            covariates=["x1", "x2"])
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.conformal_cate"


# ---------------------------------------------------------------------------
# Integration
# ---------------------------------------------------------------------------

class TestRound7LineageIntegration:
    def test_multi_estimator_pack(self, panel_df, W_matrix, tmp_path):
        rng = np.random.default_rng(4)
        df_boot = pd.DataFrame({"x": rng.normal(size=100)})

        r1 = sp.spatial_did(panel_df, y="y", treat="d",
                              unit="i", time="year", W=W_matrix)
        r2 = sp.bootstrap(df_boot,
                            statistic=lambda d: float(d["x"].mean()),
                            n_boot=20)

        rp = sp.replication_pack(
            [r1, r2], tmp_path / "round7.zip",
            data=panel_df, env=False,
        )
        import json
        import zipfile
        with zipfile.ZipFile(rp.output_path) as zf:
            assert "lineage.json" in zf.namelist()
            lin = json.loads(zf.read("lineage.json"))
            funcs = {v["function"] for v in lin["runs"].values()}
            assert "sp.spatial.spatial_did" in funcs
            assert "sp.bootstrap" in funcs
