"""Cross-validation tests: survey-enabled estimators vs R's survey::svyglm().

Validates 4 estimators against R golden values:
  S1: ImputationDiD   — control-only WLS regression (covariate coefficients)
  S2: StackedDiD      — stacked WLS with Q-weight x survey weight composition
  S3: SunAbraham      — interaction-weighted ATT
  S4: TripleDifference — three-way interaction DDD

JSON fixtures generated by: benchmarks/R/benchmark_survey_estimators.R
"""

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from diff_diff import SurveyDesign
from diff_diff.imputation import ImputationDiD
from diff_diff.linalg import LinearRegression
from diff_diff.stacked_did import StackedDiD
from diff_diff.sun_abraham import SunAbraham
from diff_diff.triple_diff import TripleDifference

GOLDEN_FILE = (
    Path(__file__).parent.parent
    / "benchmarks"
    / "data"
    / "synthetic"
    / "survey_estimator_validation_golden.json"
)

# Tolerances — observed gaps: S1 0.00%, S2 0.77%, S3 0.00%, S4 0.36%
COEF_RTOL = 1e-8
COEF_ATOL = 1e-8
SE_RTOL = 0.015  # 1.5% — covers S2 (0.77%) and S4 (0.36%) with headroom
SE_ATOL = 1e-6


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):
    """Load embedded dataset from golden JSON."""
    data = results[key]
    return pd.DataFrame({col: vals for col, vals in data.items()})


def _parse_cohort_key(key_str):
    """Parse 'g3_e0' -> (3, 0), 'g3_em2' -> (3, -2)."""
    g_part, e_part = key_str.split("_")
    g = int(g_part[1:])  # 'g3' -> 3
    e_str = e_part[1:]  # 'e0' -> '0', 'em2' -> 'm2'
    e = -int(e_str[1:]) if e_str.startswith("m") else int(e_str)
    return g, e


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


