"""Cross-validation tests: diff-diff survey variance vs R's survey package.

Compares ATT, SE, df, and CI from diff-diff's Taylor Series Linearization
(TSL) variance against pre-computed golden values from R's authoritative
``survey::svyglm()`` (Lumley 2004).

Tier 1: DifferenceInDifferences + svyglm under 4 design variants + covariates + 3 seeds
Tier 2: CallawaySantAnna + did::att_gt with survey weights (ATT only)
Tier 3: BRR replicate weights via LinearRegression + svrepdesign

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

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

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

BENCHMARK_DIR = Path(__file__).parent.parent / "benchmarks" / "data" / "synthetic"
RESULTS_FILE = BENCHMARK_DIR / "survey_crossvalidation_r_results.json"

# Tolerances
ATT_RTOL = 1e-4
ATT_ATOL = 0.01
SE_RTOL = 0.01  # 1%
SE_ATOL = 0.05
CI_RTOL = 0.01
CI_ATOL = 0.05


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})"
    )


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


def _load_scenario_data(scenario):
    """Extract DataFrame from a scenario's embedded data dict."""
    data = scenario["data"]
    df = pd.DataFrame(
        {
            col: vals
            for col, vals in data.items()
        }
    )
    # Handle R's Inf serialization (jsonlite writes Inf as "Inf" string)
    if "first_treat" in df.columns:
        raw = df["first_treat"]
        ft = pd.to_numeric(raw, errors="coerce")
        # "Inf"/"Infinity" -> inf via to_numeric; any other non-numeric -> NaN
        non_numeric = ft.isna() & raw.notna()
        if non_numeric.any():
            bad = raw[non_numeric].unique().tolist()
            raise ValueError(
                f"Unexpected non-numeric first_treat values: {bad}. "
                "Expected integers or 'Inf' for never-treated."
            )
        ft = ft.replace(np.inf, 0)
        df["first_treat"] = ft.astype(int)

    # Ensure correct dtypes
    int_cols = ["unit", "period", "treated", "post", "stratum", "psu"]
    for col in int_cols:
        if col in df.columns:
            df[col] = df[col].astype(int)
    float_cols = ["outcome", "weight", "fpc", "x1", "x2"]
    for col in float_cols:
        if col in df.columns:
            df[col] = df[col].astype(float)
    return df


def _build_survey_design(design_type, has_fpc=True):
    """Construct the SurveyDesign for a given design type."""
    if design_type == "strata_psu_fpc":
        return SurveyDesign(
            weights="weight", strata="stratum", psu="psu",
            fpc="fpc" if has_fpc else None,
        )
    elif design_type == "strata_psu_nofpc":
        return SurveyDesign(weights="weight", strata="stratum", psu="psu")
    elif design_type == "weights_only":
        return SurveyDesign(weights="weight")
    elif design_type == "strata_only":
        return SurveyDesign(weights="weight", strata="stratum")
    else:
        raise ValueError(f"Unknown design type: {design_type}")


# ---------------------------------------------------------------------------
# Tier 1: svyglm DiD — DifferenceInDifferences vs R
# ---------------------------------------------------------------------------

TIER1_SCENARIOS = [
    # Design variants at seed=42
    "strata_psu_fpc_s42",
    "strata_psu_nofpc_s42",
    "weights_only_s42",
    "strata_only_s42",
    # Covariate variant
    "strata_psu_fpc_cov_s42",
    # Multi-seed robustness
    "strata_psu_fpc_s123",
    "strata_psu_fpc_s789",
]


class TestSvyglmDiD:
    """Cross-validate DifferenceInDifferences survey variance vs R svyglm."""

    def _fit_scenario(self, r_results, key):
        """Fit diff-diff DiD for a given scenario and return result."""
        r = r_results[key]
        data = _load_scenario_data(r)
        sd = _build_survey_design(r["design_type"])

        covariates = None
        if r.get("has_covariates"):
            covariates = ["x1", "x2"]

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

    @pytest.mark.parametrize("key", TIER1_SCENARIOS)
    def test_att_matches_r(self, r_results, key):
        result, r = self._fit_scenario(r_results, key)
        _assert_close(result.att, r["att"], ATT_RTOL, ATT_ATOL,
                       f"{key} ATT")

    @pytest.mark.parametrize("key", TIER1_SCENARIOS)
    def test_se_matches_r(self, r_results, key):
        result, r = self._fit_scenario(r_results, key)
        _assert_close(result.se, r["se"], SE_RTOL, SE_ATOL,
                       f"{key} SE")

    @pytest.mark.parametrize("key", TIER1_SCENARIOS)
    def test_df_matches_r(self, r_results, key):
        result, r = self._fit_scenario(r_results, key)
        py_df = result.survey_metadata.df_survey
        r_df = r["df"]
        assert py_df == r_df, (
            f"{key} df: Python={py_df}, R={r_df}"
        )

    @pytest.mark.parametrize("key", TIER1_SCENARIOS)
    def test_ci_matches_r(self, r_results, key):
        result, r = self._fit_scenario(r_results, key)
        _assert_close(result.conf_int[0], r["ci_lower"], CI_RTOL, CI_ATOL,
                       f"{key} CI lower")
        _assert_close(result.conf_int[1], r["ci_upper"], CI_RTOL, CI_ATOL,
                       f"{key} CI upper")

    @pytest.mark.parametrize("key", TIER1_SCENARIOS)
    def test_fpc_reduces_se(self, r_results, key):
        """FPC scenario should have smaller SE than non-FPC equivalent."""
        if key != "strata_psu_fpc_s42":
            pytest.skip("Only applies to strata_psu_fpc_s42")
        r_fpc = r_results["strata_psu_fpc_s42"]
        r_nofpc = r_results["strata_psu_nofpc_s42"]
        assert r_fpc["se"] < r_nofpc["se"], (
            f"FPC SE ({r_fpc['se']:.6f}) should be < non-FPC SE ({r_nofpc['se']:.6f})"
        )


