"""Parity tests for nested random effects (/ nesting syntax).

Model: strength ~ 1 + (1|batch/cask)  ≡  (1|batch) + (1|batch:cask)
Data:  Pastes dataset (60 obs, 10 batches × 3 casks each)
Oracle: tests/fixtures/nested_pastes_r_results.json  (generated by gen_nested.R)
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

import interlace
from interlace.formula import parse_random_effects

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


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


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


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


@pytest.fixture(scope="module")
def il_result_slash(pastes_data):
    """Fit using lme4 / shorthand syntax."""
    return interlace.fit(
        "strength ~ 1",
        data=pastes_data,
        random=["(1|batch/cask)"],
    )


@pytest.fixture(scope="module")
def il_result_explicit(pastes_data):
    """Fit using explicit two-term form with pre-derived batch_cask column."""
    return interlace.fit(
        "strength ~ 1",
        data=pastes_data,
        random=["(1|batch)", "(1|batch_cask)"],
    )


# ---------------------------------------------------------------------------
# Test 1: / syntax is accepted (raises ValueError before implementation)
# ---------------------------------------------------------------------------


def test_slash_syntax_accepted(pastes_data):
    """(1|batch/cask) must be parsed without error after implementation."""
    # Before implementation: parse_random_effects raises ValueError.
    # After implementation: fit() returns a CrossedLMEResult.
    result = interlace.fit(
        "strength ~ 1",
        data=pastes_data,
        random=["(1|batch/cask)"],
    )
    assert hasattr(result, "fe_params"), "Expected a CrossedLMEResult"


# ---------------------------------------------------------------------------
# Test 2: Fixed effects parity
# ---------------------------------------------------------------------------


def test_fixed_effects_match(il_result_slash, r_results):
    r_intercept = r_results["fixef"]["(Intercept)"]
    il_intercept = il_result_slash.fe_params["Intercept"]
    diff = abs(il_intercept - r_intercept)
    assert diff < 1e-4, (
        f"Intercept abs_diff={diff:.2e} "
        f"(interlace={il_intercept:.6f}, R={r_intercept:.6f})"
    )


# ---------------------------------------------------------------------------
# Test 3: Variance components parity
# ---------------------------------------------------------------------------


def test_variance_components_match(il_result_slash, r_results):
    vc = r_results["var_components"]
    checks = [
        ("batch", vc["batch"]),
        ("batch:cask", vc["batch:cask"]),
    ]
    for grp, r_var in checks:
        il_var = il_result_slash.variance_components[grp]
        rel_diff = abs(il_var - r_var) / r_var
        assert rel_diff < 0.05, (
            f"Variance component '{grp}' rel_diff={rel_diff:.2%} "
            f"(interlace={il_var:.6f}, R={r_var:.6f})"
        )

    il_resid_var = il_result_slash.scale
    r_resid_var = vc["residual"]
    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})"
    )


# ---------------------------------------------------------------------------
# Test 4: BLUP parity
# ---------------------------------------------------------------------------


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


def test_blups_batch_cask_correlated(il_result_slash, r_results):
    r_blups = r_results["ranef"]["batch:cask"]
    labels = sorted(r_blups.keys())
    r_arr = np.array([r_blups[g] for g in labels])
    il_arr = np.array([il_result_slash.random_effects["batch:cask"][g] for g in labels])
    corr = np.corrcoef(r_arr, il_arr)[0, 1]
    assert corr > 0.99, f"batch:cask BLUP correlation={corr:.4f} < 0.99"


# ---------------------------------------------------------------------------
# Test 5: Explicit two-term form matches / shorthand
# ---------------------------------------------------------------------------


def test_explicit_matches_slash(il_result_slash, il_result_explicit):
    """Explicit two-term form must match the / shorthand."""
    # Fixed effects
    slash_fe = il_result_slash.fe_params["Intercept"]
    expl_fe = il_result_explicit.fe_params["Intercept"]
    assert abs(slash_fe - expl_fe) < 1e-10, (
        f"Intercept mismatch: slash={slash_fe:.8f}, explicit={expl_fe:.8f}"
    )

    # Variance components for batch
    slash_vc_batch = il_result_slash.variance_components["batch"]
    expl_vc_batch = il_result_explicit.variance_components["batch"]
    assert abs(slash_vc_batch - expl_vc_batch) < 1e-10, (
        f"batch VC mismatch: slash={slash_vc_batch:.8f}, explicit={expl_vc_batch:.8f}"
    )


# ---------------------------------------------------------------------------
# Test 6: Depth-3 nesting parse (unit test on parser)
# ---------------------------------------------------------------------------


def test_depth3_nesting_parse():
    """parse_random_effects(['(1|a/b/c)']) must return 3 RandomEffectSpec objects."""
    specs = parse_random_effects(["(1|a/b/c)"])
    assert len(specs) == 3, f"Expected 3 specs, got {len(specs)}"

    groups = [s.group for s in specs]
    assert groups[0] == "a", f"Expected group 'a', got '{groups[0]}'"
    assert groups[1] == "a:b", f"Expected group 'a:b', got '{groups[1]}'"
    assert groups[2] == "a:b:c", f"Expected group 'a:b:c', got '{groups[2]}'"

    # Each spec has interaction_cols set correctly
    assert specs[0].interaction_cols == [], (
        f"Top-level spec should have empty interaction_cols, "
        f"got {specs[0].interaction_cols}"
    )
    assert specs[1].interaction_cols == ["a", "b"], (
        f"Second spec interaction_cols should be ['a', 'b'], "
        f"got {specs[1].interaction_cols}"
    )
    assert specs[2].interaction_cols == ["a", "b", "c"], (
        f"Third spec interaction_cols should be ['a', 'b', 'c'], "
        f"got {specs[2].interaction_cols}"
    )
