"""Cross-validation: diff-diff survey variance vs R using real survey data.

Complements test_survey_r_crossvalidation.py (synthetic data) with real-world
survey datasets to validate numerical agreement between diff-diff and R's
authoritative ``survey`` package.

Suite A: API dataset — TSL variance (California school accountability, PSAA 1999)
Suite B: NHANES — TSL with strata+PSU (ACA young adult coverage provision)
Suite C: RECS 2020 — JK1 replicate weight variance

NOTE: Suite A uses observational program assignment (Governor's Performance Award
eligibility) as the treatment indicator. Suite B uses the ACA dependent coverage
provision (age 19-25 vs 27-34). These are real policies but the DiD estimates
should be interpreted with caution. The primary purpose is validating that
diff-diff's survey variance matches R, not causal claims.

JSON fixtures generated by:
  benchmarks/R/benchmark_realdata_api.R
  benchmarks/R/benchmark_realdata_nhanes.R
  benchmarks/R/benchmark_realdata_recs.R
"""

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from diff_diff import DifferenceInDifferences, SurveyDesign
from diff_diff.twfe import TwoWayFixedEffects

BENCHMARK_DIR = Path(__file__).parent.parent / "benchmarks" / "data" / "real"

# Golden value files
API_RESULTS_FILE = BENCHMARK_DIR / "api_realdata_golden.json"
NHANES_RESULTS_FILE = BENCHMARK_DIR / "nhanes_realdata_golden.json"
RECS_RESULTS_FILE = BENCHMARK_DIR / "recs_realdata_golden.json"

# Tight tolerances — observed gaps are < 1e-10, so 1e-8 guards against
# regressions without being so loose that a real bug could slip through.
ATT_RTOL = 1e-8
ATT_ATOL = 1e-8
SE_RTOL = 1e-8
SE_ATOL = 1e-8
CI_RTOL = 1e-8
CI_ATOL = 1e-8

# Replicate weight methods: still tight (observed gaps < 1e-10)
REP_SE_RTOL = 1e-8
REP_SE_ATOL = 1e-8
REP_CI_RTOL = 1e-8
REP_CI_ATOL = 1e-8

# Mark for easy filtering: pytest -m realdata
pytestmark = pytest.mark.realdata


def _assert_close(py_val, r_val, rtol, atol, label):
    """Assert values are close with informative error message."""
    if np.isnan(r_val):
        return
    diff = abs(py_val - r_val)
    threshold = max(atol, rtol * abs(r_val))
    assert diff < threshold, (
        f"{label}: Python={py_val:.6f}, R={r_val:.6f}, "
        f"diff={diff:.6f}, threshold={threshold:.6f} "
        f"(rtol={rtol}, atol={atol})"
    )


def _load_embedded_data(results, key="_data"):
    """Load the embedded dataset from golden JSON."""
    data = results[key]
    df = pd.DataFrame({col: vals for col, vals in data.items()})
    return df


# ============================================================================
# Fixtures
# ============================================================================


@pytest.fixture(scope="module")
def api_results():
    """Load pre-computed R golden values for API dataset."""
    if not API_RESULTS_FILE.exists():
        pytest.skip(
            f"API golden values not found: {API_RESULTS_FILE}. "
            "Run: Rscript benchmarks/R/benchmark_realdata_api.R"
        )
    with open(API_RESULTS_FILE) as f:
        return json.load(f)


@pytest.fixture(scope="module")
def nhanes_results():
    """Load pre-computed R golden values for NHANES dataset."""
    if not NHANES_RESULTS_FILE.exists():
        pytest.skip(
            f"NHANES golden values not found: {NHANES_RESULTS_FILE}. "
            "Run: python benchmarks/scripts/download_nhanes.py && "
            "Rscript benchmarks/R/benchmark_realdata_nhanes.R"
        )
    with open(NHANES_RESULTS_FILE) as f:
        return json.load(f)


@pytest.fixture(scope="module")
def recs_results():
    """Load pre-computed R golden values for RECS dataset."""
    if not RECS_RESULTS_FILE.exists():
        pytest.skip(
            f"RECS golden values not found: {RECS_RESULTS_FILE}. "
            "Run: python benchmarks/scripts/download_recs.py && "
            "Rscript benchmarks/R/benchmark_realdata_recs.R"
        )
    with open(RECS_RESULTS_FILE) as f:
        return json.load(f)


