"""R-parity tests for StackedDiD's vcov_type threading (Phase 1b 2/8).

Pins Python StackedDiD output against R `clubSandwich::vcovCR(lm(weights=Q,...),
type="CR1S"|"CR2") + coef_test()$df_Satt` and `Wald_test(test="HTZ")$df_denom`
at atol=1e-10. The clubSandwich port in `diff_diff/linalg.py` (PR #475)
supplies the underlying WLS-CR2 algebra; this file verifies that StackedDiD's
call into `solve_ols(weights=composed_weights, vcov_type=...)` produces output
consistent with the R reference, including the Bell-McCaffrey Satterthwaite
DOF threading at both the per-event-time level and the post-period-average
contrast level.

Goldens are generated by `benchmarks/R/generate_stacked_did_golden.R` (loads
the pre-stacked panel CSV at `benchmarks/data/stacked_did_test_panel.csv`,
fits `lm(weights=Q,...)`, dumps CR1S + CR2-BM + Satterthwaite DOF for the
event-study coefficients AND the post-period-average ATT contrast at both
`cluster=unit` and `cluster=unit_subexp`).

Note on `CR1` vs `CR1S`: diff-diff's HC1 + cluster uses Stata-style
`G/(G-1) * (n-1)/(n-p)` finite-sample correction, which matches
`clubSandwich::vcovCR(type="CR1S")`. The plain `clubSandwich::vcovCR(type="CR1")`
omits the `(n-1)/(n-p)` term and diverges by ~1.4% on this fixture.
"""

from __future__ import annotations

import json
import os

import numpy as np
import pytest
from scipy.optimize import brentq
from scipy.stats import t as t_dist

from diff_diff import StackedDiD
from diff_diff.prep_dgp import generate_staggered_data

_GOLDEN_PATH = os.path.join(
    os.path.dirname(__file__),
    "..",
    "benchmarks",
    "data",
    "stacked_did_golden.json",
)
_R_FIXTURE_AVAILABLE = os.path.exists(_GOLDEN_PATH)


@pytest.fixture(scope="module")
def goldens():
    if not _R_FIXTURE_AVAILABLE:
        pytest.skip(
            "stacked_did_golden.json not generated; run "
            "Rscript benchmarks/R/generate_stacked_did_golden.R"
        )
    with open(_GOLDEN_PATH) as f:
        return json.load(f)


_PANEL_CSV_PATH = os.path.join(
    os.path.dirname(__file__),
    "..",
    "benchmarks",
    "data",
    "stacked_did_test_panel.csv",
)


