"""R-parity tests for SunAbraham vcov_type threading (Phase 1b PR 1/8).

Pins SA(vcov_type=...) cohort SEs and event-study SEs against R goldens
generated by benchmarks/R/generate_clubsandwich_golden.R. The golden file
contains the panel data so Python and R fit the SAME numerical panel.

Parity targets:
- classical → R lm() summary at atol=1e-10 (full-dummy)
- hc2 → R sandwich::vcovHC(type="HC2") at atol=1e-10 (full-dummy)
- hc2_bm → R clubSandwich::vcovCR(cluster=..., type="CR2") + coef_test$df_Satt at atol=1e-10 (full-dummy)
- hc1 event-study e=0 → R fixest::sunab + cluster=~unit at atol=5e-3
  (within-transform; documented HC1 finite-sample-correction deviation,
  see REGISTRY.md SunAbraham section and TODO.md row tracking the gap)
"""

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from diff_diff import SunAbraham

GOLDEN_PATH = (
    Path(__file__).resolve().parents[1] / "benchmarks" / "data" / "clubsandwich_cr2_golden.json"
)


@pytest.fixture(scope="module")
def sa_golden():
    """Load the sun_abraham_two_cohort scenario from the clubSandwich golden."""
    if not GOLDEN_PATH.exists():
        pytest.skip(f"Golden file missing: {GOLDEN_PATH}")
    with GOLDEN_PATH.open() as f:
        all_goldens = json.load(f)
    if "sun_abraham_two_cohort" not in all_goldens:
        pytest.skip("sun_abraham_two_cohort scenario missing from golden")
    return all_goldens["sun_abraham_two_cohort"]


@pytest.fixture(scope="module")
def sa_panel(sa_golden):
    """Reconstruct the panel DataFrame from the golden arrays."""
    return pd.DataFrame(
        {
            "unit": sa_golden["unit"],
            "time": sa_golden["time"],
            "first_treat": sa_golden["first_treat"],
            "y": sa_golden["y"],
        }
    )


def _target_cohort_se(sa_results, g, e):
    """Extract SE for the (cohort=g, event_time=e) interaction from SA results."""
    return float(sa_results.cohort_effects[(g, e)]["se"])


