"""Cross-language weighted 2SLS parity tests against `estimatr::iv_robust`.

HAD Phase 4.5 B mass-point path: verifies that ``_fit_mass_point_2sls``
weighted HC1 and CR1 (Stata) sandwich SE match R's `estimatr::iv_robust`
(Blair et al. 2019) at ``atol=1e-10``. Classical SE diverges by O(1/n)
due to a different DOF / projection convention — documented in
`docs/methodology/REGISTRY.md` and skipped in parity assertions.

Fixture generated by ``benchmarks/R/generate_estimatr_iv_robust_golden.R``.
Guard per ``feedback_golden_file_pytest_skip``: CI isolated-install jobs
copy ``tests/`` only, not ``benchmarks/data/``, so a missing fixture
downgrades to pytest.skip rather than fail.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from diff_diff.had import _fit_mass_point_2sls

FIXTURE_PATH = (
    Path(__file__).parent.parent / "benchmarks" / "data" / "estimatr_iv_robust_golden.json"
)


def _load_fixture():
    if not FIXTURE_PATH.exists():
        pytest.skip(
            f"Golden fixture {FIXTURE_PATH} missing — regenerate via "
            f"`Rscript benchmarks/R/generate_estimatr_iv_robust_golden.R`."
        )
    with open(FIXTURE_PATH) as f:
        return json.load(f)


@pytest.fixture(scope="module")
def fixture():
    return _load_fixture()


class TestEstimatrIVRobustHC1Parity:
    """HC1 pweight sandwich: bit-exact match vs estimatr at atol=1e-10."""

    @pytest.mark.parametrize(
        "dgp_name",
        ["uniform_n200", "mild_n500", "informative_n500", "heavy_n1000"],
    )
    def test_hc1_beta_matches_estimatr(self, fixture, dgp_name):
        dgp = fixture["fixtures"][dgp_name]
        d = np.asarray(dgp["d"], dtype=np.float64)
        dy = np.asarray(dgp["dy"], dtype=np.float64)
        w = np.asarray(dgp["w"], dtype=np.float64)
        d_lower = float(dgp["d_lower"])

        beta, se, _psi = _fit_mass_point_2sls(
            d,
            dy,
            d_lower,
            None,
            "hc1",
            weights=w,
            return_influence=False,
        )

        np.testing.assert_allclose(
            beta,
            dgp["hc1"]["beta"],
            atol=1e-10,
            rtol=1e-10,
            err_msg=f"HC1 beta mismatch on DGP {dgp_name}",
        )

    @pytest.mark.parametrize(
        "dgp_name",
        ["uniform_n200", "mild_n500", "informative_n500", "heavy_n1000"],
    )
    def test_hc1_se_matches_estimatr(self, fixture, dgp_name):
        dgp = fixture["fixtures"][dgp_name]
        d = np.asarray(dgp["d"], dtype=np.float64)
        dy = np.asarray(dgp["dy"], dtype=np.float64)
        w = np.asarray(dgp["w"], dtype=np.float64)
        d_lower = float(dgp["d_lower"])

        _beta, se, _psi = _fit_mass_point_2sls(
            d,
            dy,
            d_lower,
            None,
            "hc1",
            weights=w,
        )

        np.testing.assert_allclose(
            se,
            dgp["hc1"]["se"],
            atol=1e-10,
            rtol=1e-10,
            err_msg=f"HC1 SE mismatch on DGP {dgp_name}",
        )


class TestEstimatrIVRobustCR1Parity:
    """CR1 pweight-cluster: bit-exact match vs estimatr (se_type='stata')."""

    def test_cr1_beta_matches_estimatr(self, fixture):
        dgp = fixture["fixtures"]["informative_cluster_n600"]
        d = np.asarray(dgp["d"], dtype=np.float64)
        dy = np.asarray(dgp["dy"], dtype=np.float64)
        w = np.asarray(dgp["w"], dtype=np.float64)
        cluster = np.asarray(dgp["cluster"], dtype=np.int64)
        d_lower = float(dgp["d_lower"])

        beta, _se, _psi = _fit_mass_point_2sls(
            d,
            dy,
            d_lower,
            cluster,
            "hc1",
            weights=w,
        )
        np.testing.assert_allclose(
            beta,
            dgp["cr1"]["beta"],
            atol=1e-10,
            rtol=1e-10,
        )

    def test_cr1_se_matches_estimatr(self, fixture):
        dgp = fixture["fixtures"]["informative_cluster_n600"]
        d = np.asarray(dgp["d"], dtype=np.float64)
        dy = np.asarray(dgp["dy"], dtype=np.float64)
        w = np.asarray(dgp["w"], dtype=np.float64)
        cluster = np.asarray(dgp["cluster"], dtype=np.int64)
        d_lower = float(dgp["d_lower"])

        _beta, se, _psi = _fit_mass_point_2sls(
            d,
            dy,
            d_lower,
            cluster,
            "hc1",
            weights=w,
        )
        np.testing.assert_allclose(
            se,
            dgp["cr1"]["se"],
            atol=1e-10,
            rtol=1e-10,
        )


class TestIFScaleInvariant:
    """IF scale invariant: compute_survey_if_variance(psi, trivial) ≈ V_HC1."""

    @pytest.mark.parametrize(
        "dgp_name",
        ["uniform_n200", "mild_n500", "informative_n500", "heavy_n1000"],
    )
    def test_if_reduces_to_hc1_variance(self, fixture, dgp_name):
        from diff_diff.survey import ResolvedSurveyDesign, compute_survey_if_variance

        dgp = fixture["fixtures"][dgp_name]
        d = np.asarray(dgp["d"], dtype=np.float64)
        dy = np.asarray(dgp["dy"], dtype=np.float64)
        w = np.asarray(dgp["w"], dtype=np.float64)
        d_lower = float(dgp["d_lower"])

        _beta, se, psi = _fit_mass_point_2sls(
            d,
            dy,
            d_lower,
            None,
            "hc1",
            weights=w,
            return_influence=True,
        )

        n = d.shape[0]
        trivial = ResolvedSurveyDesign(
            weights=w,
            weight_type="pweight",
            strata=None,
            psu=None,
            fpc=None,
            n_strata=1,
            n_psu=n,
            lonely_psu="remove",
            combined_weights=True,
            mse=False,
        )
        V_survey = compute_survey_if_variance(psi, trivial)
        np.testing.assert_allclose(
            V_survey,
            se**2,
            atol=1e-10,
            rtol=1e-10,
            err_msg=(
                "compute_survey_if_variance(psi, trivial_resolved) should "
                "equal the analytical HC1 variance under trivial survey "
                "design (PR #359 IF scale convention)."
            ),
        )
