"""Replication entry point for F6-dev-01 (Preston curve, WDI).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/figure1.png, out/table1.csv, out/exclusions.csv.
No randomness is used (random_seed recorded as null)."""
import json, os
import numpy as np
import pandas as pd
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATA, OUT = "data", "out"
MIN_COUNTRIES = 150
HC = "HC3"
Z = 1.959963984540054  # 97.5% standard-normal quantile

ind = pd.read_csv(os.path.join(DATA, "wb_indicators.csv"), keep_default_na=False, na_values=[""])
ent = pd.read_csv(os.path.join(DATA, "wb_entities.csv"), keep_default_na=False, na_values=[""])
n_input_rows = {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))}
data_issues = []
excl = []

def log_ex(step, rule, before, after, why):
    excl.append(dict(step=step, rule=rule, n_before=int(before), n_after=int(after),
                     n_removed=int(before - after), rationale=why))

# --- integrity checks ---
assert not ind.duplicated(["iso3", "year"]).any(), "duplicate iso3-year rows"
assert not ent.duplicated("iso3").any(), "duplicate entity rows"
unmatched = sorted(set(ind.iso3) - set(ent.iso3))
if unmatched:
    data_issues.append({"issue": f"{len(unmatched)} indicator iso3 codes absent from entity metadata: {unmatched}",
                        "handling": "excluded (country status cannot be verified)"})
assert (ind.gdp_pc_ppp.dropna() > 0).all(), "non-positive GDP per capita"

m = ind.merge(ent[["iso3", "region", "income_level"]], on="iso3", how="left")
n0 = len(m)
m = m[m.region.notna()]
log_ex(1, "drop indicator rows whose iso3 is not in wb_entities.csv", n0, len(m),
       "country/aggregate status cannot be established")
n1 = len(m)
n_agg = m.loc[m.region == "Aggregates", "iso3"].nunique()
c = m[m.region != "Aggregates"].copy()
log_ex(2, "drop entities with region == 'Aggregates' (all years)", n1, len(c),
       f"information rule: countries/economies only ({n_agg} aggregate entities)")
data_issues.append({"issue": f"wb_indicators.csv contains {n_agg} regional/income/other aggregate entities "
                             "mixed with countries/economies",
                    "handling": "excluded via region == 'Aggregates' in wb_entities.csv"})

# --- analysis year selection ---
c["complete"] = c.gdp_pc_ppp.notna() & c.life_expectancy.notna()
counts = c.groupby("year")["complete"].sum().astype(int)
eligible = counts[counts >= MIN_COUNTRIES]
if eligible.empty:
    raise SystemExit("No year with >= 150 complete countries: premise invalid")
year = int(eligible.index.max())

n2 = len(c)
y = c[c.year == year].copy()
log_ex(3, f"keep analysis year {year} only", n2, len(y),
       f"most recent year with >= {MIN_COUNTRIES} countries having both indicators")
n3 = len(y)
miss = y[~y.complete]
d = y[y.complete].copy()
n_excluded_missing = int(len(miss))
log_ex(4, "drop countries with GDPpc or LE missing in analysis year (no imputation)", n3, len(d),
       "information rule; missing: " + "; ".join(miss.iso3))
miss_g = int(miss.gdp_pc_ppp.isna().sum()); miss_l = int(miss.life_expectancy.isna().sum())
data_issues.append({"issue": f"{n_excluded_missing} countries/economies in {year} lack an indicator "
                             f"(GDPpc missing: {miss_g}, LE missing: {miss_l}): {', '.join(miss.iso3)}; "
                             "mostly small territories and conflict/sanctioned economies, so missingness is not random",
                    "handling": "excluded from the estimate, counted in n_excluded_missing; no imputation"})

# anomalies outside the analysis year (documented, not used)
le_p = c.pivot(index="iso3", columns="year", values="life_expectancy")
jumps = le_p.diff(axis=1).abs().max(axis=1)
big = jumps[jumps > 10].sort_values(ascending=False)
if len(big):
    data_issues.append({"issue": "implausibly large year-to-year life-expectancy swings (>10 years) in: "
                        + ", ".join(f"{k} (max |dLE|={v:.1f})" for k, v in big.items())
                        + "; e.g. CAF LE 18.8 in 2022 vs 57.4 in 2023; PSE 2023 LE 65.2 vs 76.7 in 2022 (plausibly war-related; not verified)",
                        "handling": f"only {year} values enter the estimate; values kept as published (no editing); "
                                    "influence checked with leave-one-out slopes"})
aggname = ent[ent.name != ent.name.str.strip()]
if len(aggname):
    data_issues.append({"issue": f"{len(aggname)} entity names in wb_entities.csv carry trailing whitespace",
                        "handling": "cosmetic; joins use iso3 only"})