# ============================================================================
# Suite A: API Dataset — California School Accountability (TSL Variance)
# ============================================================================


class TestApiSchoolAccountability:
    """Cross-validate survey DiD on real API data vs R survey::svyglm.

    Policy: California PSAA (1999) — Governor's Performance Award eligibility
    as treatment, api99/api00 as pre/post scores. Stratified by school type
    (E/M/H) with FPC and probability weights from apistrat sample.
    """

    def _load_api_data(self, api_results):
        """Load embedded API panel data."""
        df = _load_embedded_data(api_results)
        # Ensure dtypes
        for col in ["school_id", "period", "post", "treated"]:
            df[col] = df[col].astype(int)
        for col in ["outcome", "pw", "fpc", "meals", "ell"]:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
        return df

    def test_a1_strata_fpc_weights(self, api_results):
        """A1: Full TSL with strata + FPC + weights."""
        r = api_results["a1_strata_fpc_weights"]
        data = self._load_api_data(api_results)

        sd = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A1 ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A1 SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"A1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "A1 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "A1 CI upper")

    def test_a2_strata_no_fpc(self, api_results):
        """A2: TSL with strata + weights, no FPC."""
        r = api_results["a2_strata_weights"]
        data = self._load_api_data(api_results)

        sd = SurveyDesign(weights="pw", strata="stype")
        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A2 ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A2 SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"A2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "A2 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "A2 CI upper")

    def test_a3_weights_only(self, api_results):
        """A3: Weighted OLS baseline (no strata or FPC)."""
        r = api_results["a3_weights_only"]
        data = self._load_api_data(api_results)

        sd = SurveyDesign(weights="pw")
        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A3 ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A3 SE")
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "A3 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "A3 CI upper")

    def test_a4_twfe(self, api_results):
        """A4: TWFE plumbing with strata + FPC + weights."""
        r = api_results["a4_twfe"]
        data = self._load_api_data(api_results)

        sd = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
        est = TwoWayFixedEffects()
        result = est.fit(
            data, "outcome", "treated", "period", "school_id",
            survey_design=sd,
        )

        # ATT should match R's pooled DiD. SE differs because TWFE absorbs
        # unit fixed effects (within-transform), giving a smaller residual
        # variance than R's svyglm(outcome ~ treated * post).
        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A4 TWFE ATT")
        assert np.isfinite(result.se), "A4 TWFE SE should be finite"
        assert result.se > 0, "A4 TWFE SE should be positive"

    def test_a5_subpopulation_elementary(self, api_results):
        """A5: Subpopulation — elementary schools only."""
        r = api_results["a5_subpop_elementary"]
        data = self._load_api_data(api_results)

        # Use subpopulation approach: zero out weights for non-elementary
        sd_full = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
        sd_sub, data_sub = sd_full.subpopulation(
            data, lambda df: df["stype"] == "E",
        )

        est = DifferenceInDifferences()
        result = est.fit(
            data_sub, "outcome", "treated", "post", survey_design=sd_sub,
        )

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A5 subpop ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A5 subpop SE")
        # df differs by design: subpopulation() preserves all 3 strata (df=397),
        # while R's subset() drops to 1 stratum (df=199). ATT and SE match.

    def test_a6_covariates(self, api_results):
        """A6: Covariate-adjusted (meals, ell) with strata + FPC."""
        r = api_results["a6_covariates"]
        data = self._load_api_data(api_results)
        # Drop rows with missing covariates (matching R)
        data = data.dropna(subset=["meals", "ell"]).reset_index(drop=True)

        sd = SurveyDesign(weights="pw", strata="stype", fpc="fpc")
        est = DifferenceInDifferences()
        result = est.fit(
            data, "outcome", "treated", "post",
            covariates=["meals", "ell"],
            survey_design=sd,
        )

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A6 cov ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "A6 cov SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"A6 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "A6 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "A6 CI upper")

    def test_a7_fay_brr_replicates(self, api_results):
        """A7: Fay's BRR replicate weights from real stratified design."""
        r = api_results["a7_fay_brr"]
        data = self._load_api_data(api_results)

        # Add replicate weight columns from R's exported matrix
        rep_weights = r["replicate_weights"]
        rep_cols = sorted(rep_weights.keys())
        for col_name in rep_cols:
            data[col_name] = [float(v) for v in rep_weights[col_name]]

        sd = SurveyDesign(
            weights="pw",
            replicate_weights=rep_cols,
            replicate_method="Fay",
            fay_rho=r["fay_rho"],
        )

        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "A7 Fay ATT")
        _assert_close(result.se, r["se"], REP_SE_RTOL, REP_SE_ATOL, "A7 Fay SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"A7 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL,
                       "A7 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL,
                       "A7 CI upper")

    def test_fpc_reduces_se(self, api_results):
        """FPC should reduce SE vs no-FPC (A1 SE < A2 SE)."""
        r_fpc = api_results["a1_strata_fpc_weights"]
        r_nofpc = api_results["a2_strata_weights"]
        assert r_fpc["se"] < r_nofpc["se"], (
            f"FPC SE ({r_fpc['se']:.4f}) should be < no-FPC SE ({r_nofpc['se']:.4f})"
        )


