"""Cross-validation tests: StaggeredTripleDifference vs R triplediff package.

Compares group-time ATT(g,t) and SE values against pre-computed R golden
values generated by benchmarks/R/benchmark_staggered_triplediff.R.

CSV fixtures are generated on-the-fly via R if available, or tests skip.
"""

import json
import os
import subprocess
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from diff_diff import StaggeredTripleDifference

BENCHMARK_DIR = Path(__file__).parent.parent / "benchmarks" / "data" / "synthetic"
RESULTS_FILE = BENCHMARK_DIR / "staggered_ddd_r_results.json"
R_SCRIPT = Path(__file__).parent.parent / "benchmarks" / "R" / "benchmark_staggered_triplediff.R"

# Column mapping: R -> Python
R_TO_PY_COLS = {
    "y": "outcome",
    "time": "period",
    "id": "unit",
    "state": "first_treat",
    "partition": "eligibility",
}

# Tolerance: ATT within 0.1% relative, SE within 1% relative
ATT_RTOL = 0.001
SE_RTOL = 0.01
# Absolute tolerance for values near zero
ATT_ATOL = 0.01
SE_ATOL = 0.1


def _r_triplediff_available() -> bool:
    """Check if R + triplediff are available."""
    r_env = os.environ.get("DIFF_DIFF_R", "auto").lower()
    if r_env == "skip":
        return False
    try:
        result = subprocess.run(
            ["Rscript", "-e", "library(triplediff); library(jsonlite); cat('OK')"],
            capture_output=True, text=True, timeout=30,
        )
        return result.returncode == 0 and "OK" in result.stdout
    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
        return False


def _ensure_csv_fixtures():
    """Generate CSV fixtures via R if they don't exist and R is available."""
    # Check if any CSV is missing
    needed_keys = ["s42_dgp1", "s123_dgp1", "s99_dgp1"]
    missing = [k for k in needed_keys
               if not (BENCHMARK_DIR / f"staggered_ddd_data_{k}.csv").exists()]
    if not missing:
        return True  # All present

    if not _r_triplediff_available():
        return False  # Can't generate

    # Run R script to generate all CSVs + JSON
    try:
        result = subprocess.run(
            ["Rscript", str(R_SCRIPT)],
            capture_output=True, text=True, timeout=120,
            cwd=str(Path(__file__).parent.parent),
        )
        return result.returncode == 0
    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
        return False


@pytest.fixture(scope="module")
def r_results():
    """Load pre-computed R golden values, generating CSVs if needed."""
    if not RESULTS_FILE.exists():
        pytest.skip(
            f"R benchmark file not found: {RESULTS_FILE}. "
            "Run: Rscript benchmarks/R/benchmark_staggered_triplediff.R"
        )
    # Try to generate CSV fixtures if missing
    _ensure_csv_fixtures()
    with open(RESULTS_FILE) as f:
        return json.load(f)


def _load_r_data(seed: int, dgp: int) -> pd.DataFrame:
    """Load the CSV data that R used for a given scenario."""
    csv_path = BENCHMARK_DIR / f"staggered_ddd_data_s{seed}_dgp{dgp}.csv"
    if not csv_path.exists():
        pytest.skip(
            f"Data file not found: {csv_path}. "
            "Requires R + triplediff to generate."
        )
    data = pd.read_csv(csv_path)
    return data.rename(columns=R_TO_PY_COLS)


def _run_python(data, method, control_group):
    """Run Python estimator and return results."""
    est = StaggeredTripleDifference(
        estimation_method=method,
        control_group=control_group,
        base_period="varying",
    )
    return est.fit(
        data, "outcome", "unit", "period", "first_treat", "eligibility",
        aggregate="event_study",
    )


def _assert_close(py_val, r_val, rtol, atol, label):
    """Assert values are close with informative error message."""
    if np.isnan(r_val):
        return  # skip NaN comparisons
    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})"
    )


# ---------------------------------------------------------------------------
# Nevertreated scenarios (single comparison group, no GMM)
# ---------------------------------------------------------------------------

NT_SCENARIOS = [
    "s42_dgp1_dr_nt",
    "s42_dgp1_ipw_nt",
    "s42_dgp1_reg_nt",
    "s123_dgp1_dr_nt",
    "s99_dgp1_dr_nt",
]


class TestStaggeredDDDNevertreated:
    """Cross-validate against R with control_group='nevertreated'."""

    @pytest.mark.parametrize("key", NT_SCENARIOS)
    def test_gt_att_matches_r(self, r_results, key):
        r = r_results[key]
        data = _load_r_data(r["seed"], r["dgp_type"])
        res = _run_python(data, r["est_method"], r["control_group"])

        py_gt = sorted(res.group_time_effects.items())
        r_gt = list(zip(r["gt_groups"], r["gt_periods"]))
        assert len(py_gt) == len(r["gt_att"]), (
            f"{key}: Python has {len(py_gt)} GT cells, R has {len(r['gt_att'])}"
        )
        for i, ((g, t), eff) in enumerate(py_gt):
            assert (g, t) == (r_gt[i][0], r_gt[i][1]), (
                f"{key}: GT cell mismatch at index {i}: Python=({g},{t}), R={r_gt[i]}"
            )
            _assert_close(
                eff["effect"], r["gt_att"][i],
                ATT_RTOL, ATT_ATOL,
                f"{key} ATT(g={g},t={t})",
            )

    @pytest.mark.parametrize("key", NT_SCENARIOS)
    def test_gt_se_matches_r(self, r_results, key):
        r = r_results[key]
        data = _load_r_data(r["seed"], r["dgp_type"])
        res = _run_python(data, r["est_method"], r["control_group"])

        py_gt = sorted(res.group_time_effects.items())
        assert len(py_gt) == len(r["gt_se"]), (
            f"{key}: Python has {len(py_gt)} GT cells, R has {len(r['gt_se'])}"
        )
        for i, ((g, t), eff) in enumerate(py_gt):
            _assert_close(
                eff["se"], r["gt_se"][i],
                SE_RTOL, SE_ATOL,
                f"{key} SE(g={g},t={t})",
            )