@pytest.fixture(scope="module")
def panel():
    """Same panel descriptor as the R fixture; must match exactly.

    Per R6 P3: cross-check the regenerated panel against the committed
    `stacked_did_test_panel.csv` to detect any drift in
    `generate_staggered_data` that would otherwise silently invalidate the
    R-parity assertions. The CSV is the post-stacking panel; we verify
    that the pre-stacking columns (unit, period, outcome, first_treat)
    match column-by-column. If `generate_staggered_data` ever changes,
    this test will FAIL with a clear DGP-drift message rather than
    producing wrong parity output.
    """
    data = generate_staggered_data(
        n_units=50,
        n_periods=8,
        cohort_periods=[3, 5, 7],
        never_treated_frac=0.3,
        treatment_effect=2.0,
        dynamic_effects=False,
        seed=20260521,
    )
    # DGP drift check: load committed CSV (post-stacking, contains a subset
    # of original rows duplicated across sub-experiments) and verify that
    # the (unit, period) tuples present in the CSV have matching `outcome`
    # values in the regenerated panel. The CSV is the canonical reference
    # used by the R fixture script; any drift in `generate_staggered_data`
    # that changes the underlying outcomes would otherwise silently
    # invalidate the R-parity assertions below.
    if os.path.exists(_PANEL_CSV_PATH):
        import pandas as pd

        csv = pd.read_csv(_PANEL_CSV_PATH)
        csv_unique = csv.drop_duplicates(subset=["unit", "period"]).copy()
        merged = csv_unique.merge(
            data[["unit", "period", "outcome", "first_treat"]],
            on=["unit", "period"],
            how="inner",
            suffixes=("_csv", "_regen"),
        )
        assert len(merged) == len(csv_unique), (
            f"DGP drift: regenerated panel missing (unit, period) rows present "
            f"in committed CSV. CSV unique rows={len(csv_unique)}, "
            f"merged rows={len(merged)}. Regenerate the CSV if the DGP "
            "helper change is intentional."
        )
        # Outcome parity (drives the R-side fit).
        np.testing.assert_allclose(
            merged["outcome_regen"].astype(float).values,
            merged["outcome_csv"].astype(float).values,
            atol=1e-12,
            err_msg=(
                "DGP drift detected: generate_staggered_data(...) produces "
                "different 'outcome' values than the committed "
                "benchmarks/data/stacked_did_test_panel.csv on the "
                "matching (unit, period) rows. Regenerate the CSV "
                "if the change is intentional."
            ),
        )
        # Treatment-timing parity: normalize first_treat (CSV stores inf
        # for never-treated; regenerated panel uses 0 in some pandas
        # versions). Per R7 P3: pin the treatment-timing inputs too so
        # a future drift in cohort assignment is caught independently of
        # outcome drift.
        regen_ft = merged["first_treat_regen"].astype(float).replace(0.0, np.inf).values
        csv_ft = merged["first_treat_csv"].astype(float).values
        # CSV stores 0 for never-treated in some serializations too; normalize.
        csv_ft = np.where(csv_ft == 0.0, np.inf, csv_ft)
        np.testing.assert_array_equal(
            regen_ft,
            csv_ft,
            err_msg=(
                "DGP drift detected: generate_staggered_data(...) produces "
                "different treatment timing ('first_treat') than the committed "
                "CSV. Cohort assignment changed; regenerate CSV if intentional."
            ),
        )
    return data


def _fit(panel, vcov_type, cluster):
    est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type=vcov_type, cluster=cluster)
    return est.fit(
        panel,
        outcome="outcome",
        unit="unit",
        time="period",
        first_treat="first_treat",
        aggregate="event_study",
    )


def _es_se_vector(res, event_times_non_ref):
    """Extract event-study SE vector in the order given by event_times_non_ref."""
    return np.array([res.event_study_effects[h]["se"] for h in event_times_non_ref])


def _recover_dof_from_ci(att, se, ci_hi, alpha):
    """Recover DOF used by safe_inference from the CI half-width.

    Inverse of `ci_hi = att + se * t.ppf(1 - alpha/2, df)`. Used to assert
    the BM Satterthwaite DOF threaded correctly into the result-class
    inference fields without relying on internal estimator state.

    Raises ValueError if the inversion fails (e.g., when safe_inference
    used the normal distribution instead of t — meaning the BM DOF did
    NOT propagate and the test should FAIL, not silently skip).
    """
    half_width = ci_hi - att
    t_crit = half_width / se
    target_alpha = alpha / 2  # 0.025 for alpha=0.05

    def f(df):
        return t_dist.sf(t_crit, df) - target_alpha

    # Use a wide bracket; brentq raises ValueError if no root.
    return brentq(f, 1.0, 1e9, xtol=1e-9)


