"""Parity tests for WLS-CR2 Bell-McCaffrey against clubSandwich R goldens.

This file pins `diff_diff/linalg.py`'s weighted CR2-BM port against
`clubSandwich::vcovCR(lm(weights=w), type="CR2") + coef_test(test=
"Satterthwaite")$df_Satt` plus the compound-contrast DOF from
`Wald_test(constraints=matrix(c, 1), test="HTZ")$df_denom` (which on a 1-row
constraint matrix reduces to a Satterthwaite t-test).

Goldens are generated by `benchmarks/R/generate_clubsandwich_golden.R` (six
new weighted scenarios appended in this PR; existing six unweighted scenarios
unchanged for regression safety).

Algorithm note: clubSandwich's WLS-CR2 algebra is NOT a textbook
Pustejovsky-Tipton (2018) §3.3 transform-once derivation — the two diverge by
0.5-30% on weighted designs (see `feedback_wls_cr2_clubsandwich_parity`).
clubSandwich uses W (not sqrt(W)) in the hat matrix, W^2 in the bias-correction
term, and unweighted residuals in the score construction. The Satterthwaite
DOF uses the full H1/H2/H3 array construction (`get_arrays.R::get_GH`), NOT
the simplified `(tr B)^2 / tr(B^2)` formula that works for the unweighted case.
"""

from __future__ import annotations

import json
import os

import numpy as np
import pytest

from diff_diff.linalg import (
    _compute_bm_dof_oneway,
    _compute_cr2_bm,
    _compute_cr2_bm_contrast_dof,
    _compute_hat_diagonals,
)

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


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


def _design_x_z(d):
    """Build [1, x, z] design matrix from a golden scenario dict."""
    return np.column_stack([np.ones(len(d["x"])), np.asarray(d["x"]), np.asarray(d["z"])])


def _refit_wls(X, y, w):
    """Refit WLS to obtain coef and residuals consistent with R's lm()."""
    bread = X.T @ (X * w[:, None])
    coef = np.linalg.solve(bread, X.T @ (w * y))
    return coef, y - X @ coef, bread


class TestWLSCR2ParityClubSandwich:
    """Parity for weighted CR2 Bell-McCaffrey against clubSandwich's WLS-CR2."""

    def test_one_way_weighted_hc2_bm_se_matches_clubsandwich(self, goldens):
        d = goldens["weighted_one_way_no_cluster"]
        X = _design_x_z(d)
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        coef, resid, bread = _refit_wls(X, y, w)
        _compute_hat_diagonals(X, bread, weights=w)
        # One-way (no cluster) HC2-BM SE via the diff-diff dispatcher.
        # Per the lifted gate, _compute_bm_dof_oneway returns the DOF; the SE
        # is sqrt(diag(vcov)) where vcov is computed in `_compute_robust_vcov_numpy`.
        # We pin both the vcov (via R's vcov_hc2_bm) and the DOF here.
        r_vcov = np.asarray(d["vcov_hc2_bm"]).reshape(d["vcov_hc2_bm_shape"])
        r_se = np.asarray(d["se_hc2_bm"])
        r_dof = np.asarray(d["dof_bm_one_way"])

        # Use the public dispatcher path
        from diff_diff.linalg import compute_robust_vcov

        vcov, dof = compute_robust_vcov(
            X,
            resid,
            vcov_type="hc2_bm",
            weights=w,
            weight_type="pweight",
            return_dof=True,
        )
        np.testing.assert_allclose(
            vcov,
            r_vcov,
            atol=1e-10,
            rtol=1e-10,
            err_msg="weighted one-way HC2-BM vcov should match clubSandwich",
        )
        np.testing.assert_allclose(
            np.sqrt(np.diag(vcov)),
            r_se,
            atol=1e-10,
            err_msg="weighted one-way HC2-BM SE should match clubSandwich",
        )
        np.testing.assert_allclose(
            dof, r_dof, atol=1e-10, err_msg="weighted one-way HC2-BM DOF should match clubSandwich"
        )

    def test_balanced_clusters_weighted_cr2_se(self, goldens):
        d = goldens["weighted_balanced_clusters"]
        X = _design_x_z(d)
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster_ids = np.asarray(d["cluster"])
        coef, resid, bread = _refit_wls(X, y, w)
        r_vcov = np.asarray(d["vcov_cr2"]).reshape(d["vcov_cr2_shape"])
        r_dof = np.asarray(d["dof_bm"])

        vcov, dof = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w)
        np.testing.assert_allclose(vcov, r_vcov, atol=1e-10)
        np.testing.assert_allclose(dof, r_dof, atol=1e-10)

    def test_unbalanced_clusters_weighted_cr2_dof(self, goldens):
        d = goldens["weighted_unbalanced_clusters"]
        X = _design_x_z(d)
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster_ids = np.asarray(d["cluster"])
        coef, resid, bread = _refit_wls(X, y, w)
        r_vcov = np.asarray(d["vcov_cr2"]).reshape(d["vcov_cr2_shape"])
        r_dof = np.asarray(d["dof_bm"])

        vcov, dof = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w)
        np.testing.assert_allclose(vcov, r_vcov, atol=1e-10)
        np.testing.assert_allclose(dof, r_dof, atol=1e-10)

    def test_singletons_present_adversarial(self, goldens):
        """Adversarial case: PT2018 transform-once had ~30% gap here. The
        clubSandwich algebra must match clubSandwich at 1e-10."""
        d = goldens["weighted_singletons_present"]
        X = _design_x_z(d)
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster_ids = np.asarray(d["cluster"])
        coef, resid, bread = _refit_wls(X, y, w)
        r_vcov = np.asarray(d["vcov_cr2"]).reshape(d["vcov_cr2_shape"])
        r_dof = np.asarray(d["dof_bm"])

        vcov, dof = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w)
        np.testing.assert_allclose(
            vcov,
            r_vcov,
            atol=1e-10,
            err_msg="adversarial singleton-cluster case must match clubSandwich",
        )
        np.testing.assert_allclose(dof, r_dof, atol=1e-10)