# ============================================================================
# Suite B: NHANES — ACA Young Adult Coverage Provision (TSL + PSU)
# ============================================================================


class TestNhanesAcaCoverage:
    """Cross-validate survey DiD on real NHANES data vs R survey::svyglm.

    Policy: ACA dependent coverage provision (effective Sept 2010). Adults
    aged 19-25 could stay on parents' insurance (treatment) vs ages 27-34
    (control). Pre = NHANES 2007-2008, Post = NHANES 2015-2016.

    Outcome is binary (insured: 0/1) modeled as linear probability model,
    consistent with R's svyglm(family=gaussian). ATT is interpretable as a
    percentage-point change in insurance coverage.
    """

    def _load_nhanes_data(self, nhanes_results):
        """Load embedded NHANES panel data."""
        df = _load_embedded_data(nhanes_results)
        for col in ["unit_id", "period", "post", "treated"]:
            if col in df.columns:
                df[col] = df[col].astype(int)
        for col in ["outcome", "WTMEC2YR", "INDFMPIR"]:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
        # Ensure strata and PSU are string/categorical for SurveyDesign
        for col in ["SDMVSTRA", "SDMVPSU"]:
            if col in df.columns:
                df[col] = df[col].astype(int)
        return df

    def test_b1_strata_psu_weights(self, nhanes_results):
        """B1: Full TSL with strata + PSU + weights (nest=True)."""
        r = nhanes_results["b1_strata_psu_weights"]
        data = self._load_nhanes_data(nhanes_results)

        sd = SurveyDesign(
            weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
        )
        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B1 ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B1 SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"B1 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "B1 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "B1 CI upper")

    def test_b2_covariates(self, nhanes_results):
        """B2: Covariate-adjusted TSL with strata + PSU."""
        r = nhanes_results["b2_covariates"]
        data = self._load_nhanes_data(nhanes_results)
        data = data.dropna(subset=r.get("covariates", [])).reset_index(drop=True)

        sd = SurveyDesign(
            weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
        )
        est = DifferenceInDifferences()
        result = est.fit(
            data, "outcome", "treated", "post",
            covariates=r.get("covariates", []),
            survey_design=sd,
        )

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B2 cov ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B2 cov SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"B2 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "B2 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "B2 CI upper")

    def test_b3_weights_only(self, nhanes_results):
        """B3: Weighted OLS (no clustering), baseline comparison."""
        r = nhanes_results["b3_weights_only"]
        data = self._load_nhanes_data(nhanes_results)

        sd = SurveyDesign(weights="WTMEC2YR")
        est = DifferenceInDifferences()
        result = est.fit(data, "outcome", "treated", "post", survey_design=sd)

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B3 ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B3 SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"B3 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "B3 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "B3 CI upper")

    def test_b4_subpop_female(self, nhanes_results):
        """B4: Subpopulation — female respondents only."""
        r = nhanes_results["b4_subpop_female"]
        data = self._load_nhanes_data(nhanes_results)

        sd_full = SurveyDesign(
            weights="WTMEC2YR", strata="SDMVSTRA", psu="SDMVPSU", nest=True,
        )
        # Subpopulation: female (RIAGENDR == 2)
        sd_sub, data_sub = sd_full.subpopulation(
            data, lambda df: df["RIAGENDR"] == 2,
        )
        est = DifferenceInDifferences()
        result = est.fit(
            data_sub, "outcome", "treated", "post", survey_design=sd_sub,
        )

        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL, "B4 subpop ATT")
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "B4 subpop SE")
        assert result.survey_metadata.df_survey == r["df"], (
            f"B4 df: Python={result.survey_metadata.df_survey}, R={r['df']}"
        )
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       "B4 CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       "B4 CI upper")

    # B5 (CallawaySantAnna RC-DiD) was removed: R's did::att_gt cannot produce
    # golden values for a 2-period repeated cross-section, and the time-scale
    # mismatch between the R script (period_cs=1/2) and the Python data
    # (period=0/1) made the dormant path untestable. CallawaySantAnna survey
    # variance is validated in the synthetic-data suite instead.


