"""Round-6 estimator provenance instrumentation.

Layered on Phases 3+4+7+8+9 (48 estimators). This round adds 6 more
spanning panel / decomposition / mediation / bartik / causal_impact.
Coverage 48/925 → **54/925**.

Estimators (6):
- ``sp.panel`` — multi-method panel dispatcher (refactored: outer
  wrapper + ``_dispatch_panel_impl``).
- ``sp.causal_impact`` — Brodersen et al. (2015) BSTS-style impact.
- ``sp.mediate`` — Imai-Keele-Tingley mediation.
- ``sp.mediate_interventional`` — VanderWeele-Vansteelandt-Robins (2014).
- ``sp.bartik`` — Goldsmith-Pinkham-Sorkin-Swift (2020) shift-share IV.
- ``sp.decompose`` — Oaxaca / FFL / DFL / RIF dispatcher.
"""
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(30):
        for t in range(8):
            rows.append({"i": u, "year": t,
                          "y": rng.normal() + 0.1 * t,
                          "x1": rng.normal()})
    return pd.DataFrame(rows)


@pytest.fixture
def ts_df():
    rng = np.random.default_rng(1)
    n = 100
    intervention = np.where(np.arange(n) >= 70, 2.0, 0.0)
    return pd.DataFrame({
        "y": rng.normal(size=n) + intervention,
        "t": range(n),
    })


@pytest.fixture
def mediation_df():
    rng = np.random.default_rng(2)
    n = 200
    return pd.DataFrame({
        "y": rng.normal(size=n),
        "d": rng.binomial(1, 0.5, size=n),
        "m": rng.normal(size=n),
        "x1": rng.normal(size=n),
    })


@pytest.fixture
def decomp_df():
    rng = np.random.default_rng(3)
    n = 200
    return pd.DataFrame({
        "log_wage": rng.normal(size=n) + 0.3 * rng.binomial(1, 0.5, size=n),
        "female": rng.binomial(1, 0.5, size=n),
        "edu": rng.normal(size=n),
    })


# ---------------------------------------------------------------------------
# Per-estimator
# ---------------------------------------------------------------------------

class TestPanelProvenance:
    def test_attached_fe(self, panel_df):
        r = sp.panel(panel_df, formula="y ~ x1",
                      entity="i", time="year", method="fe")
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.panel"
        assert prov.params["method"] == "fe"

    def test_method_choice_captured(self, panel_df):
        r = sp.panel(panel_df, formula="y ~ x1",
                      entity="i", time="year", method="re")
        prov = sp.get_provenance(r)
        assert prov.params["method"] == "re"


class TestCausalImpactProvenance:
    def test_attached(self, ts_df):
        r = sp.causal_impact(ts_df, y="y", time="t", intervention_time=70)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.causal_impact"
        assert prov.params["intervention_time"] == 70


class TestMediateProvenance:
    def test_attached(self, mediation_df):
        r = sp.mediate(mediation_df, y="y", treat="d", mediator="m",
                        covariates=["x1"], n_boot=20)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.mediate"
        assert prov.params["mediator"] == "m"


class TestMediateInterventionalProvenance:
    def test_attached(self, mediation_df):
        from statspai.mediation.mediate import mediate_interventional
        r = mediate_interventional(
            mediation_df, y="y", treat="d", mediator="m",
            covariates=["x1"], n_mc=50, n_boot=20,
        )
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.mediate_interventional"
        assert prov.params["n_mc"] == 50
        assert prov.params["n_boot"] == 20


class TestDecomposeProvenance:
    def test_attached_oaxaca(self, decomp_df):
        r = sp.decompose("oaxaca", data=decomp_df,
                          y="log_wage", group="female",
                          x=["edu"])
        prov = sp.get_provenance(r)
        assert prov is not None
        # Function name surfaces the dispatched method.
        assert prov.function == "sp.decompose.oaxaca"
        assert prov.params["method"] == "oaxaca"


class TestBartikProvenance:
    def test_attached(self):
        rng = np.random.default_rng(4)
        n_regions = 80
        K = 5
        # Region-level shares (rows sum to 1); national shocks per industry.
        shares = pd.DataFrame(
            rng.dirichlet(np.ones(K), size=n_regions),
            columns=[f"ind{i}" for i in range(K)],
        )
        shocks = pd.Series(rng.normal(size=K) * 0.05,
                            index=[f"ind{i}" for i in range(K)])
        # Build region-level data
        data = pd.DataFrame({
            "region": range(n_regions),
            "y": rng.normal(size=n_regions),
            "x": rng.normal(size=n_regions),
        })
        r = sp.bartik(data=data, y="y", endog="x",
                        shares=shares, shocks=shocks)
        prov = sp.get_provenance(r)
        assert prov is not None
        assert prov.function == "sp.bartik"


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

class TestRound6LineageIntegration:
    def test_multi_estimator_pack(self, panel_df, ts_df, tmp_path):
        r1 = sp.panel(panel_df, formula="y ~ x1",
                       entity="i", time="year", method="fe")
        r2 = sp.causal_impact(ts_df, y="y", time="t",
                                intervention_time=70)

        rp = sp.replication_pack(
            [r1, r2], tmp_path / "round6.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"))
            assert lin["n_runs"] >= 2
            funcs = {v["function"] for v in lin["runs"].values()}
            assert "sp.panel" in funcs
            assert "sp.causal_impact" in funcs