class TestWLSCR2ContrastDOF:
    """Compound-contrast Satterthwaite DOF parity (MPD avg_att integration)."""

    def test_avg_att_compound_contrast_dof_matches_wald_test_htz(self, goldens):
        d = goldens["weighted_mpd_avg_att_dof"]
        # The MPD scenario has a large design matrix (28+ columns) with a
        # rank-deficient column R drops via NA (treated lies in the span of
        # the intercept + non-reference unit dummies for treated cohorts).
        # Build the rank-reduced design matching `finite_coef_names` from R.
        n = len(d["y"])
        period_arr = np.asarray(d["period"])
        unit_arr = np.asarray(d["unit"])
        treated_arr = np.asarray(d["treated"])
        finite_names = d["finite_coef_names"]
        sorted(np.unique(period_arr).tolist())
        sorted(np.unique(unit_arr).tolist())

        def _make_column(name):
            if name == "(Intercept)":
                return np.ones(n)
            if name == "treated":
                return treated_arr.astype(float)
            if name.startswith("period_f"):
                p = int(name[len("period_f") :])
                return (period_arr == p).astype(float)
            if name.startswith("treated_period_"):
                p = int(name[len("treated_period_") :])
                return (treated_arr * (period_arr == p)).astype(float)
            if name.startswith("factor(unit)"):
                u = int(name[len("factor(unit)") :])
                return (unit_arr == u).astype(float)
            raise ValueError(f"Unknown column name: {name}")

        X = np.column_stack([_make_column(name) for name in finite_names])
        np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster_ids = np.asarray(d["cluster"])

        bread = X.T @ (X * w[:, None])
        # R's `c_avg` is already subsetted to finite (non-NA) columns.
        contrasts = np.asarray(d["c_avg"])[:, None]

        r_dof_avg = float(d["dof_avg"])
        dof = _compute_cr2_bm_contrast_dof(X, cluster_ids, bread, contrasts, weights=w)
        np.testing.assert_allclose(
            dof[0],
            r_dof_avg,
            atol=1e-10,
            err_msg="weighted avg_att compound DOF should match Wald_test HTZ",
        )


