"""Acceptance test: interlace vs R lme4 for two crossed random intercepts.

Model:  y ~ x + (1 | firm) + (1 | dept)
Data:   2000 obs, 50 firms × 12 depts  (tests/fixtures/two_re_data.csv)
Oracle: tests/fixtures/two_re_r_results.json (generated by gen_two_re.R)

Acceptance criteria (from CLAUDE.md):
  - Fixed effects abs_diff < 1e-4
  - Variance components rel_diff < 5%
  - BLUP correlation > 0.99
  - Conditional residual correlation > 0.999
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

import interlace

FIXTURES = Path(__file__).parent / "fixtures"


@pytest.fixture(scope="module")
def two_re_data() -> pd.DataFrame:
    return pd.read_csv(FIXTURES / "two_re_data.csv")


@pytest.fixture(scope="module")
def r_results() -> dict:
    return json.loads((FIXTURES / "two_re_r_results.json").read_text())


@pytest.fixture(scope="module")
def il_result(two_re_data):
    return interlace.fit("y ~ x", data=two_re_data, groups=["firm", "dept"])


def test_fixed_effects_match(il_result, r_results):
    r_fe = r_results["fe_params"]
    name_map = {"(Intercept)": "Intercept", "x": "x"}
    for r_name, il_name in name_map.items():
        diff = abs(il_result.fe_params[il_name] - r_fe[r_name])
        assert diff < 1e-4, (
            f"Fixed effect '{il_name}' abs_diff={diff:.2e} "
            f"(interlace={il_result.fe_params[il_name]:.6f}, R={r_fe[r_name]:.6f})"
        )


def test_variance_components_match(il_result, r_results):
    checks = [
        ("firm", r_results["var_firm"]),
        ("dept", r_results["var_dept"]),
    ]
    for col, r_var in checks:
        il_var = il_result.variance_components[col]
        rel_diff = abs(il_var - r_var) / r_var
        assert rel_diff < 0.05, (
            f"Variance component '{col}' rel_diff={rel_diff:.2%} "
            f"(interlace={il_var:.6f}, R={r_var:.6f})"
        )

    il_resid_var = il_result.scale
    r_resid_var = r_results["var_resid"]
    rel_diff = abs(il_resid_var - r_resid_var) / r_resid_var
    assert rel_diff < 0.05, (
        f"Residual variance rel_diff={rel_diff:.2%} "
        f"(interlace={il_resid_var:.6f}, R={r_resid_var:.6f})"
    )


def test_blups_firm_correlated(il_result, r_results):
    r_blups = r_results["blups_firm"]
    labels = sorted(r_blups.keys())
    r_arr = np.array([r_blups[g] for g in labels])
    il_arr = np.array([il_result.random_effects["firm"][g] for g in labels])
    corr = np.corrcoef(r_arr, il_arr)[0, 1]
    assert corr > 0.99, f"Firm BLUP correlation={corr:.4f} < 0.99"


def test_blups_dept_correlated(il_result, r_results):
    r_blups = r_results["blups_dept"]
    labels = sorted(r_blups.keys())
    r_arr = np.array([r_blups[g] for g in labels])
    il_arr = np.array([il_result.random_effects["dept"][g] for g in labels])
    corr = np.corrcoef(r_arr, il_arr)[0, 1]
    assert corr > 0.99, f"Dept BLUP correlation={corr:.4f} < 0.99"


def test_residuals_correlated(il_result, r_results):
    r_resid = np.array(r_results["resid_cond"])
    corr = np.corrcoef(il_result.resid, r_resid)[0, 1]
    assert corr > 0.999, f"Residual correlation={corr:.6f} < 0.999"
