import warnings
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
from scipy.stats import spearmanr, mannwhitneyu, pearsonr, t as t_dist

warnings.filterwarnings("ignore")
pd.set_option("display.width", 160)
pd.set_option("display.max_columns", 12)
pd.set_option("display.float_format", "{:,.2f}".format)

# 🎨 palette and chart style
BLUE, ORANGE, AQUA, VIOLET = "#2a78d6", "#eb6834", "#1baf7a", "#6250d6"
INK, MUTED, GRID, SURFACE = "#0b0b0b", "#52514e", "#e5e4df", "#fcfcfb"
plt.rcParams.update({
    "figure.dpi": 110, "figure.facecolor": SURFACE, "axes.facecolor": SURFACE,
    "axes.edgecolor": GRID, "axes.labelcolor": MUTED, "axes.titlecolor": INK,
    "axes.titlesize": 11, "axes.titleweight": "bold", "axes.titlelocation": "left",
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.grid": True, "grid.color": GRID, "grid.linewidth": .8,
    "xtick.color": MUTED, "ytick.color": MUTED, "font.size": 10,
    "legend.frameon": False, "lines.linewidth": 2,
})

# 📁 paths: works locally (./data) and on Kaggle (/kaggle/input/...)
def find_file(pattern):
    for base in [Path("data"), Path("."), Path("/kaggle/input")]:
        hits = sorted(base.rglob(pattern)) if base.exists() else []
        if hits:
            return hits[0]
    raise FileNotFoundError(pattern)

OUT_DIR = Path("/kaggle/working") if Path("/kaggle/working").exists() else Path("output")
FIG_DIR = OUT_DIR / "figures"
FIG_DIR.mkdir(parents=True, exist_ok=True)

# 🖼️ numbered figures, saved as PNG
_FIG = {"n": 0}
def figure_title(fig, text):
    _FIG["n"] += 1
    fig.suptitle(f"Figure {_FIG['n']} · {text}", x=0.01, ha="left", fontsize=13, fontweight="bold", color=INK)
    return _FIG["n"]

def save(fig, name):
    fig.savefig(FIG_DIR / f"fig{_FIG['n']:02d}_{name}.png", dpi=150, bbox_inches="tight", facecolor=SURFACE)
    plt.show()

usd = FuncFormatter(lambda v, _: f"${v:,.0f}")

# 📐 statistics helpers
def fmt_p(p):
    return "< 0.001" if p < 0.001 else f"{p:.3f}"

def spearman_ci(x, y, n_boot=2000, seed=42):
    """Spearman rho with a percentile bootstrap 95% CI."""
    x, y = np.asarray(x), np.asarray(y)
    rho, p = spearmanr(x, y)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, len(x), (n_boot, len(x)))
    boots = [spearmanr(x[i], y[i])[0] for i in idx]
    lo, hi = np.nanpercentile(boots, [2.5, 97.5])
    return rho, p, lo, hi

def partial_spearman(x, y, z):
    """Rank correlation of x and y after removing the linear effect of rank(z) from both."""
    rx, ry, rz = (pd.Series(v).rank().to_numpy() for v in (x, y, z))
    res = lambda a: a - np.polyval(np.polyfit(rz, a, 1), rz)
    r, _ = pearsonr(res(rx), res(ry))
    n = len(rx)
    p = 2 * t_dist.sf(abs(r) * np.sqrt((n - 3) / (1 - r**2)), n - 3)
    return r, p

print("✅ setup done · figures go to", FIG_DIR)

DATA_FILE = find_file("*_Data.csv")
META_FILE = find_file("*Metadata.csv")

raw = pd.read_csv(DATA_FILE, na_values="..")
footer = raw[raw["Series Code"].isna()]
df = raw.dropna(subset=["Series Code"]).copy()
df.columns = [c.split(" [")[0] for c in df.columns]          # "2015 [YR2015]" -> "2015"
YEARS = [c for c in df.columns if c.isdigit()]
ID_COLS = [c for c in df.columns if not c.isdigit()]