class TestWLSCR2WeightTypeRejection:
    """The clubSandwich WLS-CR2 port matches `pweight` only. Other weight types
    must raise NotImplementedError rather than silently producing wrong inference.
    Addresses Codex P0 finding (weight_type contract gap)."""

    def test_aweight_clustered_hc2_bm_rejects(self):
        rng = np.random.default_rng(101)
        n, k = 30, 3
        X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.uniform(0.5, 2.0, n)
        cluster_ids = np.arange(n) % 5
        from diff_diff.linalg import compute_robust_vcov

        with pytest.raises(NotImplementedError, match="pweight"):
            compute_robust_vcov(
                X,
                y - X @ np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y)),
                cluster_ids=cluster_ids,
                vcov_type="hc2_bm",
                weights=w,
                weight_type="aweight",
            )

    def test_fweight_clustered_hc2_bm_rejects(self):
        rng = np.random.default_rng(102)
        n, k = 30, 3
        X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.integers(1, 4, size=n).astype(np.float64)
        cluster_ids = np.arange(n) % 5
        from diff_diff.linalg import compute_robust_vcov

        with pytest.raises(NotImplementedError, match="pweight"):
            compute_robust_vcov(
                X,
                y - X @ np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y)),
                cluster_ids=cluster_ids,
                vcov_type="hc2_bm",
                weights=w,
                weight_type="fweight",
            )

    def test_aweight_one_way_hc2_bm_rejects(self):
        rng = np.random.default_rng(103)
        n, k = 30, 3
        X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.uniform(0.5, 2.0, n)
        from diff_diff.linalg import compute_robust_vcov

        with pytest.raises(NotImplementedError, match="pweight"):
            compute_robust_vcov(
                X,
                y - X @ np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y)),
                vcov_type="hc2_bm",
                weights=w,
                weight_type="aweight",
            )

    def test_pweight_clustered_hc2_bm_accepted(self):
        """pweight (the matching convention) must NOT raise."""
        rng = np.random.default_rng(104)
        n, k = 30, 3
        X = np.column_stack([np.ones(n), rng.normal(size=(n, k - 1))])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.uniform(0.5, 2.0, n)
        cluster_ids = np.arange(n) % 5
        from diff_diff.linalg import compute_robust_vcov

        vcov, dof = compute_robust_vcov(
            X,
            y - X @ np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y)),
            cluster_ids=cluster_ids,
            vcov_type="hc2_bm",
            weights=w,
            weight_type="pweight",
            return_dof=True,
        )
        assert np.all(np.isfinite(vcov)) and np.all(np.diag(vcov) > 0)
        assert np.all(np.isfinite(dof)) and np.all(dof > 0)


class TestWLSCR2ZeroWeightClusterRejection:
    """Zero-total-weight clusters (subpopulation-style padding) must NOT pass
    the `G>=2` check when only one effective cluster remains. Addresses Codex
    P1 finding."""

    def test_clustered_with_one_zero_weight_cluster_raises(self):
        """A weighted CR2 fit where one cluster has all-zero weights and only
        one other cluster has positive weight must raise."""
        rng = np.random.default_rng(201)
        n = 20
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        # Two clusters: 0 (positive weight) and 1 (all-zero weight).
        cluster_ids = np.array([0] * 10 + [1] * 10)
        w = np.concatenate([np.ones(10), np.zeros(10)])  # only cluster 0 has weight
        bread = X.T @ (X * w[:, None])
        # `bread_inv` would be rank-deficient if all-zero rows reduce design,
        # but we can ensure positive bread by adding a third tiny-weight cluster.
        cluster_ids = np.array([0] * 10 + [1] * 10)
        # Reformulate: 2 unique clusters but only 1 has positive total weight
        # (cluster 1 has w=0 for all its members).
        coef = np.linalg.solve(bread, X.T @ (w * y))
        resid = y - X @ coef
        with pytest.raises(ValueError, match="positive total weight"):
            _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w)

    def test_contrast_dof_clustered_with_one_zero_weight_cluster_raises(self):
        """Same check for the contrast-DOF helper."""
        rng = np.random.default_rng(202)
        n = 20
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
        X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        cluster_ids = np.array([0] * 10 + [1] * 10)
        w = np.concatenate([np.ones(10), np.zeros(10)])
        bread = X.T @ (X * w[:, None])
        contrasts = np.eye(X.shape[1])
        with pytest.raises(ValueError, match="positive total weight"):
            _compute_cr2_bm_contrast_dof(X, cluster_ids, bread, contrasts, weights=w)

    def test_clustered_zero_weight_cluster_drops_silently_when_eff_clusters_remain(self):
        """A zero-weight cluster amongst multiple positive-weight clusters
        should be silently dropped (its scores contribute zero anyway). No raise."""
        rng = np.random.default_rng(203)
        # 4 clusters, last one zero-weight; should silently drop and use 3.
        cluster_ids = np.array([0] * 5 + [1] * 5 + [2] * 5 + [3] * 5)
        n = len(cluster_ids)
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = np.concatenate([np.ones(5), np.ones(5) * 1.5, np.ones(5) * 0.8, np.zeros(5)])
        bread = X.T @ (X * w[:, None])
        coef = np.linalg.solve(bread, X.T @ (w * y))
        resid = y - X @ coef
        vcov, dof = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w)
        assert np.all(np.isfinite(vcov)) and np.all(np.diag(vcov) > 0)
        assert np.all(np.isfinite(dof)) and np.all(dof > 0)