# ---------------------------------------------------------------------------
# Tier 1b: TWFE plumbing test
# ---------------------------------------------------------------------------

class TestSvyglmTWFE:
    """Verify TWFE survey plumbing produces correct results.

    In balanced 2x2 (2 periods, no staggered adoption), TWFE ATT
    is algebraically equivalent to the pooled DiD ATT via FWL theorem.
    """

    def test_twfe_survey_produces_finite_results(self, r_results):
        r = r_results["strata_psu_fpc_s42"]
        data = _load_scenario_data(r)
        sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")

        est = TwoWayFixedEffects()
        result = est.fit(data, "outcome", "treated", "period", "unit",
                         survey_design=sd)

        assert np.isfinite(result.att), "TWFE ATT should be finite"
        assert np.isfinite(result.se), "TWFE SE should be finite"
        assert result.se > 0, "TWFE SE should be positive"

    def test_twfe_df_matches_expected(self, r_results):
        r = r_results["strata_psu_fpc_s42"]
        data = _load_scenario_data(r)
        sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")

        est = TwoWayFixedEffects()
        result = est.fit(data, "outcome", "treated", "period", "unit",
                         survey_design=sd)

        assert result.survey_metadata.df_survey == r["df"], (
            f"TWFE df: Python={result.survey_metadata.df_survey}, expected={r['df']}"
        )

    def test_twfe_att_close_to_pooled(self, r_results):
        """In balanced 2x2, TWFE ATT ≈ pooled DiD ATT (FWL equivalence)."""
        r = r_results["strata_psu_fpc_s42"]
        data = _load_scenario_data(r)
        sd = SurveyDesign(weights="weight", strata="stratum", psu="psu", fpc="fpc")

        pooled = DifferenceInDifferences().fit(
            data, "outcome", "treated", "post", survey_design=sd)
        twfe = TwoWayFixedEffects().fit(
            data, "outcome", "treated", "period", "unit", survey_design=sd)

        _assert_close(twfe.att, pooled.att, rtol=1e-3, atol=0.01,
                       label="TWFE vs pooled ATT")


# ---------------------------------------------------------------------------
# Tier 2: CallawaySantAnna with survey weights vs R did::att_gt
# ---------------------------------------------------------------------------

CS_SCENARIOS = ["cs_weighted_nocov", "cs_weighted_cov"]


