"""Fit the hierarchical linear models reported in the manuscript.

Models (REML, statsmodels MixedLM):
- Null (unconditional) model per system and pooled -> variance components, ICC
- Country-specific full models: random intercept + random ESCS slope
- Pooled models with system fixed effects (random intercept):
    M1 student predictors; M2 + school predictors; M3 + cross-level interactions

Note: estimates are complete-case and unweighted at the model stage
(provisional; survey-weighted replication with WeMix/BRR planned).

Input : data/processed/analysis_dataset.csv
Output: output/results/res_icc.json, res_percountry.json, res_pooled.json,
        output/results/table2_desc.csv
"""
import json
import warnings
from pathlib import Path

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

warnings.filterwarnings("ignore")

ROOT = Path(__file__).resolve().parents[1]
RESULTS = ROOT / "output" / "results"
RESULTS.mkdir(parents=True, exist_ok=True)

FIX = ("ESCS_c + girl + immig + ICTQUAL_c + ICTHOME_c + ICTWKDY_c + "
       "ICTWKEND_c + ICTOUT_c + AGE_c + GRADE_c + sch_escs_g + RATCMP1_g + "
       "EDUSHORT_g + private + rural")


def dump(md, null_tau=None, null_sig=None):
    out = dict(
        params={k: float(v) for k, v in md.params.items() if not k.startswith("Group")},
        bse={k: float(v) for k, v in md.bse.items() if not k.startswith("Group")},
        pvalues={k: float(v) for k, v in md.pvalues.items() if not k.startswith("Group")},
        tau00=float(md.cov_re.iloc[0, 0]),
        sigma2=float(md.scale),
        tau11=float(md.cov_re.iloc[1, 1]) if md.cov_re.shape[0] > 1 else None,
    )
    if null_tau is not None:
        out["prv_l1"] = 1 - out["sigma2"] / null_sig
        out["prv_l2"] = 1 - out["tau00"] / null_tau
    return out


def weighted_descriptives(m: pd.DataFrame) -> pd.DataFrame:
    def wm(g, v):
        return np.average(g[v], weights=g.W_FSTUWT)

    def wsd(g, v):
        mu = wm(g, v)
        return np.sqrt(np.average((g[v] - mu) ** 2, weights=g.W_FSTUWT))

    rows = []
    for c, g in m.groupby("country"):
        rows.append(dict(
            country=c, N=len(g), J=g.CNTSCHID.nunique(),
            icteffic_m=wm(g, "ICTEFFIC"), icteffic_sd=wsd(g, "ICTEFFIC"),
            escs_m=wm(g, "ESCS"), escs_sd=wsd(g, "ESCS"),
            girl=wm(g, "girl") * 100, immig=wm(g, "immig") * 100,
            ictqual_m=wm(g, "ICTQUAL"), icthome_m=wm(g, "ICTHOME"),
            private=g.groupby("CNTSCHID").private.first().mean() * 100,
            rural=g.groupby("CNTSCHID").rural.first().mean() * 100))
    return pd.DataFrame(rows)


def main() -> None:
    m = pd.read_csv(ROOT / "data" / "processed" / "analysis_dataset.csv")

    weighted_descriptives(m).to_csv(RESULTS / "table2_desc.csv", index=False)

    # ---- null models and ICCs ----
    icc = {}
    for c, g in m.groupby("country"):
        md = smf.mixedlm("ICTEFFIC ~ 1", g, groups=g.CNTSCHID).fit(reml=True)
        t00, s2 = float(md.cov_re.iloc[0, 0]), float(md.scale)
        icc[c] = dict(tau00=t00, sigma2=s2, icc=t00 / (t00 + s2),
                      N=len(g), J=g.CNTSCHID.nunique(),
                      gamma00=float(md.params["Intercept"]),
                      se=float(md.bse["Intercept"]))
    md = smf.mixedlm("ICTEFFIC ~ C(country)", m, groups=m.CNTSCHID).fit(reml=True)
    t00, s2 = float(md.cov_re.iloc[0, 0]), float(md.scale)
    icc["Pooled (country FE)"] = dict(
        tau00=t00, sigma2=s2, icc=t00 / (t00 + s2),
        N=len(m), J=m.CNTSCHID.nunique(),
        gamma00=float(md.params["Intercept"]), se=float(md.bse["Intercept"]))
    json.dump(icc, open(RESULTS / "res_icc.json", "w"), indent=1)
    print(pd.DataFrame(icc).T.round(3))

    # ---- country-specific models (random intercept + random ESCS slope) ----
    percountry = {}
    for c, g in m.groupby("country"):
        try:
            md = smf.mixedlm(f"ICTEFFIC ~ {FIX}", g, groups=g.CNTSCHID,
                             re_formula="~ESCS_c").fit(reml=True)
            conv = True
        except Exception:
            md = smf.mixedlm(f"ICTEFFIC ~ {FIX}", g,
                             groups=g.CNTSCHID).fit(reml=True)
            conv = False
        percountry[c] = {**dump(md), "N": len(g),
                         "J": g.CNTSCHID.nunique(), "converged": conv}
    json.dump(percountry, open(RESULTS / "res_percountry.json", "w"), indent=1)

    # ---- pooled models with system fixed effects ----
    nt, ns = icc["Pooled (country FE)"]["tau00"], icc["Pooled (country FE)"]["sigma2"]
    m1 = smf.mixedlm(
        "ICTEFFIC ~ C(country) + ESCS_c + girl + immig + ICTQUAL_c + "
        "ICTHOME_c + ICTWKDY_c + ICTWKEND_c + ICTOUT_c + AGE_c + GRADE_c",
        m, groups=m.CNTSCHID).fit(reml=True)
    m2 = smf.mixedlm(f"ICTEFFIC ~ C(country) + {FIX}",
                     m, groups=m.CNTSCHID).fit(reml=True)
    m3 = smf.mixedlm(f"ICTEFFIC ~ C(country) + {FIX} + "
                     "sch_escs_g:ESCS_c + sch_escs_g:girl",
                     m, groups=m.CNTSCHID).fit(reml=True)
    json.dump({"M1": dump(m1, nt, ns), "M2": dump(m2, nt, ns),
               "M3": dump(m3, nt, ns), "N": len(m),
               "J": int(m.CNTSCHID.nunique())},
              open(RESULTS / "res_pooled.json", "w"), indent=1)
    print("pooled models done")


if __name__ == "__main__":
    main()
