"""Preston curve, F6-dev-01. Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

SEED = 20261029
rng = np.random.default_rng(SEED)
OUT = "out"
os.makedirs(OUT, exist_ok=True)
kw = dict(keep_default_na=False, na_values=[""])  # 'NA' is not a missing marker
ind = pd.read_csv("data/wb_indicators.csv", **kw)
ent = pd.read_csv("data/wb_entities.csv", **kw)
n_rows = {"wb_indicators.csv": len(ind), "wb_entities.csv": len(ent)}
issues = []

# ---- audit ----
assert not ind.duplicated(["iso3", "year"]).any()
assert set(ind.iso3) <= set(ent.iso3)
d = ind.merge(ent[["iso3", "name", "region", "income_level"]], on="iso3", suffixes=("", "_meta"))
nm = d.loc[d.name != d.name_meta, "iso3"].unique().tolist()
if nm:
    issues.append({"issue": f"Entity names differ between files for {nm} (trailing whitespace in metadata names; all are aggregates)",
                   "handling": "Joined on iso3 only; names ignored."})
n_agg = int(ind.iso3.isin(ent.iso3[ent.region == "Aggregates"]).sum())
issues.append({"issue": f"wb_indicators.csv mixes countries with regional/income aggregates ({n_agg} aggregate rows, {ind.iso3[ind.iso3.isin(ent.iso3[ent.region=='Aggregates'])].nunique()} aggregate entities)",
               "handling": "Excluded all entities with region == 'Aggregates' per the task rule."})
c = d[d.region != "Aggregates"].copy()
c["complete"] = c.gdp_pc_ppp.notna() & c.life_expectancy.notna()
per_year = c.groupby("year").complete.sum()
year = int(per_year[per_year >= 150].index.max())

# implausible LE volatility (audit, not used for estimation unless it hits the analysis year)
s = c.sort_values(["iso3", "year"]).copy()
s["dle"] = s.groupby("iso3").life_expectancy.diff()
vol = s[s.dle.abs() > 10]
for iso, g in vol.groupby("iso3"):
    ser = s[s.iso3 == iso].set_index("year").life_expectancy.round(1).to_dict()
    issues.append({"issue": f"{iso}: life expectancy jumps by >10 years between consecutive years (ending in {g.year.tolist()}); series 2015-2023 = {ser}",
                   "handling": f"Only the {year} cross-section is used; values kept as published (no basis to overwrite official data). Robustness: slope re-estimated excluding all such entities (see out/sensitivity.json)."})

y = c[c.year == year]
miss = y[~y.complete]
n_excl = int(len(miss))
issues.append({"issue": f"{n_excl} countries/economies lack an indicator in {year} (all missing GDP pc PPP): {', '.join(miss.iso3)}; missingness is not random (small territories/microstates and conflict/sanctioned economies)",
               "handling": "Listwise exclusion, no imputation (task rule); selection bias discussed in report."})
a = y[y.complete].copy()
a["lny"] = np.log(a.gdp_pc_ppp)
assert (a.gdp_pc_ppp > 0).all()

# ---- estimate ----
X = sm.add_constant(a.lny.values)
le = a.life_expectancy.values
fit = sm.OLS(le, X).fit(cov_type="HC3")
alpha, beta = fit.params
se = fit.bse[1]
ci = fit.conf_int(0.05)[1]  # normal-based (z) interval for robust cov
resid = fit.resid
baseline = float(np.sqrt(np.mean((le - le.mean()) ** 2)))
model = float(np.sqrt(np.mean(resid ** 2)))
r = float(np.corrcoef(a.lny, le)[0, 1])
concl = "positive_effect" if ci[0] > 0 else ("negative_effect" if ci[1] < 0 else "no_detectable_effect")

# ---- sensitivity (reported in report, saved separately) ----
sens = {}
for hc in ["HC0", "HC1", "HC2", "HC3"]:
    sens[f"se_{hc}"] = float(sm.OLS(le, X).fit(cov_type=hc).bse[1])
B = 2000
bs = np.empty(B)
n = len(a)
for b in range(B):
    i = rng.integers(0, n, n)
    bs[b] = np.polyfit(a.lny.values[i], le[i], 1)[0]
sens["pairs_bootstrap_ci"] = [float(np.percentile(bs, 2.5)), float(np.percentile(bs, 97.5))]
sub = ~a.iso3.isin(vol.iso3.unique())
sens["beta_excl_volatile"] = float(np.polyfit(a.lny[sub], le[sub], 1)[0])
infl = sm.OLS(le, X).fit().get_influence()
cook = infl.cooks_distance[0]
keep = cook < 4 / n
sens["beta_excl_cooks_gt_4_over_n"] = float(np.polyfit(a.lny[keep], le[keep], 1)[0])
sens["n_cooks_gt_4_over_n"] = int((~keep).sum())
sens["top_influential"] = a.iso3.values[np.argsort(-cook)[:5]].tolist()
q = sm.OLS(le, sm.add_constant(np.column_stack([a.lny, a.lny ** 2]))).fit(cov_type="HC3")
sens["quadratic_coef_lny2"] = float(q.params[2]); sens["quadratic_coef_lny2_p"] = float(q.pvalues[2])
sens["largest_neg_residuals"] = a.assign(e=resid).nsmallest(5, "e")[["iso3", "e"]].round(2).values.tolist()
for name, f in [("beta_year_2019_prepandemic", 2019)]:
    z = c[(c.year == f) & c.complete]
    sens[name] = float(np.polyfit(np.log(z.gdp_pc_ppp), z.life_expectancy, 1)[0])
with open(f"{OUT}/sensitivity.json", "w") as fh:
    json.dump(sens, fh, indent=2)

# ---- table & figure ----
order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
t = (a.groupby("income_level").agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("lny", "mean"),
                                   mean_life_expectancy=("life_expectancy", "mean"))
     .reindex([o for o in order if o in set(a.income_level)]).reset_index())
t.round(4).to_csv(f"{OUT}/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
cols = dict(zip(order, ["#d62728", "#ff7f0e", "#2ca02c", "#1f77b4"]))
for lev in order:
    m = a.income_level == lev
    ax.scatter(a.lny[m], le[m], s=14, alpha=.75, color=cols[lev], label=f"{lev} (n={m.sum()})")
xs = np.linspace(a.lny.min(), a.lny.max(), 100)
ax.plot(xs, alpha + beta * xs, "k-", lw=1.5, label=f"OLS: LE = {alpha:.1f} + {beta:.2f}·ln(GDPpc)")
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={n} countries/economies; R²={fit.rsquared:.2f})")
ax.legend(fontsize=8, loc="lower right")
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=150)

res = {
    "task_id": "F6-dev-01", "analysis_year": year, "n_excluded_missing": n_excl,
    "baseline_metric": baseline, "model_metric": model, "pearson_r": r,
    "r_squared": float(fit.rsquared), "effect_estimate": float(beta), "intercept": float(alpha),
    "effect_se": float(se), "effect_ci_low": float(ci[0]), "effect_ci_high": float(ci[1]),
    "conclusion": concl, "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": n_rows, "n_analysed": int(n), "random_seed": SEED,
}
with open(f"{OUT}/results.json", "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(sens, indent=2)); print(t)