print(f"Raw rows: {len(raw):,}  ->  data rows: {len(df):,}  (dropped {len(footer)} footer/blank lines)")
print("Footer text:", [s for s in footer["Country Name"].dropna()])
print(f"Year columns: {YEARS[0]} to {YEARS[-1]} ({len(YEARS)} years), all numeric: {all(df[y].dtype.kind == 'f' for y in YEARS)}")

# The metadata file repeats the data, then has a definitions block starting at a row whose first cell is "Code"
meta_raw = pd.read_csv(META_FILE, header=None, dtype=str, na_values="..")
start = meta_raw.index[meta_raw[0].eq("Code")][0]
DEFINITIONS = (meta_raw.iloc[start + 1:, :4]
               .set_axis(meta_raw.iloc[start, :4].tolist(), axis=1)
               .set_index("Code"))

# Integrity check: is the data part of the metadata file identical to the data file?
meta_data = pd.read_csv(META_FILE, na_values="..", nrows=start - 1).dropna(subset=["Series Code"])
meta_data.columns = [c.split(" [")[0] for c in meta_data.columns]
same = meta_data[YEARS].reset_index(drop=True).equals(df[YEARS].reset_index(drop=True))
print("Data block of the metadata file identical to the data file:", same)

DEFINITIONS[["Indicator Name", "Short definition"]].style.set_properties(**{"text-align": "left", "white-space": "normal"})

AGGREGATES = {
    "LIC": "Low income", "LMC": "Lower middle income", "LMY": "Low & middle income",
    "LDC": "Least developed countries (UN)", "MNA": "Middle East, North Africa, Afghanistan & Pakistan (excl. high income)",
}
names = df.drop_duplicates("Country Code").set_index("Country Code")["Country Name"]
countries = df[~df["Country Code"].isin(AGGREGATES)]
aggregates = df[df["Country Code"].isin(AGGREGATES)]

profile = pd.Series({
    "Rows (country × indicator)": f"{len(df):,}",
    "Indicators": df["Series Code"].nunique(),
    "Economies": countries["Country Code"].nunique(),
    "Aggregates": aggregates["Country Code"].nunique(),
    "Year columns": f"{YEARS[0]}–{YEARS[-1]}",
    "Share of cells with a value (economies)": f"{countries[YEARS].notna().mean().mean():.1%}",
}, name="value").to_frame()
display(profile)
names.loc[list(AGGREGATES)].to_frame("Aggregate name")

VARIABLES = {
    "SE.PRM.UNER.ZS":    "Children out of School",
    "SE.SEC.UNER.LO.ZS": "Adolescents out of School",
    "SE.XPD.PRIM.PC.ZS": "Govt Primary Expenditure",
    "SE.XPD.SECO.PC.ZS": "Govt Secondary Expenditure",
    "NY.GDP.PCAP.CD":    "GDP Per Capita",
}
SHOW_YEARS = [y for y in YEARS if int(y) >= 1990]
cov = (countries.set_index("Series Code")[SHOW_YEARS].notna()
       .groupby(level=0).sum().loc[list(VARIABLES)].rename(index=VARIABLES))

fig, ax = plt.subplots(figsize=(12, 3.4))
im = ax.imshow(cov.values, aspect="auto", cmap="Blues", vmin=0, vmax=217)
ax.set_yticks(range(len(cov)), cov.index)
ax.set_xticks(range(0, len(SHOW_YEARS), 5), SHOW_YEARS[::5])
ax.grid(False)
for i, row in enumerate(cov.values):           # label every 5th year so the numbers are readable
    for j in range(0, len(SHOW_YEARS), 5):
        ax.text(j, i, row[j], ha="center", va="center", fontsize=8, color="white" if row[j] > 120 else INK)
fig.colorbar(im, ax=ax, label="economies reporting (of 217)", fraction=.025, pad=.01)
figure_title(fig, "Spending data stops in 2018; out-of-school rates run to 2024")
fig.tight_layout()
save(fig, "coverage_heatmap")

trend = (aggregates[aggregates["Country Code"].isin(["LIC", "LMC"])
                    & aggregates["Series Code"].isin(["SE.PRM.UNER.ZS", "SE.SEC.UNER.LO.ZS"])]
         .set_index(["Country Code", "Series Code"])[[y for y in YEARS if 2000 <= int(y) <= 2024]])
