"""
Replication code for:
"Learning, Schooling, and Economic Growth: Revisiting a Classic Result"
C. Kirabo Jackson, 2026

This script reconstructs the analysis reported in the blog post and produces
all three figures. It uses:
  * TIMSS 1999 math and science scores
  * PISA 2000 math and science scores
  * Barro-Lee average years of schooling (age 15+) in 2000
  * World Bank real GDP per capita (NY.GDP.PCAP.KD) in 2000 and 2024

The GDP outcome is annualized log growth from the 2000 endpoint to the 2024
endpoint. Equivalently, the 24 subsequent annual growth changes are 2000->2001
through 2023->2024, so the test years (1999/2000) do not overlap the subsequent
growth years (2001-2024).

Inference uses HC1 heteroskedasticity-robust standard errors and a Student-t
reference distribution with the OLS residual degrees of freedom. This matches
the p-values and confidence intervals reported in the post.

Requirements:
    pip install requests pandas numpy scipy matplotlib statsmodels

Run:
    python learning-schooling-economic-growth.py

Outputs:
    blog_fig1_schooling_growth.png
    blog_fig2_schooling_growth_cond_scores.png
    blog_fig3_scores_growth_cond_schooling.png
    analysis_testscores_growth_2000_2024.csv
"""

from __future__ import annotations

import io
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import requests
from scipy import stats
import statsmodels.api as sm

OUT = Path(".")
GDP_START = 2000
GDP_END = 2024

# -----------------------------------------------------------------------------
# 1. Historical test scores used in the analysis
# -----------------------------------------------------------------------------
TIMSS99_MATH = {
    "SGP":604,"KOR":587,"TWN":585,"HKG":582,"JPN":579,
    "BEL":558,"NLD":540,"SVK":534,"HUN":532,"CAN":531,
    "SVN":530,"RUS":526,"AUS":525,"FIN":520,"CZE":520,
    "MYS":519,"BGR":511,"LVA":505,"USA":502,"GBR":496,
    "NZL":491,"LTU":482,"ITA":479,"CYP":476,"ROU":472,
    "MDA":469,"THA":467,"ISR":466,"TUN":448,"TUR":429,
    "IDN":403,"CHL":392,"PHL":345,"MAR":337,"ZAF":275,
}
TIMSS99_SCI = {
    "TWN":569,"SGP":568,"HUN":552,"JPN":550,"KOR":549,
    "NLD":545,"AUS":540,"CZE":539,"GBR":538,"FIN":535,
    "SVK":535,"BEL":535,"SVN":533,"CAN":533,"HKG":530,
    "RUS":529,"BGR":518,"USA":515,"NZL":510,"LVA":503,
    "ITA":493,"MYS":492,"LTU":488,"ROU":472,"ISR":468,
    "CYP":460,"MDA":459,"THA":482,"TUN":430,"TUR":433,
    "IDN":435,"CHL":420,"PHL":345,"MAR":323,"ZAF":243,
}
PISA00_MATH = {
    "JPN":557,"KOR":547,"NZL":537,"FIN":536,"AUS":533,
    "CAN":533,"CHE":529,"GBR":529,"BEL":520,"FRA":517,
    "AUT":515,"DNK":514,"ISL":514,"SWE":510,"IRL":503,
    "NOR":499,"CZE":498,"USA":493,"DEU":490,"HUN":488,
    "RUS":478,"ESP":476,"POL":470,"LVA":462,"ITA":457,
    "PRT":454,"GRC":447,"LUX":446,"MEX":387,"BRA":334,
}
PISA00_SCI = {
    "KOR":552,"JPN":550,"FIN":538,"GBR":532,"CAN":529,
    "NZL":528,"AUS":528,"AUT":519,"IRL":513,"SWE":512,
    "CZE":511,"FRA":500,"NOR":500,"USA":499,"HUN":496,
    "ISL":496,"BEL":496,"CHE":496,"ESP":491,"DEU":487,
    "POL":483,"DNK":481,"ITA":478,"GRC":461,"RUS":460,
    "LVA":460,"LUX":443,"MEX":422,"BRA":375,
}


def build_scores() -> pd.DataFrame:
    all_iso = sorted(set(TIMSS99_MATH) | set(TIMSS99_SCI) |
                     set(PISA00_MATH) | set(PISA00_SCI))
    sc = pd.DataFrame(index=all_iso)
    sc["tm99m"] = pd.Series(TIMSS99_MATH)
    sc["tm99s"] = pd.Series(TIMSS99_SCI)
    sc["pi00m"] = pd.Series(PISA00_MATH)
    sc["pi00s"] = pd.Series(PISA00_SCI)

    # Standardize each assessment across participating countries first.
    for col in ["tm99m", "tm99s", "pi00m", "pi00s"]:
        v = sc[col].dropna()
        sc[col] = (sc[col] - v.mean()) / v.std(ddof=1)

    score_cols = ["tm99m", "tm99s", "pi00m", "pi00s"]
    sc["n_tests"] = sc[score_cols].notna().sum(axis=1)
    sc = sc.loc[sc["n_tests"] >= 2].copy()
    sc["test_score_raw"] = sc[score_cols].mean(axis=1)
    return sc[["test_score_raw", "n_tests"]].reset_index().rename(
        columns={"index": "iso3"}
    )


