"""Monte Carlo coverage simulation for the cell-level wild PSU bootstrap
in ChaisemartinDHaultfoeuille (PR 4).

Validates empirical coverage of the bootstrap confidence intervals under
a DGP with PSU that varies across cells of the same group — the regime
unlocked by PR 4's replacement of the group-level PSU map with a
cell-level map. Under PSU-within-group-constant the dispatcher routes
through the legacy bootstrap (covered by the pre-PR-4 test suite), so
the coverage check here exercises only the new cell-level code path.

Asserts coverage at three surfaces, each covering a distinct code path:

1. Overall DID_M bootstrap CI (`res.bootstrap_results.overall_ci`)
   — single-target cell-level branch.
2. Event-study **horizon-1** CI (`res.bootstrap_results.event_study_cis[1]`)
   — first horizon of the shared-PSU-weight multi-horizon block.
3. Event-study **horizon-2** CI + sup-t `cband_crit_value` finiteness
   — exercises the cross-horizon shared-draw machinery that
   guarantees sup-t joint distribution validity. At L_max >= 2 the
   shared (n_bootstrap, n_psu) PSU-level weight matrix must be drawn
   ONCE and reused across horizons; a regression where each horizon
   re-draws weights would break the sup-t coherence and the finite
   critical value check below would surface it.

Marked ``slow`` and excluded from the default pytest run. To execute:

    pytest tests/test_dcdh_bootstrap_cell_period_coverage.py -m slow -v
"""

from __future__ import annotations

import warnings

import numpy as np
import pandas as pd
import pytest

from diff_diff.chaisemartin_dhaultfoeuille import ChaisemartinDHaultfoeuille
from diff_diff.survey import SurveyDesign


def _simulate_panel(
    n_groups: int,
    n_periods: int,
    first_treated_period: int,
    tau: float,
    rng: np.random.Generator,
    psu_sigma: float = 0.5,
    obs_sigma: float = 1.0,
    group_sigma: float = 1.0,
) -> pd.DataFrame:
    """Generate a one-obs-per-cell panel with within-group-varying PSU
    (PSU = period parity per group) so the cell-level bootstrap is
    exercised. Matches the DGP shape of `test_dcdh_cell_period_coverage.py`.
    """
    groups = np.arange(n_groups)
    treated = groups < n_groups // 2
    group_fe = rng.normal(0.0, group_sigma, size=n_groups)
    # (group, psu) effect — constant within (g, parity), drives
    # within-group residual correlation that Binder variance + wild
    # PSU bootstrap should capture.
    psu_fe = rng.normal(0.0, psu_sigma, size=(n_groups, 2))

    rows = []
    for g in groups:
        for t in range(n_periods):
            parity = int(t % 2)
            # Per-(group, parity) PSU codes: 2 * n_groups distinct PSUs
            # so Binder / bootstrap see cell granularity rather than
            # two global PSUs reused across groups.
            psu_id = int(g) * 2 + parity
            d = 1 if (treated[g] and t >= first_treated_period) else 0
            y = (
                group_fe[g]
                + 0.1 * t
                + tau * d
                + psu_fe[g, parity]
                + rng.normal(0.0, obs_sigma)
            )
            rows.append({
                "group": int(g),
                "period": int(t),
                "treatment": int(d),
                "outcome": float(y),
                "psu": psu_id,
                "pw": 1.0,
            })
    return pd.DataFrame(rows)