class TestSurveyEstimatorValidation:
    """Validate survey-enabled estimators against R's survey::svyglm()."""

    def _load_staggered_data(self, golden):
        """Load staggered panel from _staggered_data key."""
        df = _load_embedded_data(golden, "_staggered_data")
        # Handle Inf serialization: R's toJSON writes Inf as "Inf" (string)
        df["first_treat"] = pd.to_numeric(df["first_treat"], errors="coerce")
        for col in ["unit", "period"]:
            df[col] = df[col].astype(int)
        for col in ["outcome", "weight", "fpc", "x1", "x2"]:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df["stratum"] = df["stratum"].astype(int)
        df["psu"] = df["psu"].astype(int)
        return df

    def _load_ddd_data(self, golden):
        """Load DDD cross-section from _ddd_data key."""
        df = _load_embedded_data(golden, "_ddd_data")
        for col in ["group", "partition", "time"]:
            df[col] = df[col].astype(int)
        for col in ["outcome", "weight", "fpc"]:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df["stratum"] = df["stratum"].astype(int)
        df["psu"] = df["psu"].astype(int)
        return df

    # ------------------------------------------------------------------
    # S1: ImputationDiD — control-only WLS regression
    # ------------------------------------------------------------------

    def test_s1_imputation_control_regression(self, golden):
        """Validate control-only WLS covariate coefficients against R's svyglm."""
        r = golden["s1_imputation_did"]
        data = self._load_staggered_data(golden)

        # Build Omega_0 (untreated observations)
        # Matches imputation.py line 341: omega_0_mask = ~df["_treated"]
        omega_0 = data[
            (np.isinf(data["first_treat"]))
            | (data["period"] < data["first_treat"])
        ].copy()
        omega_0 = omega_0.reset_index(drop=True)

        sd = SurveyDesign(
            weights="weight", strata="stratum", psu="psu", fpc="fpc"
        )
        resolved = sd.resolve(omega_0)
        survey_weights = omega_0["weight"].values.astype(float)

        # Build design matrix: unit + time dummies + covariates
        # Algebraically equivalent to within-transformation for covariate coefs
        unit_dummies = pd.get_dummies(
            omega_0["unit"].astype(str), prefix="u", drop_first=True
        )
        time_dummies = pd.get_dummies(
            omega_0["period"].astype(str), prefix="t", drop_first=True
        )
        X = np.column_stack(
            [
                unit_dummies.values.astype(float),
                time_dummies.values.astype(float),
                omega_0["x1"].values,
                omega_0["x2"].values,
            ]
        )
        y = omega_0["outcome"].values

        reg = LinearRegression(
            include_intercept=True,
            survey_design=resolved,
            weights=survey_weights,
            weight_type="pweight",
        )
        reg.fit(X, y)

        # x1 and x2 are last two columns
        _assert_close(
            reg.coefficients_[-2],
            r["coef_x1"],
            COEF_RTOL,
            COEF_ATOL,
            "S1 x1 coef",
        )
        _assert_close(
            reg.coefficients_[-1],
            r["coef_x2"],
            COEF_RTOL,
            COEF_ATOL,
            "S1 x2 coef",
        )
        _assert_close(
            np.sqrt(reg.vcov_[-2, -2]),
            r["se_x1"],
            SE_RTOL,
            SE_ATOL,
            "S1 x1 SE",
        )
        _assert_close(
            np.sqrt(reg.vcov_[-1, -1]),
            r["se_x2"],
            SE_RTOL,
            SE_ATOL,
            "S1 x2 SE",
        )

    def test_s1_imputation_end_to_end(self, golden):
        """Smoke test: ImputationDiD.fit() with survey weights produces finite ATT."""
        data = self._load_staggered_data(golden)
        sd = SurveyDesign(
            weights="weight", strata="stratum", psu="psu", fpc="fpc"
        )
        est = ImputationDiD()
        result = est.fit(
            data,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            covariates=["x1", "x2"],
            survey_design=sd,
        )
        assert np.isfinite(result.overall_att), "ImputationDiD ATT should be finite"
        assert np.isfinite(result.overall_se), "ImputationDiD SE should be finite"
        assert result.overall_se > 0, "ImputationDiD SE should be positive"

    # ------------------------------------------------------------------
    # S2: StackedDiD — stacked WLS with Q-weight x survey weight
    # ------------------------------------------------------------------

    def test_s2_stacked_did(self, golden):
        """Validate StackedDiD ATT against R's svyglm on stacked data."""
        r = golden["s2_stacked_did"]
        data = self._load_staggered_data(golden)

        sd = SurveyDesign(
            weights="weight", strata="stratum", psu="psu", fpc="fpc"
        )
        est = StackedDiD(
            kappa_pre=1,
            kappa_post=1,
            weighting="sample_share",
            clean_control="never_treated",
        )
        result = est.fit(
            data,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            aggregate="event_study",
            survey_design=sd,
        )

        _assert_close(
            result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S2 ATT"
        )
        _assert_close(
            result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S2 SE"
        )

        # Compare individual event-study coefficients and SEs
        if "event_study" in r:
            for e_str, r_eff in r["event_study"].items():
                e = int(e_str)
                py_eff = result.event_study_effects.get(e)
                if py_eff is not None:
                    _assert_close(
                        py_eff["effect"],
                        r_eff["coef"],
                        COEF_RTOL,
                        COEF_ATOL,
                        f"S2 delta_{e}",
                    )
                    _assert_close(
                        py_eff["se"],
                        r_eff["se"],
                        SE_RTOL,
                        SE_ATOL,
                        f"S2 SE_{e}",
                    )

    # ------------------------------------------------------------------
    # S3: SunAbraham — interaction-weighted ATT
    # ------------------------------------------------------------------

    def test_s3_sun_abraham(self, golden):
        """Validate SunAbraham IW-aggregated ATT against R's svyglm."""
        r = golden["s3_sun_abraham"]
        data = self._load_staggered_data(golden)

        sd = SurveyDesign(
            weights="weight", strata="stratum", psu="psu", fpc="fpc"
        )
        est = SunAbraham()
        result = est.fit(
            data,
            outcome="outcome",
            unit="unit",
            time="period",
            first_treat="first_treat",
            survey_design=sd,
        )

        # Compare overall IW-aggregated ATT
        _assert_close(
            result.overall_att, r["att"], COEF_RTOL, COEF_ATOL, "S3 ATT"
        )
        _assert_close(
            result.overall_se, r["se"], SE_RTOL, SE_ATOL, "S3 SE"
        )

        # Compare individual cohort x rel-time effects and SEs (post-treatment only)
        if "cohort_effects" in r:
            for key_str, r_delta in r["cohort_effects"].items():
                g, e = _parse_cohort_key(key_str)
                py_entry = result.cohort_effects.get((g, e))
                if py_entry is not None:
                    _assert_close(
                        py_entry["effect"],
                        r_delta["coef"],
                        COEF_RTOL,
                        COEF_ATOL,
                        f"S3 delta_{g}_{e}",
                    )
                    _assert_close(
                        py_entry["se"],
                        r_delta["se"],
                        SE_RTOL,
                        SE_ATOL,
                        f"S3 SE_{g}_{e}",
                    )

    # ------------------------------------------------------------------
    # S4: TripleDifference — three-way interaction DDD
    # ------------------------------------------------------------------

    def test_s4_triple_diff(self, golden):
        """Validate TripleDifference DDD against R's svyglm three-way interaction."""
        r = golden["s4_triple_diff"]
        data = self._load_ddd_data(golden)

        sd = SurveyDesign(
            weights="weight", strata="stratum", psu="psu", fpc="fpc"
        )
        est = TripleDifference(estimation_method="reg")
        result = est.fit(
            data,
            outcome="outcome",
            group="group",
            partition="partition",
            time="time",
            survey_design=sd,
        )

        _assert_close(
            result.att, r["att"], COEF_RTOL, COEF_ATOL, "S4 DDD coef"
        )
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL, "S4 DDD SE")
        _assert_close(
            result.t_stat, r["t_stat"], SE_RTOL, SE_ATOL, "S4 t_stat"
        )
        _assert_close(
            result.conf_int[0], r["ci_lower"], SE_RTOL, SE_ATOL, "S4 CI lower"
        )
        _assert_close(
            result.conf_int[1], r["ci_upper"], SE_RTOL, SE_ATOL, "S4 CI upper"
        )
