"""Extract the Latin American subsample from the PISA 2022 international database.

Input : data/raw/CY08MSP_STU_QQQ.SAV, data/raw/CY08MSP_SCH_QQQ.SAV
        (download from https://www.oecd.org/en/data/datasets/pisa-2022-database.html)
Output: data/interim/stu_lac.pkl, data/interim/sch_lac.pkl
"""
from pathlib import Path

import pandas as pd
import pyreadstat

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

LAC = ["ARG", "BRA", "CHL", "COL", "CRI", "DOM", "SLV", "GTM",
       "JAM", "MEX", "PAN", "PRY", "PER", "URY"]

STU_COLS = (
    ["CNT", "CNTSCHID", "CNTSTUID",
     "ICTEFFIC", "ICTQUAL", "ICTHOME", "ICTWKDY", "ICTWKEND", "ICTOUT",
     "ESCS", "ST004D01T", "IMMIG", "AGE", "GRADE", "W_FSTUWT"]
    + [f"W_FSTURWT{i}" for i in range(1, 81)]
)
SCH_COLS = ["CNT", "CNTSCHID", "DIGPREP", "RATCMP1", "EDUSHORT",
            "SCHLTYPE", "SC001Q01TA", "SCHSIZE"]


def main() -> None:
    stu, _ = pyreadstat.read_sav(str(RAW / "CY08MSP_STU_QQQ.SAV"), usecols=STU_COLS)
    stu = stu[stu.CNT.isin(LAC)]
    stu.to_pickle(INTERIM / "stu_lac.pkl")
    print("students:", stu.shape)
    print(stu.groupby("CNT").agg(n=("CNTSTUID", "size"),
                                 icteffic_n=("ICTEFFIC", "count")))

    sch, _ = pyreadstat.read_sav(str(RAW / "CY08MSP_SCH_QQQ.SAV"), usecols=SCH_COLS)
    sch = sch[sch.CNT.isin(LAC)]
    sch.to_pickle(INTERIM / "sch_lac.pkl")
    print("schools:", sch.shape)


if __name__ == "__main__":
    main()