# -----------------------------------------------------------------------------
# 2. World Bank GDP per capita
# -----------------------------------------------------------------------------
def wb_year(indicator: str, year: int, label: str) -> pd.DataFrame:
    url = (
        f"https://api.worldbank.org/v2/country/all/indicator/{indicator}"
        f"?date={year}&format=json&per_page=20000"
    )
    r = requests.get(url, timeout=60)
    r.raise_for_status()
    payload = r.json()
    rows = [
        {
            "iso3": x["countryiso3code"],
            "country": x["country"]["value"],
            label: float(x["value"]),
        }
        for x in payload[1]
        if x["value"] is not None and x["countryiso3code"]
    ]
    return pd.DataFrame(rows)


# -----------------------------------------------------------------------------
# 3. Barro-Lee schooling in 2000 (male + female, age 15+)
# -----------------------------------------------------------------------------
BARRO_LEE_URL = (
    "https://raw.githubusercontent.com/barrolee/BarroLeeDataSet/"
    "refs/heads/master/BLData/BL2013_MF1599_v2.2.csv"
)


def load_schooling() -> pd.DataFrame:
    r = requests.get(BARRO_LEE_URL, timeout=60)
    r.raise_for_status()
    bl = pd.read_csv(io.BytesIO(r.content))
    return (
        bl.loc[(bl["year"] == 2000) & (bl["sex"] == "MF"),
               ["WBcode", "yr_sch", "region_code"]]
          .rename(columns={
              "WBcode": "iso3",
              "yr_sch": "schooling_2000",
              "region_code": "region",
          })
    )


# -----------------------------------------------------------------------------
# 4. HC1 inference using finite-sample t critical values
# -----------------------------------------------------------------------------
def fit_model(df: pd.DataFrame, regressors: list[str]):
    X = sm.add_constant(df[regressors])
    y = df["growth_2000_2024"]
    ols = sm.OLS(y, X).fit()
    hc1 = sm.OLS(y, X).fit(cov_type="HC1")
    return ols, hc1


def inference(ols, hc1, name: str) -> dict[str, float]:
    b = float(hc1.params[name])
    se = float(hc1.bse[name])
    df_resid = float(ols.df_resid)
    tstat = b / se
    p = float(2 * stats.t.sf(abs(tstat), df_resid))
    crit = float(stats.t.ppf(0.975, df_resid))
    return {
        "coef": b,
        "se": se,
        "p": p,
        "lo": b - crit * se,
        "hi": b + crit * se,
        "df": df_resid,
    }


def residualize(df: pd.DataFrame, variable: str, controls: list[str]) -> pd.Series:
    X = sm.add_constant(df[controls])
    return sm.OLS(df[variable], X).fit().resid


def annotate_points(ax, df, xcol: str, ycol: str):
    for _, row in df.iterrows():
        ax.annotate(
            row["iso3"], (row[xcol], row[ycol]), xytext=(3, 3),
            textcoords="offset points", fontsize=8, alpha=0.72
        )


def stats_box(ax, n: int, label: str, inf: dict[str, float]):
    text = (
        f"N = {n}\n"
        f"{label} = {inf['coef']:.3f} pp/year\n"
        f"HC1 SE = {inf['se']:.3f}\n"
        f"p = {inf['p']:.3f}\n"
        f"95% CI [{inf['lo']:.3f}, {inf['hi']:.3f}]"
    )
    ax.text(
        0.985, 0.04, text, transform=ax.transAxes, ha="right", va="bottom",
        fontsize=11, bbox=dict(boxstyle="round", alpha=0.70)
    )


def make_partial_plot(
    df: pd.DataFrame,
    xvar: str,
    controls: list[str],
    slope: float,
    inf: dict[str, float],
    title: str,
    subtitle: str,
    xlabel: str,
    stats_label: str,
    filename: str,
):
    work = df.copy()
    work["x_resid"] = residualize(work, xvar, controls)
    work["y_resid"] = residualize(work, "growth_2000_2024", controls)

    fig, ax = plt.subplots(figsize=(12, 8))
    ax.scatter(work["x_resid"], work["y_resid"], s=70, alpha=0.82)
    annotate_points(ax, work, "x_resid", "y_resid")

    xline = np.linspace(work["x_resid"].min() - 0.1,
                        work["x_resid"].max() + 0.1, 200)
    ax.plot(xline, slope * xline, linewidth=1.8)
    ax.axhline(0, linewidth=0.7, alpha=0.5)
    ax.axvline(0, linewidth=0.7, alpha=0.5)
    ax.set_title(title, loc="left", fontweight="bold", fontsize=19, pad=72)
    ax.text(0, 1.02, subtitle, transform=ax.transAxes, ha="left", va="bottom",
            fontsize=12)
    ax.set_xlabel(xlabel, fontsize=12)
    ax.set_ylabel(
        "Conditional average annual GDP per capita growth, 2000–2024 "
        "(percentage points)", fontsize=12
    )
    stats_box(ax, len(work), stats_label, inf)
    ax.grid(alpha=0.15)
    fig.tight_layout()
    fig.savefig(OUT / filename, dpi=220, bbox_inches="tight")
    plt.close(fig)