# --- estimation ---
d["log_gdp_pc"] = np.log(d.gdp_pc_ppp)
x = d.log_gdp_pc.to_numpy(); le = d.life_expectancy.to_numpy()
n = len(d)
baseline_rmse = float(np.sqrt(np.mean((le - le.mean()) ** 2)))
X = sm.add_constant(x)
fit = sm.OLS(le, X).fit(cov_type=HC)
alpha, beta = (float(v) for v in fit.params)
se = float(fit.bse[1])
resid = le - (alpha + beta * x)
model_rmse = float(np.sqrt(np.mean(resid ** 2)))
r = float(np.corrcoef(x, le)[0, 1])
r2 = float(fit.rsquared)
lo, hi = beta - Z * se, beta + Z * se
se_hc = {h: float(sm.OLS(le, X).fit(cov_type=h).bse[1]) for h in ["HC0", "HC1", "HC2", "HC3"]}
se_classical = float(sm.OLS(le, X).fit().bse[1])

# leave-one-out influence
loo = []
for i in range(n):
    k = np.arange(n) != i
    b = np.polyfit(x[k], le[k], 1)[0]
    loo.append(b)
loo = np.array(loo)
iso = d.iso3.to_numpy()

if lo > 0:
    concl = "positive_effect"
elif hi < 0:
    concl = "negative_effect"
else:
    concl = "no_detectable_effect"

results = {
    "task_id": "F6-dev-01",
    "analysis_year": year,
    "n_excluded_missing": n_excluded_missing,
    "baseline_metric": baseline_rmse,
    "model_metric": model_rmse,
    "pearson_r": r,
    "r_squared": r2,
    "effect_estimate": beta,
    "intercept": alpha,
    "effect_se": se,
    "effect_ci_low": float(lo),
    "effect_ci_high": float(hi),
    "conclusion": concl,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": int(n),
    "random_seed": None,
    # supplementary diagnostics
    "se_type": HC,
    "ci_method": "beta +/- 1.959964 * HC3 SE (normal approximation)",
    "complete_countries_by_year": {str(k): int(v) for k, v in counts.items()},
    "n_countries_in_year": int(n3),
    "n_aggregate_entities_excluded": int(n_agg),
    "effect_se_hc0": se_hc["HC0"], "effect_se_hc1": se_hc["HC1"],
    "effect_se_hc2": se_hc["HC2"], "effect_se_hc3": se_hc["HC3"],
    "effect_se_classical": se_classical,
    "loo_slope_min": float(loo.min()), "loo_slope_min_dropped": str(iso[loo.argmin()]),
    "loo_slope_max": float(loo.max()), "loo_slope_max_dropped": str(iso[loo.argmax()]),
    "rmse_reduction_share": float(1 - model_rmse / baseline_rmse),
    "le_gain_per_doubling_gdp": float(beta * np.log(2)),
    "mean_life_expectancy": float(le.mean()),
    "mean_log_gdp_pc": float(x.mean()),
}

with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

# --- table ---
order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
t = (d.groupby("income_level")
       .agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("log_gdp_pc", "mean"),
            mean_life_expectancy=("life_expectancy", "mean"))
       .reset_index())
t["k"] = t.income_level.map({v: i for i, v in enumerate(order)}).fillna(99)
t = t.sort_values(["k", "income_level"]).drop(columns="k")
assert t.n_countries.sum() == n
t.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

# --- exclusions ledger ---
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
    os.path.join(OUT, "exclusions.csv"), index=False)

# --- figure ---
fig, ax = plt.subplots(figsize=(7.5, 5.2))
cols = dict(zip(order, ["#d7191c", "#fdae61", "#abd9e9", "#2c7bb6"]))
for lvl in order:
    s = d[d.income_level == lvl]
    ax.scatter(s.log_gdp_pc, s.life_expectancy, s=18, alpha=0.8, color=cols[lvl], label=f"{lvl} (n={len(s)})")
xx = np.linspace(x.min(), x.max(), 100)
ax.plot(xx, alpha + beta * xx, color="k", lw=1.5,
        label=f"OLS: LE = {alpha:.1f} + {beta:.2f} ln(GDPpc), R² = {r2:.2f}")
ax.set_xlabel("ln(GDP per capita, PPP, constant 2021 int'l $)")
ax.set_ylabel("Life expectancy at birth (years)")
ax.set_title(f"Preston curve, {year}: {n} countries/economies (WDI)")
ax.legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)

print(f"year={year} n={n} excluded_missing={n_excluded_missing} beta={beta:.4f} se({HC})={se:.4f} "
      f"CI=[{lo:.4f},{hi:.4f}] r={r:.4f} R2={r2:.4f} RMSE base={baseline_rmse:.4f} model={model_rmse:.4f} -> {concl}")