class TestSunAbrahamRParity:
    """SA vcov_type R-parity tests against clubSandwich/sandwich/fixest goldens."""

    def test_sa_classical_cohort_se_matches_lm_summary(self, sa_panel, sa_golden):
        """SA(vcov_type='classical') cohort SE matches R lm summary() classical SE."""
        sa = SunAbraham(vcov_type="classical")
        res = sa.fit(sa_panel, outcome="y", unit="unit", time="time", first_treat="first_treat")
        sa_se = _target_cohort_se(
            res, sa_golden["target_cohort_g"], sa_golden["target_event_time_e"]
        )
        r_se = sa_golden["classical_se"]
        assert np.isclose(sa_se, r_se, atol=1e-10), (
            f"SA classical SE {sa_se} does not match R lm() classical SE {r_se} "
            f"at atol=1e-10 (diff={abs(sa_se - r_se):.2e})"
        )

    def test_sa_hc2_cohort_se_matches_lm_vcovHC(self, sa_panel, sa_golden):
        """SA(vcov_type='hc2') cohort SE matches R sandwich::vcovHC(type='HC2')."""
        sa = SunAbraham(vcov_type="hc2")
        res = sa.fit(sa_panel, outcome="y", unit="unit", time="time", first_treat="first_treat")
        sa_se = _target_cohort_se(
            res, sa_golden["target_cohort_g"], sa_golden["target_event_time_e"]
        )
        r_se = sa_golden["hc2_se"]
        assert np.isclose(sa_se, r_se, atol=1e-10), (
            f"SA hc2 SE {sa_se} does not match R vcovHC(type='HC2') SE {r_se} "
            f"at atol=1e-10 (diff={abs(sa_se - r_se):.2e})"
        )
        # Sanity: hc2 == cr2_bm_singleton (one-way HC2-BM via singleton-cluster trick)
        assert np.isclose(sa_se, sa_golden["cr2_bm_singleton_se"], atol=1e-10), (
            "SA hc2 SE should match clubSandwich CR2 with singleton clusters "
            "(both reduce to one-way HC2-BM)"
        )

    def test_sa_hc2_bm_cohort_se_matches_clubsandwich_cr2(self, sa_panel, sa_golden):
        """SA(vcov_type='hc2_bm') cohort SE matches R clubSandwich::vcovCR(cluster=unit, type='CR2').

        Auto-cluster fires (no explicit cluster=); routes to CR2-BM at unit.
        """
        sa = SunAbraham(vcov_type="hc2_bm")
        res = sa.fit(sa_panel, outcome="y", unit="unit", time="time", first_treat="first_treat")
        sa_se = _target_cohort_se(
            res, sa_golden["target_cohort_g"], sa_golden["target_event_time_e"]
        )
        r_se = sa_golden["cr2_bm_unit_se"]
        assert np.isclose(sa_se, r_se, atol=1e-10), (
            f"SA hc2_bm SE {sa_se} does not match R clubSandwich CR2 (cluster=unit) "
            f"SE {r_se} at atol=1e-10 (diff={abs(sa_se - r_se):.2e})"
        )

    def test_bm_dof_matches_clubsandwich_singleton_and_unit(self, sa_panel, sa_golden):
        """Bell-McCaffrey Satterthwaite DOF parity vs R clubSandwich.

        SA's cohort_effects dict doesn't surface per-coefficient BM DOF, so
        we exercise the underlying LinearRegression machinery directly on
        the same full-dummy design SA's Part G builds. Pins both the
        singleton-cluster (one-way HC2-BM) DOF and the cluster=unit DOF at
        atol=1e-10. A regression in BM DOF logic would change inference
        output (p-values, CIs) without changing SEs; this test guards that.
        """
        import pandas as pd
        from diff_diff.linalg import LinearRegression

        # Reconstruct the full-dummy design SA's Part G builds (same shape
        # as benchmarks/R/generate_clubsandwich_golden.R `sun_abraham_two_cohort`).
        df = sa_panel.copy()
        df["rel_time"] = np.where(df["first_treat"] > 0, df["time"] - df["first_treat"], -999)
        treatment_groups = sorted(g for g in df["first_treat"].unique() if g > 0)
        all_rel = sorted(e for e in df.loc[df["first_treat"] > 0, "rel_time"].unique() if e != -1)
        interaction_cols = []
        for g in treatment_groups:
            for e in all_rel:
                name = f"D_{g}_{e}"
                ind = ((df["first_treat"] == g) & (df["rel_time"] == e)).astype(float)
                if ind.sum() > 0:
                    df[name] = ind
                    interaction_cols.append(name)

        target_col = f"D_{sa_golden['target_cohort_g']}_{sa_golden['target_event_time_e']}"
        target_within_idx = interaction_cols.index(target_col)
        # Full-dummy: intercept at 0, interactions at 1..n_interactions
        target_full_idx = target_within_idx + 1

        interaction_arrs = [df[c].values.astype(np.float64) for c in interaction_cols]
        unit_dummies = pd.get_dummies(df["unit"], prefix="_fe_unit", drop_first=True).values.astype(
            np.float64
        )
        time_dummies = pd.get_dummies(df["time"], prefix="_fe_time", drop_first=True).values.astype(
            np.float64
        )
        intercept = np.ones(len(df))
        X = np.column_stack([intercept] + interaction_arrs + [unit_dummies, time_dummies])
        y = df["y"].values.astype(np.float64)

        # Singleton-cluster CR2 (one-way HC2-BM): cluster_ids = seq_len(n)
        reg_singleton = LinearRegression(
            include_intercept=False,
            cluster_ids=np.arange(len(df)),
            vcov_type="hc2_bm",
        ).fit(X, y)
        bm_dof_singleton = float(reg_singleton._bm_dof[target_full_idx])
        r_dof_singleton = sa_golden["dof_bm_singleton"]
        assert np.isclose(bm_dof_singleton, r_dof_singleton, atol=1e-10), (
            f"Singleton-cluster BM DOF {bm_dof_singleton} does not match R "
            f"{r_dof_singleton} at atol=1e-10"
        )

        # CR2-BM at unit (the SA auto-cluster default for hc2_bm)
        reg_unit = LinearRegression(
            include_intercept=False,
            cluster_ids=df["unit"].values,
            vcov_type="hc2_bm",
        ).fit(X, y)
        bm_dof_unit = float(reg_unit._bm_dof[target_full_idx])
        r_dof_unit = sa_golden["dof_bm_unit"]
        assert np.isclose(bm_dof_unit, r_dof_unit, atol=1e-10), (
            f"CR2-BM (cluster=unit) DOF {bm_dof_unit} does not match R "
            f"{r_dof_unit} at atol=1e-10"
        )

    def test_sa_hc2_bm_aggregated_inference_uses_bm_contrast_dof(self, sa_panel, sa_golden):
        """SA(vcov_type='hc2_bm') user-facing aggregated inference uses
        Bell-McCaffrey contrast DOF, NOT normal-distribution DOF.

        Both the IW-aggregated event-study effects and the overall ATT
        are linear contrasts of cohort × event-time coefficients. The
        registry contract for ``vcov_type="hc2_bm"`` requires CR2 BM
        Satterthwaite DOF on these aggregated outputs (matching
        ``clubSandwich::Wald_test(constraints=matrix(c, 1), test="HTZ")``).
        This test guards against a regression to ``df=None`` →
        normal-distribution inference, which would silently change
        ``p_value``/``conf_int`` even though SE matches R at atol=1e-10.

        Reverse-engineers the DOF from ``p_value`` and ``|t|`` to pin
        the threading end-to-end.
        """
        from scipy import stats

        sa = SunAbraham(vcov_type="hc2_bm")
        res = sa.fit(sa_panel, outcome="y", unit="unit", time="time", first_treat="first_treat")

        def _dof_from_p(t_stat, p_value, lo=2.0, hi=1e6):
            target = 1 - p_value / 2
            for _ in range(80):
                mid = (lo + hi) / 2
                if stats.t.cdf(abs(t_stat), mid) < target:
                    lo = mid
                else:
                    hi = mid
            return 0.5 * (lo + hi)

        # Event-study e=0 contrast DOF: 4 cohorts contribute (g=4,5,6,7),
        # IW weights 0.25 each. R clubSandwich Wald_test(test='HTZ') = 35.0.
        eff_0 = res.event_study_effects[0]
        t_es0 = eff_0["effect"] / eff_0["se"]
        df_es0 = _dof_from_p(t_es0, eff_0["p_value"])
        r_df_es0 = sa_golden["dof_bm_contrast_es0_unit"]
        assert np.isclose(df_es0, r_df_es0, atol=1e-6), (
            f"SA event-study e=0 inferred BM DOF {df_es0:.6f} does not match "
            f"R clubSandwich Wald_test HTZ DOF {r_df_es0:.6f}. If this regresses, "
            "the user-facing event_study_effects[0]['p_value']/['conf_int'] is "
            "using normal-distribution DOF instead of CR2-BM contrast DOF."
        )

        # Overall ATT contrast DOF.
        t_all = res.overall_att / res.overall_se
        df_all = _dof_from_p(t_all, res.overall_p_value)
        r_df_all = sa_golden["dof_bm_contrast_overall_unit"]
        assert np.isclose(df_all, r_df_all, atol=1e-6), (
            f"SA overall ATT inferred BM DOF {df_all:.6f} does not match "
            f"R clubSandwich Wald_test HTZ DOF {r_df_all:.6f}. The overall_att "
            "inference is using normal-distribution DOF instead of CR2-BM "
            "contrast DOF."
        )

    def test_sa_hc1_event_study_e0_se_close_to_fixest_sunab(self, sa_panel, sa_golden):
        """SA(vcov_type='hc1') event-study e=0 SE matches fixest::sunab at atol=5e-3.

        SA hc1 uses within-transform path; fixest::sunab also uses within-transform.
        **Deviation from R:** fixest applies a finite-sample correction that
        counts absorbed unit + time FE in the effective parameter count
        (n / (n - k_total)) whereas SA's solve_ols counts only the
        within-transformed design columns (n / (n - k_dm)) — the SE values
        differ by ~1-2% (typically ~2e-3 in absolute terms on this panel).
        The IW aggregation step is otherwise identical. Tolerance set to
        atol=5e-3 (matches REGISTRY.md deviation note and the TODO.md
        follow-up row that tracks harmonizing the finite-sample correction).

        Tests the IW-aggregated event-study coefficient (matches fixest's
        natural sunab() output), not the per-cohort coefficient.
        """
        sa = SunAbraham(vcov_type="hc1")
        res = sa.fit(sa_panel, outcome="y", unit="unit", time="time", first_treat="first_treat")
        sa_se = float(res.event_study_effects[0]["se"])
        r_se = sa_golden["sunab_hc1_event_study_e0_se"]
        if r_se is None or (isinstance(r_se, float) and np.isnan(r_se)):
            pytest.skip("fixest::sunab parity SE missing from golden")
        # Absolute tolerance pin matching the documented deviation; harmonized
        # with REGISTRY.md SunAbraham deviation note and TODO.md row.
        assert abs(sa_se - r_se) < 5e-3, (
            f"SA hc1 event-study e=0 SE {sa_se:.6f} diverges from fixest::sunab "
            f"event-study SE {r_se:.6f} by {abs(sa_se - r_se):.2e} > 5e-3. "
            "Expected ~2e-3 from finite-sample correction (see REGISTRY.md "
            "deviation note); investigate if larger."
        )