# -----------------------------------------------------------------------------
# 5. Build data, estimate models, and make figures
# -----------------------------------------------------------------------------
def main():
    scores = build_scores()
    gdp0 = wb_year("NY.GDP.PCAP.KD", GDP_START, "gdp_pc_2000")
    gdp1 = wb_year("NY.GDP.PCAP.KD", GDP_END, "gdp_pc_2024")
    schooling = load_schooling()

    df = (
        scores
        .merge(gdp0[["iso3", "country", "gdp_pc_2000"]], on="iso3")
        .merge(gdp1[["iso3", "gdp_pc_2024"]], on="iso3")
        .merge(schooling, on="iso3")
        .dropna()
        .copy()
    )

    # Average annual continuously-compounded real GDP-per-capita growth.
    df["growth_2000_2024"] = (
        100 * (np.log(df["gdp_pc_2024"]) - np.log(df["gdp_pc_2000"]))
        / (GDP_END - GDP_START)
    )
    df["log_gdp_pc_2000"] = np.log(df["gdp_pc_2000"])

    # A one-unit score difference is exactly one SD in the final estimation sample.
    df["test_score_sd"] = (
        (df["test_score_raw"] - df["test_score_raw"].mean())
        / df["test_score_raw"].std(ddof=1)
    )

    # Figure 1 / traditional schooling specification.
    ols_s, hc1_s = fit_model(df, ["log_gdp_pc_2000", "schooling_2000"])
    inf_s = inference(ols_s, hc1_s, "schooling_2000")

    # Figures 2 and 3 / joint schooling + learning specification.
    ols_f, hc1_f = fit_model(
        df, ["test_score_sd", "log_gdp_pc_2000", "schooling_2000"]
    )
    inf_school_full = inference(ols_f, hc1_f, "schooling_2000")
    inf_score = inference(ols_f, hc1_f, "test_score_sd")

    print("PRIMARY RESULTS")
    print(f"N = {len(df)}")
    print(
        f"Schooling, without scores: {inf_s['coef']:.6f}, "
        f"HC1 SE={inf_s['se']:.6f}, p={inf_s['p']:.6f}, "
        f"95% CI=[{inf_s['lo']:.6f}, {inf_s['hi']:.6f}]"
    )
    print(
        f"Schooling, with scores:    {inf_school_full['coef']:.6f}, "
        f"HC1 SE={inf_school_full['se']:.6f}, p={inf_school_full['p']:.6f}, "
        f"95% CI=[{inf_school_full['lo']:.6f}, {inf_school_full['hi']:.6f}]"
    )
    print(
        f"Test score, with schooling: {inf_score['coef']:.6f}, "
        f"HC1 SE={inf_score['se']:.6f}, p={inf_score['p']:.6f}, "
        f"95% CI=[{inf_score['lo']:.6f}, {inf_score['hi']:.6f}]"
    )
    mean_growth = float(df["growth_2000_2024"].mean())
    print(f"Mean annual GDP-per-capita growth = {mean_growth:.6f}%")
    print(f"1-SD score effect / mean growth = {100*inf_score['coef']/mean_growth:.2f}%")

    make_partial_plot(
        df=df,
        xvar="schooling_2000",
        controls=["log_gdp_pc_2000"],
        slope=inf_s["coef"],
        inf=inf_s,
        title="More schooling predicts higher subsequent economic growth",
        subtitle="GDP per capita growth, 2000–2024; conditioning on log GDP per capita in 2000",
        xlabel="Conditional years of schooling in 2000",
        stats_label="Schooling coefficient",
        filename="blog_fig1_schooling_growth.png",
    )

    make_partial_plot(
        df=df,
        xvar="schooling_2000",
        controls=["test_score_sd", "log_gdp_pc_2000"],
        slope=inf_school_full["coef"],
        inf=inf_school_full,
        title="Accounting for learning attenuates the schooling–growth relationship",
        subtitle="GDP per capita growth, 2000–2024; conditioning on test scores and log GDP per capita in 2000",
        xlabel="Conditional years of schooling in 2000",
        stats_label="Schooling coefficient",
        filename="blog_fig2_schooling_growth_cond_scores.png",
    )

    make_partial_plot(
        df=df,
        xvar="test_score_sd",
        controls=["schooling_2000", "log_gdp_pc_2000"],
        slope=inf_score["coef"],
        inf=inf_score,
        title="Test scores strongly predict subsequent economic growth",
        subtitle="1999/2000 test scores; GDP per capita growth, 2000–2024; conditioning on schooling and initial GDP",
        xlabel="Conditional test score (SD units; 1999 TIMSS + 2000 PISA)",
        stats_label="1 SD score coefficient",
        filename="blog_fig3_scores_growth_cond_schooling.png",
    )

    df.to_csv(OUT / "analysis_testscores_growth_2000_2024.csv", index=False)
    print("\nSaved figures and analysis sample.")


if __name__ == "__main__":
    main()