class TestCSWeighted:
    """Cross-validate CallawaySantAnna weighted ATT vs R did::att_gt.

    SE comparison note: R's did::att_gt does not implement design-based
    variance — its SEs are survey-naive (just using weighted point estimates).
    diff-diff computes design-based SEs via compute_survey_if_variance().
    Therefore we compare ATT only and verify that diff-diff's design-based
    SE is finite, positive, and differs from the naive (unweighted) SE.
    For cs_weighted_cov, the reg+covariates SE additionally differs due to
    the documented deviation in REGISTRY.md (conservative plug-in IF vs
    semiparametrically efficient IF).
    """

    def _fit_scenario(self, r_results, key):
        r = r_results[key]
        data = _load_scenario_data(r)

        covariates = None
        if r.get("has_covariates"):
            covariates = ["x1"]

        sd = SurveyDesign(weights="weight")

        est = CallawaySantAnna(
            estimation_method="reg",
            control_group="never_treated",
            base_period="varying",
            n_bootstrap=0,
        )
        result = est.fit(
            data, "outcome", "unit", "period", "first_treat",
            covariates=covariates,
            aggregate="simple",
            survey_design=sd,
        )
        return result, r

    @pytest.mark.parametrize("key", CS_SCENARIOS)
    def test_att_matches_r_did(self, r_results, key):
        """Group-time ATTs should match R's att_gt with weightsname."""
        if key not in r_results:
            pytest.skip(f"Scenario {key} not in R results (did package not available)")
        result, r = self._fit_scenario(r_results, key)

        py_gt = result.group_time_effects

        # Build R (group, time) -> ATT mapping, sorted for stable comparison
        r_gt = {
            (int(g), int(t)): att
            for g, t, att in zip(r["gt_groups"], r["gt_periods"], r["gt_att"])
        }

        assert len(py_gt) == len(r_gt), (
            f"{key}: Python has {len(py_gt)} GT cells, R has {len(r_gt)}"
        )

        for (g, t), r_att in sorted(r_gt.items()):
            assert (g, t) in py_gt, f"{key}: missing Python GT cell ({g},{t})"
            _assert_close(
                py_gt[(g, t)]["effect"], r_att,
                rtol=1e-3, atol=0.05,
                label=f"{key} ATT(g={g},t={t})",
            )

    @pytest.mark.parametrize("key", CS_SCENARIOS)
    def test_overall_att_matches(self, r_results, key):
        """Simple-aggregated overall ATT should match R."""
        if key not in r_results:
            pytest.skip(f"Scenario {key} not in R results")
        result, r = self._fit_scenario(r_results, key)

        _assert_close(
            result.overall_att, r["overall_att"],
            rtol=1e-3, atol=0.05,
            label=f"{key} overall ATT",
        )

    @pytest.mark.parametrize("key", CS_SCENARIOS)
    def test_se_is_design_based(self, r_results, key):
        """Python SE should be positive and finite (design-based)."""
        if key not in r_results:
            pytest.skip(f"Scenario {key} not in R results")
        result, _ = self._fit_scenario(r_results, key)

        assert np.isfinite(result.overall_se), "Overall SE should be finite"
        assert result.overall_se > 0, "Overall SE should be positive"

    @pytest.mark.parametrize("key", CS_SCENARIOS)
    def test_design_se_differs_from_naive(self, r_results, key):
        """Design-based SE should differ from naive (unweighted) SE."""
        if key not in r_results:
            pytest.skip(f"Scenario {key} not in R results")
        r = r_results[key]
        data = _load_scenario_data(r)

        covariates = ["x1"] if r.get("has_covariates") else None

        # With survey weights
        sd = SurveyDesign(weights="weight")
        est_survey = CallawaySantAnna(
            estimation_method="reg", control_group="never_treated",
            base_period="varying", n_bootstrap=0,
        )
        result_survey = est_survey.fit(
            data, "outcome", "unit", "period", "first_treat",
            covariates=covariates, aggregate="simple", survey_design=sd,
        )

        # Without survey weights
        est_naive = CallawaySantAnna(
            estimation_method="reg", control_group="never_treated",
            base_period="varying", n_bootstrap=0,
        )
        result_naive = est_naive.fit(
            data, "outcome", "unit", "period", "first_treat",
            covariates=covariates, aggregate="simple",
        )

        assert result_survey.overall_se != result_naive.overall_se, (
            f"{key}: Design-based SE ({result_survey.overall_se:.6f}) should "
            f"differ from naive SE ({result_naive.overall_se:.6f})"
        )


# ---------------------------------------------------------------------------
# Tier 3: BRR replicate weights
# ---------------------------------------------------------------------------

class TestReplicateWeights:
    """Cross-validate BRR replicate weight variance vs R svrepdesign."""

    def _fit_brr(self, r_results):
        r = r_results["brr_replicate"]
        data = _load_scenario_data(r)

        # 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="weight",
            replicate_weights=rep_cols,
            replicate_method="BRR",
        )
        resolved = sd.resolve(data)

        # Build design matrix manually (no intercept — LR adds it)
        from diff_diff.linalg import LinearRegression

        X = np.column_stack([
            data["treated"].values.astype(float),
            data["post"].values.astype(float),
            (data["treated"] * data["post"]).values.astype(float),
        ])
        y = data["outcome"].values.astype(float)

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

        # Interaction coefficient is at index 3 (intercept=0, treated=1, post=2, treated:post=3)
        return reg, r

    def test_att_matches_svrepglm(self, r_results):
        if "brr_replicate" not in r_results:
            pytest.skip("BRR scenario not in R results")
        reg, r = self._fit_brr(r_results)

        py_att = reg.coefficients_[3]
        _assert_close(py_att, r["att"], ATT_RTOL, ATT_ATOL,
                       "BRR ATT")

    def test_se_matches_svrepglm(self, r_results):
        if "brr_replicate" not in r_results:
            pytest.skip("BRR scenario not in R results")
        reg, r = self._fit_brr(r_results)

        py_se = np.sqrt(reg.vcov_[3, 3])
        # Wider tolerance for replicate variance (approximate method)
        _assert_close(py_se, r["se"], rtol=0.05, atol=0.1,
                       label="BRR SE")

    def test_df_matches_svrepglm(self, r_results):
        if "brr_replicate" not in r_results:
            pytest.skip("BRR scenario not in R results")
        reg, r = self._fit_brr(r_results)

        assert reg.survey_df_ == r["df"], (
            f"BRR df: Python={reg.survey_df_}, R={r['df']}"
        )