class TestLinearRegressionWeightedClusterHC2BM:
    """Regression for the R2 P0: `LinearRegression(vcov_type="hc2_bm")` on a
    weighted clustered fit must report Satterthwaite DOF via `_bm_dof`, not
    fall back to `n - k`. The previous skip in `fit()` left `_bm_dof = None`
    on the weighted-cluster branch, so `get_inference()` silently used the
    generic DOF and produced anti-conservative p-values / CIs."""

    def test_lr_weighted_cluster_hc2_bm_df_matches_helper(self):
        """`LinearRegression.get_inference()` df must match
        `compute_robust_vcov(..., return_dof=True)` DOF, not `n - k`."""
        from diff_diff.linalg import LinearRegression, compute_robust_vcov

        rng = np.random.default_rng(301)
        n = 40
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
        cluster_ids = np.arange(n) % 8  # 8 clusters
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.uniform(0.5, 2.0, n)

        # Helper-level: directly compute the Satterthwaite DOF
        coef = np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y))
        resid = y - X @ coef
        _, helper_dof = compute_robust_vcov(
            X,
            resid,
            cluster_ids=cluster_ids,
            vcov_type="hc2_bm",
            weights=w,
            weight_type="pweight",
            return_dof=True,
        )

        # Public surface: fit through LinearRegression. weights/cluster live
        # on the constructor (not fit()), and the constructor needs
        # include_intercept=False because X already has an intercept column.
        lr = LinearRegression(
            include_intercept=False,
            vcov_type="hc2_bm",
            cluster_ids=cluster_ids,
            weights=w,
            weight_type="pweight",
        )
        lr.fit(X, y)
        # The Satterthwaite DOF must be threaded into get_inference(), NOT the
        # n-k fallback. Helper DOF should be small (cluster-driven); fallback
        # would be ~37 (n - k = 40 - 3).
        # _bm_dof is the internal attribute populated by fit() and consumed by
        # get_inference(); also check per-coefficient inference for thread.
        assert lr._bm_dof is not None, (
            "LinearRegression must populate _bm_dof on weighted-cluster path "
            "(previous code skipped this branch)"
        )
        np.testing.assert_allclose(
            lr._bm_dof,
            helper_dof,
            atol=1e-10,
            err_msg="_bm_dof must match helper Satterthwaite DOF",
        )
        # Sanity: cluster-driven DOF should be much smaller than n - k.
        assert np.all(np.asarray(lr._bm_dof) < n - X.shape[1]), (
            f"Expected cluster-driven Satterthwaite DOF << n-k={n - X.shape[1]}, "
            f"got {lr._bm_dof}"
        )
        # End-to-end: get_inference for each coefficient yields the right df.
        for i in range(X.shape[1]):
            inf_i = lr.get_inference(index=i)
            assert (
                inf_i.df == lr._bm_dof[i]
            ), f"get_inference(index={i}).df should equal _bm_dof[{i}]"

    def test_lr_weighted_oneway_hc2_bm_df_matches_helper(self):
        """One-way weighted HC2-BM via LinearRegression must also thread DOF."""
        from diff_diff.linalg import LinearRegression, compute_robust_vcov

        rng = np.random.default_rng(302)
        n = 25
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n)])
        y = X @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n)
        w = rng.uniform(0.5, 2.0, n)

        coef = np.linalg.solve(X.T @ (X * w[:, None]), X.T @ (w * y))
        resid = y - X @ coef
        _, helper_dof = compute_robust_vcov(
            X,
            resid,
            vcov_type="hc2_bm",
            weights=w,
            weight_type="pweight",
            return_dof=True,
        )

        lr = LinearRegression(
            include_intercept=False,
            vcov_type="hc2_bm",
            weights=w,
            weight_type="pweight",
        )
        lr.fit(X, y)
        assert lr._bm_dof is not None
        np.testing.assert_allclose(
            lr._bm_dof,
            helper_dof,
            atol=1e-10,
            err_msg="One-way LinearRegression _bm_dof must match helper Satterthwaite DOF",
        )