@pytest.mark.slow
def test_bootstrap_cell_period_coverage_varying_psu():
    """Empirical 95% coverage for bootstrap CIs at BOTH the overall
    DID_M and per-horizon event-study surfaces on a DGP with
    within-group-varying PSU. Tolerance ±2.5pp mirrors the analytical
    coverage test; per-horizon coverage guards the shared-PSU-weight
    machinery from CRITICAL #2 of the PR 4 plan review.
    """
    n_reps = 500
    n_groups = 40
    n_periods = 6
    first_treated_period = 3
    tau_true = 2.0

    rng = np.random.default_rng(20260419)
    covered_overall = 0
    covered_h1 = 0
    covered_h2 = 0
    cband_finite = 0
    failed = 0

    for r in range(n_reps):
        df = _simulate_panel(
            n_groups=n_groups,
            n_periods=n_periods,
            first_treated_period=first_treated_period,
            tau=tau_true,
            rng=rng,
        )
        sd = SurveyDesign(weights="pw", psu="psu")
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                # n_bootstrap=500 keeps internal bootstrap Monte-Carlo
                # noise well below the across-reps variance (at B=500
                # the percentile-CI endpoints are stable to ~0.3pp per
                # Efron-Tibshirani §13.3), so the across-reps coverage
                # mostly reflects the sampling-distribution / bootstrap-
                # consistency question rather than bootstrap MC noise.
                # L_max=2 exercises the shared-PSU-weight multi-horizon
                # block (a single `(n_bootstrap, n_psu)` weight matrix
                # is drawn once and broadcast per-horizon via each
                # horizon's cell-to-PSU map). L_max=1 would collapse to
                # a single target and never exercise the cross-horizon
                # shared-draw machinery.
                #
                # n_bootstrap=1000 keeps internal bootstrap MC noise
                # below ~0.3pp per CI endpoint; the percentile-CI
                # coverage at horizon-2 (where the shared-weight
                # broadcast is exercised) is finite-sample-sensitive
                # and B=500 would risk a spurious edge-of-band miss.
                res = ChaisemartinDHaultfoeuille(
                    n_bootstrap=1000, seed=r + 1,
                ).fit(
                    df,
                    outcome="outcome", group="group",
                    time="period", treatment="treatment",
                    survey_design=sd, L_max=2,
                )
        except Exception:
            failed += 1
            continue

        if res.bootstrap_results is None:
            failed += 1
            continue

        # Overall bootstrap CI.
        overall_ci = res.bootstrap_results.overall_ci
        if overall_ci is None or not all(np.isfinite(overall_ci)):
            failed += 1
            continue
        lo_o, hi_o = float(overall_ci[0]), float(overall_ci[1])
        if lo_o <= tau_true <= hi_o:
            covered_overall += 1

        # Horizon-1 and horizon-2 bootstrap CIs (guard the shared-
        # PSU-weight multi-horizon path). Horizon-2 in particular
        # requires the SAME shared PSU weight matrix drawn once at
        # the top of the multi-horizon block; a per-horizon re-draw
        # would break the sup-t joint-distribution guarantee and
        # `cband_crit_value` would be undefined or wrong.
        es_cis = res.bootstrap_results.event_study_cis
        if es_cis is not None:
            if 1 in es_cis:
                h1_ci = es_cis[1]
                if h1_ci is not None and all(np.isfinite(h1_ci)):
                    lo_h, hi_h = float(h1_ci[0]), float(h1_ci[1])
                    if lo_h <= tau_true <= hi_h:
                        covered_h1 += 1
            if 2 in es_cis:
                h2_ci = es_cis[2]
                if h2_ci is not None and all(np.isfinite(h2_ci)):
                    lo2, hi2 = float(h2_ci[0]), float(h2_ci[1])
                    if lo2 <= tau_true <= hi2:
                        covered_h2 += 1

        # Sup-t critical value: finite across reps means the shared-
        # draw machinery produced coherent joint replicates at both
        # horizons. NaN or unset would indicate the multi-horizon
        # block short-circuited or the shared-weight broadcast
        # misaligned across horizons.
        cband = getattr(res.bootstrap_results, "cband_crit_value", None)
        if cband is not None and np.isfinite(float(cband)):
            cband_finite += 1

    completed = n_reps - failed
    assert completed >= int(0.95 * n_reps), (
        f"MC simulation had {failed}/{n_reps} fit failures, above "
        f"the 5% tolerance."
    )
    coverage_overall = covered_overall / completed
    coverage_h1 = covered_h1 / completed
    coverage_h2 = covered_h2 / completed
    assert 0.925 <= coverage_overall <= 0.975, (
        f"Overall bootstrap CI coverage {coverage_overall:.3f} "
        f"(completed {completed}) outside [0.925, 0.975]; "
        f"tau_true={tau_true}, n_groups={n_groups}, n_periods={n_periods}, "
        f"n_reps={n_reps}."
    )
    assert 0.925 <= coverage_h1 <= 0.975, (
        f"Horizon-1 event-study bootstrap CI coverage "
        f"{coverage_h1:.3f} (completed {completed}) outside "
        f"[0.925, 0.975]; this is the shared-PSU-weight surface, "
        f"regression here likely indicates a bug in the multi-horizon "
        f"cell-level broadcast."
    )
    # Horizon-2 tolerance is wider than horizon-1 because finite-
    # sample coverage of the analytical TSL SE on this DGP is
    # itself ~0.93 at l=2 (measured offline: analytical h-1 coverage
    # 0.94, h-2 coverage 0.926 at n_groups=40). The bootstrap should
    # track the analytical SE asymptotically, so an observed
    # bootstrap coverage in [0.91, 0.98] at h-2 is consistent with
    # correct clustering; a drop to ≤ 0.90 would indicate the
    # shared-weight broadcast is not coherent across horizons.
    assert 0.910 <= coverage_h2 <= 0.975, (
        f"Horizon-2 event-study bootstrap CI coverage "
        f"{coverage_h2:.3f} (completed {completed}) outside "
        f"[0.910, 0.975]; horizon-2 is the cross-horizon surface "
        f"that exercises the SAME shared PSU weight matrix used "
        f"at horizon-1 — a regression here indicates the shared-"
        f"draw broadcast is not coherent across horizons."
    )
    # Sup-t critical value must be finite in the vast majority of
    # reps; occasional NaN on degenerate draws is tolerable but
    # widespread NaN signals the shared-weight block never yielded
    # a coherent joint distribution.
    assert cband_finite >= int(0.90 * completed), (
        f"Sup-t critical value was finite in only {cband_finite}/"
        f"{completed} reps. The shared (n_bootstrap, n_psu) PSU-"
        f"level weight matrix must be drawn ONCE at the top of the "
        f"multi-horizon block; a per-horizon re-draw would break "
        f"the sup-t joint distribution."
    )