yrs = trend.columns.astype(int)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.2), sharey=True)
for ax, code in zip(axes, ["SE.PRM.UNER.ZS", "SE.SEC.UNER.LO.ZS"]):
    for grp, color in [("LIC", ORANGE), ("LMC", BLUE)]:
        s = trend.loc[(grp, code)]
        ax.plot(yrs, s.values, color=color)
        ax.annotate(f"{AGGREGATES[grp]}\n{s.iloc[-1]:.0f}% in 2024", (yrs[-1], s.iloc[-1]),
                    xytext=(6, 0), textcoords="offset points", va="center", fontsize=8.5, color=INK)
        lo_year = s.idxmin()
        if grp == "LIC" and code == "SE.SEC.UNER.LO.ZS":
            ax.scatter(int(lo_year), s.min(), s=40, color=color, edgecolor=SURFACE, lw=2, zorder=3)
            ax.annotate(f"low point {lo_year}: {s.min():.1f}%", (int(lo_year), s.min()),
                        xytext=(0, -16), textcoords="offset points", ha="center", fontsize=8, color=MUTED)
    ax.set_title(VARIABLES[code] + " (% of age group)")
    ax.set_xlim(2000, 2031)
    ax.set_ylim(0, 55)
figure_title(fig, "Progress stalled in low-income countries after the early 2010s")
fig.tight_layout()
save(fig, "income_group_trends")

trend.round(1).loc[:, ["2000", "2010", "2015", "2020", "2024"]].rename(index=AGGREGATES | VARIABLES)

def wide_for(year, frame=countries):
    w = frame.pivot(index="Country Code", columns="Series Code", values=year)
    return w.rename(columns=VARIABLES)[list(VARIABLES.values())]

rows = []
for year in map(str, range(2008, 2019)):
    w = wide_for(year)
    rows.append({"Year": year, **w.notna().sum().to_dict(), "Complete (all 5)": len(w.dropna())})
year_coverage = pd.DataFrame(rows).set_index("Year")

MIN_N = 60
YEAR = max(y for y, n in year_coverage["Complete (all 5)"].items() if n >= MIN_N)
print(f"Rule: most recent year with at least {MIN_N} complete economies  ->  YEAR = {YEAR}")
year_coverage.style.highlight_max("Complete (all 5)", color="#cfe2f8").apply(
    lambda r: ["font-weight:bold" if r.name == YEAR else "" for _ in r], axis=1)

wide = wide_for(YEAR)
missing = pd.DataFrame({"Missing": wide.isna().sum(), "% missing": wide.isna().mean() * 100})
print(f"{YEAR}: {wide.shape[0]} economies × {wide.shape[1]} indicators")
missing

clean = wide.dropna().sort_values("Children out of School", ascending=False)
dropped = wide.drop(clean.index)
gdp_kept, gdp_drop = clean["GDP Per Capita"], dropped["GDP Per Capita"].dropna()
u, p_mw = mannwhitneyu(gdp_kept, gdp_drop)

pd.DataFrame({
    "Economies": [len(clean), len(dropped)],
    "with GDP data": [gdp_kept.size, gdp_drop.size],
    "Median GDP per capita (US$)": [gdp_kept.median(), gdp_drop.median()],
}, index=["Kept (all 5)", "Dropped"]).style.format({"Median GDP per capita (US$)": "${:,.0f}"}).set_caption(
    f"Mann-Whitney U test on GDP per capita, kept vs dropped: p = {fmt_p(p_mw)}")

fig, ax = plt.subplots(figsize=(11, 3.6))
rng = np.random.default_rng(1)
for i, (label, s, color) in enumerate([("Dropped", gdp_drop, ORANGE), ("Kept (all 5)", gdp_kept, BLUE)]):
    jitter = i + rng.uniform(-.18, .18, len(s))
    ax.scatter(s, jitter, s=22, alpha=.55, color=color, edgecolor=SURFACE, lw=.6)
    ax.plot([s.median()] * 2, [i - .3, i + .3], color=INK, lw=2.5)
    ax.annotate(f"median ${s.median():,.0f}  (n = {len(s)})", (s.median(), i + .32), fontsize=8.5, color=INK, ha="center")