class TestWLSCR2FEDoFNoiseGuard:
    """Regression for R6 P1: weighted CR2-BM Satterthwaite DOF for high-
    leverage FE-dummy contrasts (where the contrast vector projects to
    near-zero on the design) is at the float64 noise floor and disagrees
    with R's clubSandwich by 15-30% due to BLAS reduction-order differences.
    The mitigation is to detect the noise floor and return NaN with a
    UserWarning rather than ship arbitrarily-different small-sample DOF.

    Other (non-noise-floor) coefficients still match clubSandwich exactly.
    """

    def test_did_absorbed_fe_nan_guards_treated_unit_dummies(self, goldens):
        """The weighted_did_absorbed_fe scenario: treated-unit dummies are at
        the noise floor (vcov is fine, but DOF formula is noise/noise). NaN-
        guard fires; non-noise-floor coefs still match R."""
        import warnings

        d = goldens["weighted_did_absorbed_fe"]
        n = len(d["y"])
        unit_int = np.asarray(d["unit"])
        period_int = np.asarray(d["period"])
        treat_post = np.asarray(d["treat_post"])
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster = unit_int

        cols = [np.ones(n), treat_post.astype(float)]
        col_names = ["(Intercept)", "treat_post"]
        for u in sorted(np.unique(unit_int))[1:]:
            cols.append((unit_int == u).astype(float))
            col_names.append(f"unit{u}")
        for p in sorted(np.unique(period_int))[1:]:
            cols.append((period_int == p).astype(float))
            col_names.append(f"period{p}")
        X = np.column_stack(cols)

        bread = X.T @ (X * w[:, None])
        coef = np.linalg.solve(bread, X.T @ (w * y))
        resid = y - X @ coef

        with warnings.catch_warnings(record=True) as ws:
            warnings.simplefilter("always")
            vcov, dof = _compute_cr2_bm(X, resid, cluster, bread, weights=w)

        # VCOV matches R at machine precision.
        r_vcov = np.asarray(d["vcov_cr2"]).reshape(d["vcov_cr2_shape"])
        np.testing.assert_allclose(vcov, r_vcov, atol=1e-10)

        # DOF: non-noise-floor coefs match R at 1e-10; treated-unit dummies
        # (unit2/3/4 — indices 2, 3, 4 in column ordering) are NaN.
        r_dof = np.asarray(d["dof_cr2"])
        nan_mask = np.isnan(dof)
        # The 3 treated-unit dummies should be flagged as noise-floor.
        treated_unit_idx = [col_names.index(f"unit{u}") for u in [2, 3, 4]]
        for idx in treated_unit_idx:
            assert nan_mask[idx], f"Expected NaN at noise-floor index {idx} ({col_names[idx]})"
        # All other coefs match R exactly.
        finite_idx = ~nan_mask
        np.testing.assert_allclose(
            dof[finite_idx],
            r_dof[finite_idx],
            atol=1e-10,
            err_msg="Non-noise-floor coefficients must match clubSandwich at atol=1e-10",
        )
        # Warning fired.
        noise_warnings = [w for w in ws if "noise floor" in str(w.message)]
        assert (
            len(noise_warnings) == 1
        ), f"Expected 1 noise-floor UserWarning, got {len(noise_warnings)}"