# ============================================================================
# Suite C: RECS 2020 — JK1 Replicate Weight Variance
# ============================================================================


class TestRecsReplicateWeights:
    """Cross-validate JK1 replicate weight variance on real RECS 2020 data.

    RECS is a federal energy survey by EIA with 60 pre-computed JK1 replicate
    weights. This validates the replicate weight variance path with real
    federal survey weights using a weighted regression (not a full DiD).
    """

    def _load_recs_data(self, recs_results):
        """Load embedded RECS subset data."""
        df = _load_embedded_data(recs_results)
        # Ensure numeric columns
        for col in ["TOTALBTU", "DOLLAREL", "NWEIGHT"]:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
        # Identify replicate weight columns
        rep_cols = sorted([c for c in df.columns if c.startswith("NWEIGHT")
                           and c != "NWEIGHT"])
        for col in rep_cols:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        return df, rep_cols

    def test_c1_simple_regression(self, recs_results):
        """C1: Simple regression TOTALBTU ~ KOWNRENT with JK1 replicate SEs."""
        r = recs_results["c1_simple"]
        data, rep_cols = self._load_recs_data(recs_results)
        data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(
            drop=True
        )

        # Build design matrix: intercept + KOWNRENT
        from diff_diff.linalg import LinearRegression as LR

        X = data[["KOWNRENT"]].values.astype(float)
        y = data["TOTALBTU"].values.astype(float)

        sd = SurveyDesign(
            weights="NWEIGHT",
            replicate_weights=rep_cols,
            replicate_method="JK1",
        )
        resolved = sd.resolve(data)

        reg = LR(include_intercept=True, survey_design=resolved)
        reg.fit(X, y)

        # KOWNRENT coefficient is at index 1 (intercept=0, KOWNRENT=1)
        py_coef = reg.coefficients_[1]
        py_se = np.sqrt(reg.vcov_[1, 1])

        _assert_close(py_coef, r["coef"], ATT_RTOL, ATT_ATOL, "C1 coef")
        _assert_close(py_se, r["se"], REP_SE_RTOL, REP_SE_ATOL, "C1 SE")
        assert reg.survey_df_ == r["df"], (
            f"C1 df: Python={reg.survey_df_}, R={r['df']}"
        )
        # Compute CI from coef, SE, and survey df
        from scipy.stats import t as t_dist
        t_crit = t_dist.ppf(0.975, reg.survey_df_)
        py_ci_lower = py_coef - t_crit * py_se
        py_ci_upper = py_coef + t_crit * py_se
        _assert_close(py_ci_lower, r["ci_lower"], REP_CI_RTOL, REP_CI_ATOL,
                       "C1 CI lower")
        _assert_close(py_ci_upper, r["ci_upper"], REP_CI_RTOL, REP_CI_ATOL,
                       "C1 CI upper")

    def test_c2_full_regression(self, recs_results):
        """C2: Full regression TOTALBTU ~ KOWNRENT + TYPEHUQ + REGIONC."""
        r = recs_results["c2_full"]
        data, rep_cols = self._load_recs_data(recs_results)
        data = data.dropna(
            subset=["TOTALBTU", "KOWNRENT", "TYPEHUQ", "REGIONC", "NWEIGHT"]
        ).reset_index(drop=True)

        from diff_diff.linalg import LinearRegression as LR

        # Create dummy variables for categorical predictors
        X_parts = [data[["KOWNRENT"]].values.astype(float)]
        # TYPEHUQ dummies (drop first level)
        typehuq_dummies = pd.get_dummies(data["TYPEHUQ"], prefix="TYPEHUQ", drop_first=True)
        X_parts.append(typehuq_dummies.values.astype(float))
        # REGIONC dummies (drop first level)
        regionc_dummies = pd.get_dummies(data["REGIONC"], prefix="REGIONC", drop_first=True)
        X_parts.append(regionc_dummies.values.astype(float))

        X = np.column_stack(X_parts)
        y = data["TOTALBTU"].values.astype(float)

        sd = SurveyDesign(
            weights="NWEIGHT",
            replicate_weights=rep_cols,
            replicate_method="JK1",
        )
        resolved = sd.resolve(data)

        reg = LR(include_intercept=True, survey_design=resolved)
        reg.fit(X, y)

        # KOWNRENT coefficient is at index 1
        py_coef = reg.coefficients_[1]
        py_se = np.sqrt(reg.vcov_[1, 1])

        _assert_close(py_coef, r["coef_kownrent"], ATT_RTOL, ATT_ATOL,
                       "C2 KOWNRENT coef")
        _assert_close(py_se, r["se_kownrent"], REP_SE_RTOL, REP_SE_ATOL,
                       "C2 KOWNRENT SE")
        assert reg.survey_df_ == r["df"], (
            f"C2 df: Python={reg.survey_df_}, R={r['df']}"
        )
        from scipy.stats import t as t_dist
        t_crit = t_dist.ppf(0.975, reg.survey_df_)
        py_ci_lower = py_coef - t_crit * py_se
        py_ci_upper = py_coef + t_crit * py_se
        _assert_close(py_ci_lower, r["ci_lower_kownrent"], REP_CI_RTOL, REP_CI_ATOL,
                       "C2 CI lower")
        _assert_close(py_ci_upper, r["ci_upper_kownrent"], REP_CI_RTOL, REP_CI_ATOL,
                       "C2 CI upper")

    def test_c3_deff_diagnostics(self, recs_results):
        """C3: DEFF diagnostics from real JK1 replicate SEs.

        DEFF = (survey_SE / naive_SE)^2. The naive SE baseline differs between
        R (weighted lm) and Python (compute_deff_diagnostics SRS baseline), so
        DEFF values aren't directly comparable. This test verifies DEFFs are
        finite, positive, and > 1 (expected for a complex survey design).
        """
        data, rep_cols = self._load_recs_data(recs_results)
        data = data.dropna(subset=["TOTALBTU", "KOWNRENT", "NWEIGHT"]).reset_index(
            drop=True
        )

        from diff_diff.linalg import LinearRegression as LR

        X = data[["KOWNRENT"]].values.astype(float)
        y = data["TOTALBTU"].values.astype(float)

        sd = SurveyDesign(
            weights="NWEIGHT",
            replicate_weights=rep_cols,
            replicate_method="JK1",
        )
        resolved = sd.resolve(data)

        reg = LR(include_intercept=True, survey_design=resolved)
        reg.fit(X, y)

        # Compute DEFF diagnostics
        from diff_diff.survey import compute_deff_diagnostics

        deff = compute_deff_diagnostics(
            X=np.column_stack([np.ones(len(y)), X]),
            residuals=y - reg.predict(X),
            survey_vcov=reg.vcov_,
            weights=resolved.weights,
        )

        # DEFF should be finite and positive for all coefficients
        assert all(np.isfinite(deff.deff)), "DEFF values should be finite"
        assert all(deff.deff > 0), "DEFF values should be positive"
        assert all(np.isfinite(deff.effective_n)), "Effective n should be finite"
        assert all(deff.effective_n > 0), "Effective n should be positive"

        # DEFF > 1 expected for a complex survey with unequal weights
        assert deff.deff[1] > 1.0, (
            f"DEFF(KOWNRENT) = {deff.deff[1]:.4f}, expected > 1.0 for survey data"
        )
