"""Case 006: reproduce the R comparison using its shared synthetic CSV.

Download case-006-small-icc-data.csv beside this script, then run:
    python case-006-small-icc-multilevel.py
Or pass the CSV path as the first argument. Python standard library only.

For this balanced, intercept-only random-intercept model, interior REML
variance estimates have closed-form ANOVA expressions. This is not a
general mixed-model fitter for covariates, unequal sizes, or random slopes.
"""
import csv
import math
from pathlib import Path
import statistics
import sys

path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name(
    "case-006-small-icc-data.csv"
)
groups = {}
with path.open(newline="", encoding="utf-8-sig") as source:
    for row in csv.DictReader(source):
        groups.setdefault(row["cluster"], []).append(float(row["y"]))
J = len(groups)
sizes = {len(values) for values in groups.values()}
if J < 2 or len(sizes) != 1:
    raise ValueError("This example requires balanced clusters and J > 1.")
m = sizes.pop()
if m < 2:
    raise ValueError("This example requires at least two observations per cluster.")
values = [y for group in groups.values() for y in group]
n = len(values)
mean = statistics.mean(values)
means = [statistics.mean(group) for group in groups.values()]
ms_within = sum(
    sum((y - statistics.mean(group)) ** 2 for y in group)
    for group in groups.values()
) / (n - J)
ms_between = m * sum((value - mean) ** 2 for value in means) / (J - 1)
tau2 = (ms_between - ms_within) / m
if tau2 <= 0:
    raise ValueError("Interior REML formula not applicable: variance is on boundary.")
sigma2 = ms_within
icc = tau2 / (tau2 + sigma2)
naive_se = math.sqrt(statistics.variance(values) / n)
mixed_se = math.sqrt((tau2 + sigma2 / m) / J)
for name, value in {
    "clusters": J, "cluster_size": m, "observations": n,
    "mean_both_models": mean, "between_variance_REML": tau2,
    "within_variance_REML": sigma2, "fitted_ICC": icc,
    "independence_SE": naive_se, "random_intercept_SE": mixed_se,
    "fitted_design_effect": 1 + (m - 1) * icc,
    "SE_ratio": mixed_se / naive_se,
}.items():
    print(f"{name}: {value:.8f}")
assert n == 4000 and J == 80 and m == 50
assert math.isclose(icc, 0.02612601, abs_tol=1e-7)
assert math.isclose(naive_se, 0.01606634, abs_tol=1e-8)
assert math.isclose(mixed_se, 0.02426445, abs_tol=1e-8)
print("All checks passed. Python", sys.version.split()[0])