class TestWLSCR2SubpopulationInvariance:
    """Regression for CI codex R8 P0 (round 2): weighted clustered CR2 must
    be subpopulation-invariant on mixed-zero clusters. Previously, zero-weight
    rows inside positive-weight clusters still entered the CR2 adjustment
    matrices (H_gg, G_g, A_g, bias_term), silently changing vcov/DOF. The fix
    physically filters `weights > 0` rows before building any per-cluster
    matrix. This test pins parity between (a) computing on the full design
    with zero-weight padding and (b) computing on the physically dropped
    positive-weight subset."""

    def test_per_coefficient_dof_invariant_to_zero_weight_padding(self, goldens):
        """`_compute_cr2_bm` on (X, w) with mixed-zero padding rows must yield
        identical vcov + per-coefficient DOF to running on (X_pos, w_pos)
        with zero-weight rows physically removed."""
        # Use a simple weighted clustered fixture.
        rng = np.random.default_rng(401)
        n_clusters, cluster_size = 5, 6
        n_pos = n_clusters * cluster_size
        X_pos = np.column_stack([np.ones(n_pos), rng.normal(size=n_pos), rng.normal(size=n_pos)])
        y_pos = X_pos @ np.array([1.0, 0.5, 0.3]) + rng.normal(0, 0.5, n_pos)
        w_pos = rng.uniform(0.5, 2.5, n_pos)
        cluster_pos = np.repeat(np.arange(n_clusters), cluster_size)

        # Pad with 8 zero-weight rows, sprinkled across existing clusters.
        n_pad = 8
        X_pad = rng.normal(size=(n_pad, 3))
        X_pad[:, 0] = 1.0  # intercept
        y_pad = rng.normal(size=n_pad)
        w_pad = np.zeros(n_pad)
        cluster_pad = rng.integers(0, n_clusters, size=n_pad)
        # Interleave (shuffle order to verify it's not order-dependent).
        order = rng.permutation(n_pos + n_pad)
        X_full = np.vstack([X_pos, X_pad])[order]
        y_full = np.concatenate([y_pos, y_pad])[order]
        w_full = np.concatenate([w_pos, w_pad])[order]
        cluster_full = np.concatenate([cluster_pos, cluster_pad])[order]

        # Refit WLS on both designs. Bread is invariant to zero-weight rows.
        bread = X_full.T @ (X_full * w_full[:, None])
        coef = np.linalg.solve(bread, X_full.T @ (w_full * y_full))
        resid_full = y_full - X_full @ coef
        resid_pos = y_pos - X_pos @ coef
        bread_pos = X_pos.T @ (X_pos * w_pos[:, None])
        np.testing.assert_allclose(bread, bread_pos, atol=1e-12)

        vcov_full, dof_full = _compute_cr2_bm(
            X_full, resid_full, cluster_full, bread, weights=w_full
        )
        vcov_pos, dof_pos = _compute_cr2_bm(X_pos, resid_pos, cluster_pos, bread_pos, weights=w_pos)
        np.testing.assert_allclose(
            vcov_full,
            vcov_pos,
            atol=1e-12,
            err_msg="vcov must be invariant to zero-weight padding rows",
        )
        np.testing.assert_allclose(
            dof_full,
            dof_pos,
            atol=1e-12,
            err_msg="DOF must be invariant to zero-weight padding rows",
        )

    def test_contrast_dof_invariant_to_zero_weight_padding(self, goldens):
        """Same invariance for `_compute_cr2_bm_contrast_dof`."""
        rng = np.random.default_rng(402)
        n_clusters, cluster_size = 5, 6
        n_pos = n_clusters * cluster_size
        X_pos = np.column_stack([np.ones(n_pos), rng.normal(size=n_pos), rng.normal(size=n_pos)])
        w_pos = rng.uniform(0.5, 2.5, n_pos)
        cluster_pos = np.repeat(np.arange(n_clusters), cluster_size)

        n_pad = 6
        X_pad = rng.normal(size=(n_pad, 3))
        X_pad[:, 0] = 1.0
        w_pad = np.zeros(n_pad)
        cluster_pad = rng.integers(0, n_clusters, size=n_pad)
        order = rng.permutation(n_pos + n_pad)
        X_full = np.vstack([X_pos, X_pad])[order]
        w_full = np.concatenate([w_pos, w_pad])[order]
        cluster_full = np.concatenate([cluster_pos, cluster_pad])[order]

        bread = X_full.T @ (X_full * w_full[:, None])
        bread_pos = X_pos.T @ (X_pos * w_pos[:, None])

        # Compound contrast: average of x and z coefficients (indices 1 and 2).
        contrasts = np.array([[0.0], [0.5], [0.5]])
        dof_full = _compute_cr2_bm_contrast_dof(
            X_full, cluster_full, bread, contrasts, weights=w_full
        )
        dof_pos = _compute_cr2_bm_contrast_dof(
            X_pos, cluster_pos, bread_pos, contrasts, weights=w_pos
        )
        np.testing.assert_allclose(
            dof_full,
            dof_pos,
            atol=1e-12,
            err_msg="Contrast DOF must be invariant to zero-weight padding rows",
        )