ax.set_xscale("log")
ax.xaxis.set_major_formatter(usd)
ax.set_yticks([0, 1], ["Dropped", "Kept (all 5)"])
ax.set_ylim(-.6, 1.7)
ax.set_xlabel(f"GDP per capita {YEAR} (US$, log scale)")
ax.grid(axis="y", visible=False)
figure_title(fig, "Who survives listwise deletion?")
fig.tight_layout()
save(fig, "selection_bias")

TESTS = [
    ("GDP Per Capita",             "Children out of School",    "H1"),
    ("GDP Per Capita",             "Adolescents out of School", "H1"),
    ("Govt Primary Expenditure",   "Children out of School",    "H2"),
    ("Govt Secondary Expenditure", "Adolescents out of School", "H2"),
]
def run_tests(data):
    out = []
    for x, y, h in TESTS:
        rho, p, lo, hi = spearman_ci(data[x], data[y])
        out.append({"Hypothesis": h, "Explanatory": x, "Outcome": y, "rho": rho,
                    "CI low": lo, "CI high": hi, "p-value": fmt_p(p), "n": len(data)})
    return pd.DataFrame(out)

results = run_tests(clean)
results

fig, ax = plt.subplots(figsize=(10, 3.6))
labels = [f"{h}: {x.replace('Govt ', '').replace(' Expenditure', ' spending')} → {y.split(' out')[0]}"
          for h, x, y in zip(results["Hypothesis"], results["Explanatory"], results["Outcome"])]
ys = np.arange(len(results))[::-1]
for yy, (_, r) in zip(ys, results.iterrows()):
    color = BLUE if r["Hypothesis"] == "H1" else ORANGE
    ax.plot([r["CI low"], r["CI high"]], [yy, yy], color=color, lw=3, solid_capstyle="round")
    ax.scatter(r["rho"], yy, s=70, color=color, edgecolor=SURFACE, lw=2, zorder=3)
    ax.annotate(f"{r['rho']:.2f}", (r["rho"], yy), xytext=(0, 9), textcoords="offset points", ha="center", fontsize=8.5)
ax.axvline(0, color=MUTED, lw=1, ls="--")
ax.set_yticks(ys, labels)
ax.set_xlim(-1, .5)
ax.set_xlabel("Spearman rho with bootstrap 95% CI")
ax.grid(axis="y", visible=False)
figure_title(fig, f"Wealth separates countries far more than spending effort ({YEAR}, n = {len(clean)})")
fig.tight_layout()
save(fig, "forest_spearman")

fig, axes = plt.subplots(2, 2, figsize=(12, 8.5))
label_codes = list(clean.index[:5]) + list(clean.sort_values("GDP Per Capita").index[-2:])
for ax, (x, y, h) in zip(axes.ravel(), TESTS):
    color = BLUE if h == "H1" else ORANGE
    rho, p = spearmanr(clean[x], clean[y])
    ax.scatter(clean[x], clean[y], s=34, alpha=.8, color=color, edgecolor=SURFACE, lw=.8)
    if x == "GDP Per Capita":
        ax.set_xscale("log")
        ax.xaxis.set_major_formatter(usd)
    for code in label_codes:
        ax.annotate(code, (clean.loc[code, x], clean.loc[code, y]), xytext=(4, 3),
                    textcoords="offset points", fontsize=7.5, color=MUTED)
    ax.set_xlabel(x + (" (US$, log scale)" if x == "GDP Per Capita" else " (% of GDP per capita)"))
    ax.set_ylabel(y + " (%)")
    ax.set_title(f"{h} · rho = {rho:.2f}, p {'< 0.001' if p < .001 else '= ' + format(p, '.3f')}")
figure_title(fig, f"Out-of-school rates against wealth and spending effort, {YEAR}")
fig.tight_layout()
save(fig, "scatter_panels")

