"""Build the analytic dataset for the six-system HLM analysis.

Sample decisions (documented in the manuscript):
- Seven LAC systems administered the optional ICT familiarity questionnaire:
  ARG, BRA, CHL, CRI, DOM, PAN, URY.
- Costa Rica (CRI) is excluded: ESCS is unavailable in the international database.
- DIGPREP was not administered in Panama; school digital resources are
  measured with RATCMP1 and EDUSHORT, available in all systems.
- Complete cases on all model variables.

Input : data/interim/stu_lac.pkl, data/interim/sch_lac.pkl
Output: data/processed/analysis_dataset.csv
"""
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
INTERIM = ROOT / "data" / "interim"
PROCESSED = ROOT / "data" / "processed"
PROCESSED.mkdir(parents=True, exist_ok=True)

SEL = ["ARG", "BRA", "CHL", "DOM", "PAN", "URY"]

MODEL_VARS = ["ESCS", "girl", "immig", "ICTQUAL", "ICTHOME", "ICTWKDY",
              "ICTWKEND", "ICTOUT", "AGE", "GRADE", "sch_escs", "RATCMP1",
              "EDUSHORT", "private", "rural"]

# System-level descriptors.
# GDP per capita PPP (current international $, 2022) and fixed broadband
# subscriptions per 100 people (2022): World Bank WDI (NY.GDP.PCAP.PP.CD,
# IT.NET.BBND.P2), retrieved July 2026. One-to-one program coding: author,
# from national policy documents.
SYSTEM = pd.DataFrame({
    "CNT": ["ARG", "BRA", "CHL", "CRI", "DOM", "PAN", "URY"],
    "country": ["Argentina", "Brazil", "Chile", "Costa Rica",
                "Dominican Rep.", "Panama", "Uruguay"],
    "gdppc": [29804, 19874, 30919, 26892, 24622, 36327, 33034],
    "broadband": [24.7, 21.0, 22.8, 21.8, 10.7, 16.4, 33.5],
    "one2one": ["Conectar Igualdad (2010, relaunched 2022)",
                "PROUCA (2010, discontinued)",
                "Enlaces (1992-2018, not 1:1)",
                "PRONIE MEP-FOD (1988- )",
                "Republica Digital (2016- )",
                "None sustained",
                "Plan Ceibal (2007- )"],
})


def main() -> None:
    stu = pd.read_pickle(INTERIM / "stu_lac.pkl")
    sch = pd.read_pickle(INTERIM / "sch_lac.pkl")

    stu = stu[stu.CNT.isin(SEL) & stu.ICTEFFIC.notna()].copy()

    # recodes
    stu["girl"] = (stu.ST004D01T == 1).astype(float)          # 1 = female
    stu["immig"] = np.where(stu.IMMIG.isna(), np.nan,
                            (stu.IMMIG > 1).astype(float))     # 2nd/1st gen = 1
    sch["private"] = np.where(sch.SCHLTYPE.isna(), np.nan,
                              (sch.SCHLTYPE < 3).astype(float))  # 1,2 = private
    sch["rural"] = np.where(sch.SC001Q01TA.isna(), np.nan,
                            (sch.SC001Q01TA == 1.0).astype(float))  # <3,000 people

    d = stu.merge(sch.drop(columns=["CNT"]), on="CNTSCHID", how="left")
    d["sch_escs"] = d.groupby("CNTSCHID").ESCS.transform("mean")
    d = d.merge(SYSTEM, on="CNT")

    # complete cases
    d = d.dropna(subset=MODEL_VARS).copy()

    # weights rescaled within system (mean 1)
    d["w1"] = d.groupby("CNT").W_FSTUWT.transform(lambda x: x / x.mean())

    # centering: group-mean (school) for level-1 continuous predictors,
    # grand-mean within system for level-2 predictors (Enders & Tofighi, 2007)
    for v in ["ESCS", "ICTQUAL", "ICTHOME", "ICTWKDY", "ICTWKEND",
              "ICTOUT", "AGE", "GRADE"]:
        d[v + "_c"] = d[v] - d.groupby("CNTSCHID")[v].transform("mean")
    for v in ["sch_escs", "RATCMP1", "EDUSHORT"]:
        d[v + "_g"] = d[v] - d.groupby("CNT")[v].transform("mean")

    d.to_csv(PROCESSED / "analysis_dataset.csv", index=False)
    SYSTEM.to_csv(PROCESSED / "system_level_data.csv", index=False)
    print("analytic sample:", len(d), "students;",
          d.CNTSCHID.nunique(), "schools")
    print(d.groupby("country").agg(N=("CNTSTUID", "size"),
                                   schools=("CNTSCHID", "nunique")))


if __name__ == "__main__":
    main()