class TestWLSCR2SingleContrastNoiseFloor:
    """Regression for CI codex R8 P2: the noise-floor NaN-guard must trigger
    on single-contrast calls to `_compute_cr2_bm_contrast_dof` too, not just
    batched per-coefficient calls. The earlier batch-relative-only criterion
    couldn't catch single contrasts because `max|P|_overall == contrast's own
    scale` for m=1. The absolute (eps × n × k × bread_scale)² floor closes
    the gap."""

    def test_single_fe_dummy_contrast_triggers_nan_guard(self, goldens):
        """A direct single-contrast call for a high-leverage FE-dummy unit
        dummy must return NaN, not a BLAS-dependent finite DOF."""
        import warnings

        d = goldens["weighted_did_absorbed_fe"]
        n = len(d["y"])
        unit_int = np.asarray(d["unit"])
        period_int = np.asarray(d["period"])
        treat_post = np.asarray(d["treat_post"])
        np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster = unit_int

        cols = [np.ones(n), treat_post.astype(float)]
        col_names = ["(Intercept)", "treat_post"]
        for u in sorted(np.unique(unit_int))[1:]:
            cols.append((unit_int == u).astype(float))
            col_names.append(f"unit{u}")
        for p in sorted(np.unique(period_int))[1:]:
            cols.append((period_int == p).astype(float))
            col_names.append(f"period{p}")
        X = np.column_stack(cols)
        bread = X.T @ (X * w[:, None])

        # Single contrast = unit2 dummy (the noise-floor case from the
        # weighted_did_absorbed_fe fixture).
        unit2_idx = col_names.index("unit2")
        c_unit2 = np.zeros(X.shape[1])
        c_unit2[unit2_idx] = 1.0
        contrasts = c_unit2[:, None]  # (k, 1) — single contrast

        with warnings.catch_warnings(record=True) as ws:
            warnings.simplefilter("always")
            dof = _compute_cr2_bm_contrast_dof(X, cluster, bread, contrasts, weights=w)
        assert np.isnan(dof[0]), (
            "Single weighted FE-dummy contrast must trigger absolute noise-"
            "floor guard; got finite DOF instead."
        )
        noise_warnings = [w for w in ws if "noise floor" in str(w.message)]
        assert (
            len(noise_warnings) == 1
        ), f"Expected 1 noise-floor UserWarning, got {len(noise_warnings)}"

    def test_single_non_noise_contrast_returns_finite_dof(self, goldens):
        """Sanity check: a single contrast that's NOT at the noise floor
        (e.g., treat_post) must return finite DOF, not NaN."""
        d = goldens["weighted_did_absorbed_fe"]
        n = len(d["y"])
        unit_int = np.asarray(d["unit"])
        period_int = np.asarray(d["period"])
        treat_post = np.asarray(d["treat_post"])
        np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster = unit_int

        cols = [np.ones(n), treat_post.astype(float)]
        col_names = ["(Intercept)", "treat_post"]
        for u in sorted(np.unique(unit_int))[1:]:
            cols.append((unit_int == u).astype(float))
            col_names.append(f"unit{u}")
        for p in sorted(np.unique(period_int))[1:]:
            cols.append((period_int == p).astype(float))
            col_names.append(f"period{p}")
        X = np.column_stack(cols)
        bread = X.T @ (X * w[:, None])

        # Single contrast = treat_post (NOT a noise-floor contrast).
        treat_post_idx = col_names.index("treat_post")
        c_treat = np.zeros(X.shape[1])
        c_treat[treat_post_idx] = 1.0
        contrasts = c_treat[:, None]

        dof = _compute_cr2_bm_contrast_dof(X, cluster, bread, contrasts, weights=w)
        assert np.isfinite(dof[0]) and dof[0] > 0, (
            f"Single non-noise contrast should return finite positive DOF; " f"got {dof[0]}"
        )
        # Sanity: should match R's df_satt for treat_post (3.754) at atol=1e-10.
        r_dof_treat_post = float(np.asarray(d["dof_cr2"])[treat_post_idx])
        np.testing.assert_allclose(dof[0], r_dof_treat_post, atol=1e-10)


class TestLinearRegressionFENanGuardEndToEnd:
    """Regression for R7 P0: NaN BM DOF from the noise-floor guard must
    propagate to NaN inference fields in `LinearRegression.get_inference()`,
    not fall through to normal-theory inference (which would silently emit
    huge t-stats and zero-width CIs for the affected coefficients)."""

    def test_did_absorbed_fe_lr_inference_nan_for_guarded_coefs(self, goldens):
        from diff_diff.linalg import LinearRegression

        d = goldens["weighted_did_absorbed_fe"]
        n = len(d["y"])
        unit_int = np.asarray(d["unit"])
        period_int = np.asarray(d["period"])
        treat_post = np.asarray(d["treat_post"])
        y = np.asarray(d["y"])
        w = np.asarray(d["weights"])
        cluster = unit_int

        cols = [np.ones(n), treat_post.astype(float)]
        col_names = ["(Intercept)", "treat_post"]
        for u in sorted(np.unique(unit_int))[1:]:
            cols.append((unit_int == u).astype(float))
            col_names.append(f"unit{u}")
        for p in sorted(np.unique(period_int))[1:]:
            cols.append((period_int == p).astype(float))
            col_names.append(f"period{p}")
        X = np.column_stack(cols)

        lr = LinearRegression(
            include_intercept=False,
            vcov_type="hc2_bm",
            cluster_ids=cluster,
            weights=w,
            weight_type="pweight",
        )
        lr.fit(X, y)

        # Guarded coefficients: treated-unit dummies (unit2/3/4).
        treated_unit_idx = [col_names.index(f"unit{u}") for u in [2, 3, 4]]
        for i in treated_unit_idx:
            inf_i = lr.get_inference(index=i)
            assert np.isnan(inf_i.t_stat), (
                f"NaN guard must produce NaN t_stat at index {i} ({col_names[i]}); "
                f"got {inf_i.t_stat}"
            )
            assert np.isnan(
                inf_i.p_value
            ), f"NaN guard must produce NaN p_value at index {i}; got {inf_i.p_value}"
            assert all(
                np.isnan(b) for b in inf_i.conf_int
            ), f"NaN guard must produce NaN conf_int at index {i}; got {inf_i.conf_int}"
            # SE and coefficient remain valid (vcov matches at machine precision).
            assert np.isfinite(inf_i.se) and inf_i.se > 0
            assert np.isfinite(inf_i.coefficient)

        # Non-guarded coefficients still emit finite inference.
        non_guarded_idx = [i for i in range(X.shape[1]) if i not in treated_unit_idx]
        for i in non_guarded_idx:
            inf_i = lr.get_inference(index=i)
            assert np.isfinite(inf_i.t_stat), (
                f"Non-guarded coef index {i} ({col_names[i]}) should have finite t_stat; "
                f"got {inf_i.t_stat}"
            )
            assert np.isfinite(inf_i.p_value)
            assert all(np.isfinite(b) for b in inf_i.conf_int)