class TestStackedDiDParityR:
    """Pin StackedDiD output against clubSandwich at atol=1e-10.

    Coverage:
      - hc1 + cluster=unit SE matches CR1S (Stata-style correction)
      - hc2_bm + cluster=unit SE matches CR2 Bell-McCaffrey
      - hc2_bm + cluster=unit per-event-time DOF matches coef_test()$df_Satt
      - hc2_bm + cluster=unit overall ATT DOF matches Wald_test(HTZ)$df_denom
      - hc2_bm + cluster=unit_subexp SE matches CR2 (symmetric)

    The DOF parity tests use CI-width inversion to recover the DOF that
    `safe_inference` actually used. If the BM DOF doesn't propagate (the
    R1 P0 bug), `safe_inference` uses the normal distribution and the
    half-width is `se * 1.96`, which inverts to `df → infinity` — that
    inversion FAILS within the bracket [1, 1e9] and raises ValueError,
    which surfaces as a test failure (not a silent skip).
    """

    def test_hc1_se_matches_clubsandwich_cr1s_unit_cluster(self, goldens, panel):
        """Default cluster=unit, vcov_type=hc1 matches R `vcovCR(type='CR1S')`.

        Pins BOTH per-event-time SE AND the post-period-average overall_se
        (per R6 P3: prior fixture only pinned per-event SE; the overall
        delta-method SE was untested, leaving a gap where post-period
        covariance regressions could slip through).

        See module docstring: 'CR1S' (Stata-style) is the correct comparison
        target; plain 'CR1' (clubSandwich's no-(n-1)/(n-p) form) would diverge
        by ~1.4% on this fixture.
        """
        event_times = goldens["meta"]["event_times_non_ref"]
        r_se = np.array(goldens["unit"]["se_cr1_es"])
        r_overall_se = float(goldens["unit"]["se_overall_cr1"])
        res = _fit(panel, vcov_type="hc1", cluster="unit")
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "StackedDiD hc1 event-study SE must match clubSandwich CR1S at atol=1e-10. "
                "Note: target is CR1S (Stata-style), not plain CR1."
            ),
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "StackedDiD hc1 overall_se (post-period delta-method) must "
                "match clubSandwich CR1S contrast variance at atol=1e-10"
            ),
        )

    def test_hc2_bm_se_matches_clubsandwich_cr2_unit_cluster(self, goldens, panel):
        """cluster=unit, vcov_type=hc2_bm matches R `vcovCR(type='CR2')`.

        Pins BOTH per-event-time SE AND overall_se contract variance
        (per R6 P3).
        """
        event_times = goldens["meta"]["event_times_non_ref"]
        r_se = np.array(goldens["unit"]["se_cr2_es"])
        r_overall_se = float(goldens["unit"]["se_overall_cr2"])
        res = _fit(panel, vcov_type="hc2_bm", cluster="unit")
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg="StackedDiD hc2_bm event-study SE must match clubSandwich CR2 at atol=1e-10",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "StackedDiD hc2_bm overall_se must match clubSandwich CR2 "
                "contrast variance at atol=1e-10"
            ),
        )

    def test_hc2_bm_per_event_dof_matches_coef_test_df_satt_unit_cluster(self, goldens, panel):
        """Per-event-time BM DOF propagated via safe_inference matches
        `coef_test(..., vcov=CR2)$df_Satt`.

        Recovers the DOF from CI half-width inversion. If BM DOF didn't
        propagate, `safe_inference` would use the normal distribution
        (df=None) and the recovery would fail with ValueError — which
        surfaces as a test failure, NOT a silent skip. This was the R1 P0
        failure mode caught by CI codex.
        """
        event_times = goldens["meta"]["event_times_non_ref"]
        r_dof = np.array(goldens["unit"]["dof_bm_es"])
        res = _fit(panel, vcov_type="hc2_bm", cluster="unit")
        validated_count = 0
        for h, r_dof_h in zip(event_times, r_dof):
            es = res.event_study_effects[h]
            att, se, ci_hi = es["effect"], es["se"], es["conf_int"][1]
            # If propagation failed, this raises ValueError → test fails.
            df_python = _recover_dof_from_ci(att, se, ci_hi, res.alpha)
            np.testing.assert_allclose(
                df_python,
                r_dof_h,
                atol=1e-6,
                rtol=1e-6,
                err_msg=(
                    f"BM DOF for event_time={h} must match clubSandwich coef_test "
                    f"df_Satt; recovered df_python={df_python}, R df_Satt={r_dof_h}"
                ),
            )
            validated_count += 1
        # Hard guard: ensure we actually validated every coefficient.
        assert validated_count == len(event_times), (
            f"Expected to validate DOF for all {len(event_times)} event-time "
            f"coefficients, only validated {validated_count}"
        )

    def test_hc2_bm_overall_att_dof_matches_wald_test_htz_unit_cluster(self, goldens, panel):
        """Overall ATT BM DOF matches `Wald_test(HTZ)$df_denom` (per review P2:
        the post-period-average contrast was missing from the prior fixture)."""
        r_dof_overall = float(goldens["unit"]["dof_bm_overall"])
        res = _fit(panel, vcov_type="hc2_bm", cluster="unit")
        att, se, ci_hi = res.overall_att, res.overall_se, res.overall_conf_int[1]
        df_python = _recover_dof_from_ci(att, se, ci_hi, res.alpha)
        np.testing.assert_allclose(
            df_python,
            r_dof_overall,
            atol=1e-6,
            rtol=1e-6,
            err_msg=(
                f"Overall ATT BM DOF must match clubSandwich Wald_test(HTZ) "
                f"df_denom; recovered df_python={df_python}, R df={r_dof_overall}"
            ),
        )

    def test_hc2_bm_se_matches_clubsandwich_cr2_unit_subexp_cluster(self, goldens, panel):
        """Same parity for the alternate cluster level (per-event SE +
        overall_se contract variance)."""
        event_times = goldens["meta"]["event_times_non_ref"]
        r_se = np.array(goldens["unit_subexp"]["se_cr2_es"])
        r_overall_se = float(goldens["unit_subexp"]["se_overall_cr2"])
        res = _fit(panel, vcov_type="hc2_bm", cluster="unit_subexp")
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "StackedDiD hc2_bm + cluster=unit_subexp must match "
                "clubSandwich CR2 at atol=1e-10"
            ),
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "StackedDiD hc2_bm + cluster=unit_subexp overall_se must "
                "match clubSandwich CR2 at atol=1e-10"
            ),
        )

    def test_hc2_bm_overall_att_dof_matches_wald_test_htz_unit_subexp_cluster(self, goldens, panel):
        """Symmetric overall ATT DOF parity at the alternate cluster level."""
        r_dof_overall = float(goldens["unit_subexp"]["dof_bm_overall"])
        res = _fit(panel, vcov_type="hc2_bm", cluster="unit_subexp")
        att, se, ci_hi = res.overall_att, res.overall_se, res.overall_conf_int[1]
        df_python = _recover_dof_from_ci(att, se, ci_hi, res.alpha)
        np.testing.assert_allclose(
            df_python,
            r_dof_overall,
            atol=1e-6,
            rtol=1e-6,
            err_msg=(
                f"Overall ATT BM DOF (unit_subexp) must match clubSandwich; "
                f"recovered df_python={df_python}, R df={r_dof_overall}"
            ),
        )