# ---------------------------------------------------------------------------
# Notyettreated scenarios (multiple comparison groups, GMM)
# ---------------------------------------------------------------------------

NYT_SCENARIOS = [
    "s42_dgp1_dr_nyt",
    "s42_dgp1_ipw_nyt",
    "s42_dgp1_reg_nyt",
    "s123_dgp1_dr_nyt",
    "s99_dgp1_dr_nyt",
]


class TestStaggeredDDDNotyettreated:
    """Cross-validate against R with control_group='notyettreated'."""

    @pytest.mark.parametrize("key", NYT_SCENARIOS)
    def test_gt_att_matches_r(self, r_results, key):
        r = r_results[key]
        data = _load_r_data(r["seed"], r["dgp_type"])
        res = _run_python(data, r["est_method"], r["control_group"])

        py_gt = sorted(res.group_time_effects.items())
        r_gt = list(zip(r["gt_groups"], r["gt_periods"]))
        assert len(py_gt) == len(r["gt_att"]), (
            f"{key}: Python has {len(py_gt)} GT cells, R has {len(r['gt_att'])}"
        )
        for i, ((g, t), eff) in enumerate(py_gt):
            assert (g, t) == (r_gt[i][0], r_gt[i][1]), (
                f"{key}: GT cell mismatch at index {i}: Python=({g},{t}), R={r_gt[i]}"
            )
            _assert_close(
                eff["effect"], r["gt_att"][i],
                ATT_RTOL, ATT_ATOL,
                f"{key} ATT(g={g},t={t})",
            )

    @pytest.mark.parametrize("key", NYT_SCENARIOS)
    def test_gt_se_matches_r(self, r_results, key):
        r = r_results[key]
        data = _load_r_data(r["seed"], r["dgp_type"])
        res = _run_python(data, r["est_method"], r["control_group"])

        py_gt = sorted(res.group_time_effects.items())
        assert len(py_gt) == len(r["gt_se"]), (
            f"{key}: Python has {len(py_gt)} GT cells, R has {len(r['gt_se'])}"
        )
        for i, ((g, t), eff) in enumerate(py_gt):
            _assert_close(
                eff["se"], r["gt_se"][i],
                SE_RTOL, SE_ATOL,
                f"{key} SE(g={g},t={t})",
            )


# ---------------------------------------------------------------------------
# Event study and overall ATT aggregation
# ---------------------------------------------------------------------------

ALL_SCENARIOS = NT_SCENARIOS + NYT_SCENARIOS


class TestStaggeredDDDAggregation:
    """Cross-validate event study and overall ATT aggregation.

    Note: Aggregation weights differ between R's agg_ddd() and the
    CallawaySantAnna mixin we reuse. The group-time ATT(g,t) values
    match R exactly (0.00%), but aggregation introduces small differences
    due to the weighting scheme. We use 10% tolerance here.
    """

    AGG_RTOL = 0.25  # 25% — aggregation weighting differs from R's agg_ddd
    AGG_ATOL = 3.0

    @pytest.mark.parametrize("key", ALL_SCENARIOS)
    def test_event_study_att_close_to_r(self, r_results, key):
        r = r_results[key]
        if r.get("es_att") is None:
            pytest.skip("No event study results in R")
        data = _load_r_data(r["seed"], r["dgp_type"])
        res = _run_python(data, r["est_method"], r["control_group"])

        assert res.event_study_effects is not None
        r_es = dict(zip(r["es_event_times"], r["es_att"]))
        for e, eff in res.event_study_effects.items():
            if e in r_es:
                _assert_close(
                    eff["effect"], r_es[e],
                    self.AGG_RTOL, self.AGG_ATOL,
                    f"{key} ES ATT(e={e})",
                )

    @pytest.mark.parametrize("key", ALL_SCENARIOS)
    def test_overall_att_close_to_r(self, r_results, key):
        r = r_results[key]
        if np.isnan(r.get("overall_att_simple", float("nan"))):
            pytest.skip("No simple aggregation in R")
        data = _load_r_data(r["seed"], r["dgp_type"])
        est = StaggeredTripleDifference(
            estimation_method=r["est_method"],
            control_group=r["control_group"],
            base_period="varying",
        )
        res = est.fit(
            data, "outcome", "unit", "period", "first_treat", "eligibility",
        )
        _assert_close(
            res.overall_att, r["overall_att_simple"],
            self.AGG_RTOL, self.AGG_ATOL,
            f"{key} overall ATT (simple)",
        )