class TestUnweightedRegressionStillBitEqual:
    """Regression safety: existing unweighted goldens must still match bit-equal."""

    def test_unweighted_balanced_small_unchanged(self, goldens):
        d = goldens["balanced_small"]
        X = np.column_stack([np.ones(len(d["x"])), np.asarray(d["x"])])
        y = np.asarray(d["y"])
        cluster_ids = np.asarray(d["cluster"])
        bread = X.T @ X
        coef = np.linalg.solve(bread, X.T @ y)
        resid = y - X @ coef
        r_vcov = np.asarray(d["vcov_cr2"]).reshape(d["vcov_shape"])
        r_dof = np.asarray(d["dof_bm"])

        # Explicitly pass weights=None (the regression-safety path).
        vcov, dof = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=None)
        # Pin at atol=1e-12 (existing goldens at atol=1e-10; bit-equal at 1e-14).
        np.testing.assert_allclose(
            vcov,
            r_vcov,
            atol=1e-12,
            err_msg="unweighted CR2-BM vcov must remain bit-equal to prior",
        )
        np.testing.assert_allclose(
            dof, r_dof, atol=1e-12, err_msg="unweighted CR2-BM DOF must remain bit-equal to prior"
        )


class TestDOFFormulaDualPathEquivalence:
    """Dual-path equivalence: the existing simple (tr B)^2 / tr(B^2) formula
    must agree with the new full P_array form at the unweighted limit.

    This addresses the reviewer's MEDIUM #2 concern about parallel DOF formulas
    drifting over time. The new P_array form is only used when weights is
    not None; the unweighted path keeps the simple formula. At the unweighted
    limit (weights=ones), both formulas mathematically reduce to the same
    Pustejovsky-Tipton (2018) Theorem 1 value, and this test pins their
    numerical agreement.
    """

    def test_unweighted_dof_path_equivalence(self, goldens):
        d = goldens["unbalanced_medium"]
        X = np.column_stack([np.ones(len(d["x"])), np.asarray(d["x"])])
        y = np.asarray(d["y"])
        cluster_ids = np.asarray(d["cluster"])
        bread = X.T @ X
        coef = np.linalg.solve(bread, X.T @ y)
        resid = y - X @ coef

        # Path A: unweighted simple formula (weights=None).
        _, dof_simple = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=None)
        # Path B: P_array formula via weights=ones (forces weighted branch).
        n = X.shape[0]
        w_ones = np.ones(n)
        _, dof_p_array = _compute_cr2_bm(X, resid, cluster_ids, bread, weights=w_ones)
        np.testing.assert_allclose(
            dof_simple,
            dof_p_array,
            atol=1e-12,
            err_msg=(
                "At the unweighted limit, the simple (tr B)^2/tr(B^2) and the "
                "full P_array formulas must agree. Drift indicates the P_array "
                "implementation has a bug (likely chol direction or M_U "
                "normalization)."
            ),
        )

    def test_unweighted_one_way_dof_path_equivalence(self, goldens):
        d = goldens["singletons_present"]
        X = np.column_stack([np.ones(len(d["x"])), np.asarray(d["x"])])
        y = np.asarray(d["y"])
        bread = X.T @ X
        coef = np.linalg.solve(bread, X.T @ y)
        y - X @ coef
        h_diag = _compute_hat_diagonals(X, bread, weights=None)

        # Path A: unweighted simple formula.
        dof_simple = _compute_bm_dof_oneway(X, bread, h_diag, weights=None)
        # Path B: P_array via singleton-cluster reduction with weights=ones.
        n = X.shape[0]
        w_ones = np.ones(n)
        h_diag_ones = _compute_hat_diagonals(X, bread, weights=w_ones)
        dof_p_array = _compute_bm_dof_oneway(X, bread, h_diag_ones, weights=w_ones)
        np.testing.assert_allclose(
            dof_simple,
            dof_p_array,
            atol=1e-12,
            err_msg=(
                "One-way unweighted DOF: simple formula must agree with the "
                "P_array singleton-cluster reduction at the unweighted limit."
            ),
        )