# =============================================================================
# Variant parity (per R8 P1): pin non-default vcov_type surfaces to R goldens.
# =============================================================================


@pytest.fixture(scope="module")
def population_panel():
    """Heterogeneous-population panel matching the R variant fixture.

    Same seed/population assignment as the Python helper that dumped
    `stacked_did_population_panel.csv`. Per R8 P1: the prior `population`
    smoke test used a constant pop=1.0 which is numerically degenerate
    with aggregate weighting. This fixture forces heterogeneous Q-weights.
    """
    rng = np.random.default_rng(99)
    data = generate_staggered_data(
        n_units=50,
        n_periods=8,
        cohort_periods=[3, 5, 7],
        never_treated_frac=0.3,
        treatment_effect=2.0,
        dynamic_effects=False,
        seed=20260521,
    )
    unit_pops = rng.uniform(1.0, 100.0, size=data["unit"].nunique())
    data = data.copy()
    data["pop"] = data["unit"].map(lambda u: unit_pops[int(u)])
    return data


class TestStackedDiDParityRVariants:
    """R-parity for non-default vcov_type surfaces (R8 P1 coverage)."""

    def test_population_weighting_hc1_matches_clubsandwich_cr1s(self, goldens, population_panel):
        """weighting='population' with heterogeneous Q-weights produces
        SE that differs from the aggregate-weight case AND matches R."""
        if "population" not in goldens:
            pytest.skip("population variant not in goldens; regenerate fixture")
        r_se_es = np.array(goldens["population"]["unit"]["se_cr1_es"])
        r_overall_se = float(goldens["population"]["unit"]["se_overall_cr1"])
        event_times = goldens["population"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc1", weighting="population")
        res = est.fit(
            population_panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            population="pop",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="StackedDiD weighting='population' hc1 event-study SE must match R CR1S",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="StackedDiD weighting='population' hc1 overall_se must match R",
        )

    def test_population_weighting_hc2_bm_matches_clubsandwich_cr2(self, goldens, population_panel):
        if "population" not in goldens:
            pytest.skip("population variant not in goldens")
        r_se_es = np.array(goldens["population"]["unit"]["se_cr2_es"])
        r_overall_se = float(goldens["population"]["unit"]["se_overall_cr2"])
        event_times = goldens["population"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc2_bm", weighting="population")
        res = est.fit(
            population_panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            population="pop",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="StackedDiD weighting='population' hc2_bm event-study SE must match R CR2",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="StackedDiD weighting='population' hc2_bm overall_se must match R",
        )
        # DOF parity (per R9 P1)
        r_dof_es = np.array(goldens["population"]["unit"]["dof_bm_es"])
        r_dof_overall = float(goldens["population"]["unit"]["dof_bm_overall"])
        for h, r_dof_h in zip(event_times, r_dof_es):
            es = res.event_study_effects[h]
            if es["se"] == 0:
                continue
            df_py = _recover_dof_from_ci(es["effect"], es["se"], es["conf_int"][1], res.alpha)
            np.testing.assert_allclose(
                df_py,
                r_dof_h,
                atol=1e-6,
                rtol=1e-6,
                err_msg=f"population hc2_bm BM DOF at h={h} must match R",
            )
        df_overall = _recover_dof_from_ci(
            res.overall_att, res.overall_se, res.overall_conf_int[1], res.alpha
        )
        np.testing.assert_allclose(
            df_overall,
            r_dof_overall,
            atol=1e-6,
            rtol=1e-6,
            err_msg="population hc2_bm overall BM DOF must match R Wald_test(HTZ)",
        )

    def test_anticipation1_hc1_matches_clubsandwich_cr1s(self, goldens, panel):
        """anticipation=1: reference period shifts to e=-2, post-period
        starts at h=-1, event window expands by 1 pre-period. Verify R
        parity on both event-study and overall SE."""
        if "anticipation1" not in goldens:
            pytest.skip("anticipation1 variant not in goldens")
        r_se_es = np.array(goldens["anticipation1"]["unit"]["se_cr1_es"])
        r_overall_se = float(goldens["anticipation1"]["unit"]["se_overall_cr1"])
        event_times = goldens["anticipation1"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc1", anticipation=1)
        res = est.fit(
            panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="StackedDiD anticipation=1 hc1 event-study SE must match R CR1S",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="StackedDiD anticipation=1 hc1 overall_se must match R",
        )

    def test_anticipation1_hc2_bm_matches_clubsandwich_cr2(self, goldens, panel):
        if "anticipation1" not in goldens:
            pytest.skip("anticipation1 variant not in goldens")
        r_se_es = np.array(goldens["anticipation1"]["unit"]["se_cr2_es"])
        r_overall_se = float(goldens["anticipation1"]["unit"]["se_overall_cr2"])
        event_times = goldens["anticipation1"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc2_bm", anticipation=1)
        res = est.fit(
            panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="StackedDiD anticipation=1 hc2_bm event-study SE must match R CR2",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="StackedDiD anticipation=1 hc2_bm overall_se must match R",
        )
        # DOF parity (per R9 P1): under anticipation=1, the post-period
        # contrast includes h=-1, so the BM DOF should reflect that extension.
        r_dof_es = np.array(goldens["anticipation1"]["unit"]["dof_bm_es"])
        r_dof_overall = float(goldens["anticipation1"]["unit"]["dof_bm_overall"])
        for h, r_dof_h in zip(event_times, r_dof_es):
            es = res.event_study_effects[h]
            if es["se"] == 0:
                continue
            df_py = _recover_dof_from_ci(es["effect"], es["se"], es["conf_int"][1], res.alpha)
            np.testing.assert_allclose(
                df_py,
                r_dof_h,
                atol=1e-6,
                rtol=1e-6,
                err_msg=f"anticipation1 hc2_bm BM DOF at h={h} must match R",
            )
        df_overall = _recover_dof_from_ci(
            res.overall_att, res.overall_se, res.overall_conf_int[1], res.alpha
        )
        np.testing.assert_allclose(
            df_overall,
            r_dof_overall,
            atol=1e-6,
            rtol=1e-6,
            err_msg="anticipation1 hc2_bm overall BM DOF must match R Wald_test(HTZ)",
        )

    def test_sample_share_weighting_hc1_matches_clubsandwich_cr1s(self, goldens, panel):
        """weighting='sample_share' uses a distinct Q-weight formula
        `Q_sa = ((N^D_a + N^C_a)/(N^D+N^C)) / (N^C_a/N^C)` for controls.
        Pin R parity at atol=1e-10 (per R9 P1)."""
        if "sample_share" not in goldens:
            pytest.skip("sample_share variant not in goldens")
        r_se_es = np.array(goldens["sample_share"]["unit"]["se_cr1_es"])
        r_overall_se = float(goldens["sample_share"]["unit"]["se_overall_cr1"])
        event_times = goldens["sample_share"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc1", weighting="sample_share")
        res = est.fit(
            panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="sample_share hc1 event-study SE must match R CR1S",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="sample_share hc1 overall_se must match R",
        )

    def test_sample_share_weighting_hc2_bm_matches_clubsandwich_cr2(self, goldens, panel):
        """sample_share + hc2_bm: pin SE + BM DOF parity."""
        if "sample_share" not in goldens:
            pytest.skip("sample_share variant not in goldens")
        r_se_es = np.array(goldens["sample_share"]["unit"]["se_cr2_es"])
        r_overall_se = float(goldens["sample_share"]["unit"]["se_overall_cr2"])
        r_dof_es = np.array(goldens["sample_share"]["unit"]["dof_bm_es"])
        r_dof_overall = float(goldens["sample_share"]["unit"]["dof_bm_overall"])
        event_times = goldens["sample_share"]["meta"]["event_times_non_ref"]
        est = StackedDiD(kappa_pre=2, kappa_post=2, vcov_type="hc2_bm", weighting="sample_share")
        res = est.fit(
            panel,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
        )
        py_se = _es_se_vector(res, event_times)
        np.testing.assert_allclose(
            py_se,
            r_se_es,
            atol=1e-10,
            rtol=1e-10,
            err_msg="sample_share hc2_bm event-study SE must match R CR2",
        )
        np.testing.assert_allclose(
            res.overall_se,
            r_overall_se,
            atol=1e-10,
            err_msg="sample_share hc2_bm overall_se must match R",
        )
        for h, r_dof_h in zip(event_times, r_dof_es):
            es = res.event_study_effects[h]
            if es["se"] == 0:
                continue
            df_py = _recover_dof_from_ci(es["effect"], es["se"], es["conf_int"][1], res.alpha)
            np.testing.assert_allclose(
                df_py,
                r_dof_h,
                atol=1e-6,
                rtol=1e-6,
                err_msg=f"sample_share hc2_bm BM DOF at h={h} must match R",
            )
        df_overall = _recover_dof_from_ci(
            res.overall_att, res.overall_se, res.overall_conf_int[1], res.alpha
        )
        np.testing.assert_allclose(
            df_overall,
            r_dof_overall,
            atol=1e-6,
            rtol=1e-6,
            err_msg="sample_share hc2_bm overall BM DOF must match R Wald_test(HTZ)",
        )