clean["GDP quartile"] = pd.qcut(clean["GDP Per Capita"], 4, labels=["Q1 poorest", "Q2", "Q3", "Q4 richest"])
quart = clean.groupby("GDP quartile", observed=True).agg(
    countries=("GDP Per Capita", "size"),
    gdp_range=("GDP Per Capita", lambda s: f"${s.min():,.0f}–${s.max():,.0f}"),
    children=("Children out of School", "median"),
    adolescents=("Adolescents out of School", "median"),
    primary_effort=("Govt Primary Expenditure", "median"),
)

fig, ax = plt.subplots(figsize=(10, 4))
xpos, w = np.arange(len(quart)), .36
for off, col, color, lab in [(-w / 2 - .01, "children", BLUE, "Children (primary age)"),
                             (w / 2 + .01, "adolescents", AQUA, "Adolescents (lower secondary age)")]:
    bars = ax.bar(xpos + off, quart[col], w, color=color, label=lab)
    ax.bar_label(bars, fmt="%.1f%%", fontsize=8.5, padding=2, color=INK)
ax.set_xticks(xpos, [f"{q}\n{r}".replace("$", r"\$") for q, r in zip(quart.index, quart["gdp_range"])])
ax.set_ylabel("median % out of school")
ax.legend(loc="upper right")
ax.grid(axis="x", visible=False)
figure_title(fig, "The poorest quarter of countries has by far the most children out of school")
fig.tight_layout()
save(fig, "gdp_quartiles")
quart

clean["Primary spend per student (US$)"] = clean["Govt Primary Expenditure"] / 100 * clean["GDP Per Capita"]
clean["Secondary spend per student (US$)"] = clean["Govt Secondary Expenditure"] / 100 * clean["GDP Per Capita"]

rows = []
for spend, outcome in [("Govt Primary Expenditure", "Children out of School"),
                       ("Govt Secondary Expenditure", "Adolescents out of School"),
                       ("Primary spend per student (US$)", "Children out of School"),
                       ("Secondary spend per student (US$)", "Adolescents out of School")]:
    rho, p = spearmanr(clean[spend], clean[outcome])
    pr, pp = partial_spearman(clean[spend], clean[outcome], clean["GDP Per Capita"])
    rows.append({"Spending measure": spend, "Outcome": outcome,
                 "rho": rho, "p": fmt_p(p), "partial rho | GDP": pr, "partial p": fmt_p(pp)})
rows.append({"Spending measure": "Primary spend per student (US$)", "Outcome": "GDP Per Capita",
             "rho": spearmanr(clean["Primary spend per student (US$)"], clean["GDP Per Capita"])[0],
             "p": "", "partial rho | GDP": np.nan, "partial p": ""})
pd.DataFrame(rows)

h1_cols = ["GDP Per Capita", "Children out of School", "Adolescents out of School"]
h1_cov = pd.Series({y: len(wide_for(y)[h1_cols].dropna()) for y in map(str, range(2018, 2025))}, name="complete economies")
RECENT = h1_cov.idxmax()
recent = wide_for(RECENT)[h1_cols].dropna()
print(f"Best-covered recent year for H1: {RECENT} (n = {len(recent)})")
display(h1_cov.to_frame().T)

recent_results = pd.DataFrame([
    dict(zip(["rho", "p", "CI low", "CI high"], spearman_ci(recent["GDP Per Capita"], recent[y]))) | {"Outcome": y}
    for y in h1_cols[1:]
]).set_index("Outcome")
recent_results["p"] = recent_results["p"].map(fmt_p)
recent_results

fig, ax = plt.subplots(figsize=(11, 5))
for y, color in [("Adolescents out of School", AQUA), ("Children out of School", BLUE)]:
    ax.scatter(recent["GDP Per Capita"], recent[y], s=26, alpha=.7, color=color, edgecolor=SURFACE, lw=.6,
               label=y.replace(" out of School", f" (rho = {spearmanr(recent['GDP Per Capita'], recent[y])[0]:.2f})"))
# rolling median of the adolescent rate along log GDP, to show the shape
srt = recent.sort_values("GDP Per Capita")
roll = srt["Adolescents out of School"].rolling(21, center=True, min_periods=11).median()
ax.plot(srt["GDP Per Capita"], roll, color=INK, lw=2, label="Adolescents, rolling median (21 countries)")
for code in srt.sort_values("Adolescents out of School").index[-5:]:
    ax.annotate(f"{code}", (recent.loc[code, "GDP Per Capita"], recent.loc[code, "Adolescents out of School"]),
                xytext=(4, 2), textcoords="offset points", fontsize=7.5, color=MUTED)
ax.set_xscale("log")
ax.xaxis.set_major_formatter(usd)
ax.set_xlabel(f"GDP per capita {RECENT} (US$, log scale)")
ax.set_ylabel("% out of school")
ax.legend(loc="upper right")
figure_title(fig, f"The wealth gradient still holds in {RECENT}, with a much larger sample (n = {len(recent)})")
fig.tight_layout()
save(fig, "recent_h1")

by_year = []
for year in map(str, range(2008, 2018)):
    w = wide_for(year).dropna()
    if len(w) < 40:
        continue
    r = run_tests(w)
    r["Year"] = int(year)
    by_year.append(r)
by_year = pd.concat(by_year)

fig, axes = plt.subplots(1, 4, figsize=(14, 3.8), sharey=True)
for ax, (x, y, h) in zip(axes, TESTS):
    s = by_year[(by_year["Explanatory"] == x) & (by_year["Outcome"] == y)]
    color = BLUE if h == "H1" else ORANGE
    ax.fill_between(s["Year"], s["CI low"], s["CI high"], color=color, alpha=.15, lw=0)
    ax.plot(s["Year"], s["rho"], color=color, marker="o", ms=5, mec=SURFACE)
    ax.axhline(0, color=MUTED, lw=1, ls="--")
    ax.set_title(f"{h}: {x.split(' ')[1] if h == 'H2' else 'GDP'} → {y.split(' ')[0]}", fontsize=10)
    ax.set_xticks(s["Year"][::2])
axes[0].set_ylabel("Spearman rho (95% CI)")
figure_title(fig, "The wealth link is stable across years; the spending-effort link stays weak")
fig.tight_layout()
save(fig, "robustness_years")
by_year.pivot_table(index="Year", columns=["Hypothesis", "Outcome", "Explanatory"], values="rho").round(2).assign(
    n=by_year.groupby("Year")["n"].first())

# Latest value per country and indicator within 2014–2018
WINDOW = [str(y) for y in range(2014, 2019)]
latest = (countries.set_index(["Country Code", "Series Code"])[WINDOW]
          .apply(lambda r: r.dropna().iloc[-1] if r.notna().any() else np.nan, axis=1)
          .unstack().rename(columns=VARIABLES)[list(VARIABLES.values())].dropna())
print(f"Latest value in {WINDOW[0]}–{WINDOW[-1]}: n = {len(latest)} economies (vs {len(clean)} in {YEAR} alone)")
window_results = run_tests(latest)
window_results["partial rho | GDP"] = [
    partial_spearman(latest[x], latest[y], latest["GDP Per Capita"])[0] if h == "H2" else np.nan for x, y, h in TESTS]
window_results["partial p"] = [
    fmt_p(partial_spearman(latest[x], latest[y], latest["GDP Per Capita"])[1]) if h == "H2" else "" for x, y, h in TESTS]
window_results

import shutil
out_cols = list(VARIABLES.values()) + ["Primary spend per student (US$)", "Secondary spend per student (US$)", "GDP quartile"]
clean[out_cols].to_csv(OUT_DIR / f"school_dropout_{YEAR}_clean.csv")
recent.to_csv(OUT_DIR / f"school_dropout_{RECENT}_h1.csv")
zip_path = shutil.make_archive(str(OUT_DIR / "figures"), "zip", FIG_DIR)
print("Saved:")
for f in sorted(OUT_DIR.glob("*.csv")) + [Path(zip_path)]:
    print("  ", f)
print(f"   {len(list(FIG_DIR.glob('*.png')))} figures in {FIG_DIR}")