"""Replication entry point for F6-dev-01 (Preston curve). Run from workspace root:
    ~/wp29/venv/bin/python out/analysis.py
"""
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

SEED = 0  # no randomness used; fixed for completeness
np.random.seed(SEED)
DATA, OUT = "data", "out"
os.makedirs(OUT, exist_ok=True)

# keep_default_na=False so that ISO code "NAM" (Namibia) etc. is never parsed as NaN
rd = dict(keep_default_na=False, na_values=[""])
ind = pd.read_csv(os.path.join(DATA, "wb_indicators.csv"), **rd)
ent = pd.read_csv(os.path.join(DATA, "wb_entities.csv"), **rd)
n_input_rows = {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))}
issues = []

# ---- audit ----
assert not ind.duplicated(["iso3", "year"]).any()
assert not ent.duplicated("iso3").any()
unmatched = sorted(set(ind.iso3) - set(ent.iso3))
assert not unmatched, unmatched
ent["region"] = ent["region"].str.strip()
ent["income_level"] = ent["income_level"].str.strip()
ent["name"] = ent["name"].str.strip()

m = ind.merge(ent[["iso3", "region", "income_level"]], on="iso3", how="left")
n_agg_ent = int(m.loc[m.region == "Aggregates", "iso3"].nunique())
issues.append({"issue": f"wb_indicators.csv mixes countries/economies with {n_agg_ent} regional/income/other aggregates (region == 'Aggregates').",
               "handling": "All aggregate entities excluded before any computation, per TASK rules."})
issues.append({"issue": "Some entity names carry trailing whitespace in wb_entities.csv (e.g. 'Latin America & Caribbean ', 'Sub-Saharan Africa ').",
               "handling": "Strings stripped; join is on iso3 so no effect on the analysis."})
issues.append({"issue": "ISO3 code 'NAM' (Namibia) would be read as missing by pandas' default NA parsing.",
               "handling": "CSV read with keep_default_na=False and only empty strings treated as missing."})
c = m[m.region != "Aggregates"].copy()

# ---- analysis year ----
c["complete"] = c.gdp_pc_ppp.notna() & c.life_expectancy.notna()
per_year = c.groupby("year")["complete"].sum()
eligible = per_year[per_year >= 150]
analysis_year = int(eligible.index.max())

y = c[c.year == analysis_year]
missing = y[~y.complete]
n_excluded_missing = int(len(missing))
issues.append({"issue": f"In {analysis_year}, {n_excluded_missing} countries/economies lack an indicator (GDP pc PPP missing for: "
               + ", ".join(missing.iso3) + "; LE missing for "
               + str(int(missing.life_expectancy.isna().sum())) + "). Missingness is non-random (small territories, conflict/sanctioned economies such as SSD, SYR, YEM, VEN, PRK, CUB, ERI).",
               "handling": "Listwise exclusion, no imputation, counted in n_excluded_missing; discussed as a selection threat."})

# implausible LE swings in the panel (outside analysis year)
swing = c.groupby("iso3").life_expectancy.agg(lambda s: s.max() - s.min())
big = swing[swing > 15].index.tolist()
issues.append({"issue": "Very large year-to-year swings in life expectancy for " + ", ".join(big)
               + " (e.g. CAF 18.8 years in 2022 vs 57.4 in 2023; SSD 35-40 in 2015-2017), reflecting crisis-mortality estimates/revisions.",
               "handling": f"Only {analysis_year} values enter the estimate; CAF {analysis_year} value kept (plausible); a sensitivity excluding CAF is reported in the report."})

d = y[y.complete].copy()
d["log_gdp"] = np.log(d.gdp_pc_ppp)
assert (d.gdp_pc_ppp > 0).all()
n = len(d)

# ---- baseline and model ----
le = d.life_expectancy.values
baseline_rmse = float(np.sqrt(np.mean((le - le.mean()) ** 2)))
X = sm.add_constant(d.log_gdp.values)
fit = sm.OLS(le, X).fit(cov_type="HC3")
alpha, beta = map(float, fit.params)
resid = fit.resid
model_rmse = float(np.sqrt(np.mean(resid ** 2)))
se = float(fit.bse[1])
ci = fit.conf_int(alpha=0.05)[1]  # normal-based 95% CI with HC3 SE
r = float(np.corrcoef(d.log_gdp, le)[0, 1])
r2 = float(fit.rsquared)

lo, hi = float(ci[0]), float(ci[1])
conclusion = "positive_effect" if lo > 0 else ("negative_effect" if hi < 0 else "no_detectable_effect")

# ---- sensitivity (printed, saved to out/sensitivity.json) ----
sens = {}
for k in ["nonrobust", "HC0", "HC1", "HC2", "HC3"]:
    f = sm.OLS(le, X).fit(cov_type=k)
    sens[f"se_{k}"] = float(f.bse[1])
noc = d[d.iso3 != "CAF"]
f = sm.OLS(noc.life_expectancy.values, sm.add_constant(noc.log_gdp.values)).fit(cov_type="HC3")
sens["beta_excl_CAF"] = float(f.params[1]); sens["se_excl_CAF_HC3"] = float(f.bse[1])
infl = sm.OLS(le, X).fit().get_influence()
cooks = infl.cooks_distance[0]
top = np.argsort(cooks)[::-1][:5]
sens["top5_cooks"] = {d.iso3.iloc[i]: float(cooks[i]) for i in top}
keep = cooks < 4 / n
f = sm.OLS(le[keep], X[keep]).fit(cov_type="HC3")
sens["beta_excl_cooks_gt_4_over_n"] = float(f.params[1]); sens["n_excl_cooks"] = int((~keep).sum())
Xq = np.column_stack([X, (d.log_gdp.values - d.log_gdp.mean()) ** 2])
fq = sm.OLS(le, Xq).fit(cov_type="HC3")
sens["quadratic_term"] = float(fq.params[2]); sens["quadratic_term_p_HC3"] = float(fq.pvalues[2])
sens["rho_spearman"] = float(d[["log_gdp", "life_expectancy"]].corr(method="spearman").iloc[0, 1])
# slope by year (all eligible years) for stability
sens["beta_by_year"] = {}
for yr in eligible.index:
    dy = c[(c.year == yr) & c.complete]
    fy = sm.OLS(dy.life_expectancy.values, sm.add_constant(np.log(dy.gdp_pc_ppp.values))).fit()
    sens["beta_by_year"][int(yr)] = float(fy.params[1])
sens["complete_per_year"] = {int(k): int(v) for k, v in per_year.items()}
sens["largest_neg_resid"] = {d.iso3.iloc[i]: float(resid[i]) for i in np.argsort(resid)[:5]}
sens["largest_pos_resid"] = {d.iso3.iloc[i]: float(resid[i]) for i in np.argsort(resid)[::-1][:5]}
with open(os.path.join(OUT, "sensitivity.json"), "w") as fh:
    json.dump(sens, fh, indent=2)

# ---- table ----
order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
tab = (d.groupby("income_level")
         .agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("log_gdp", "mean"),
              mean_life_expectancy=("life_expectancy", "mean"))
         .reindex([o for o in order if o in set(d.income_level)] +
                  sorted(set(d.income_level) - set(order)))
         .reset_index())
tab.round(4).to_csv(os.path.join(OUT, "table1.csv"), index=False)

# ---- figure ----
fig, ax = plt.subplots(figsize=(7, 5))
cols = {"Low income": "tab:red", "Lower middle income": "tab:orange",
        "Upper middle income": "tab:green", "High income": "tab:blue"}
for lvl, g in d.groupby("income_level"):
    ax.scatter(g.log_gdp, g.life_expectancy, s=16, alpha=0.75, label=lvl, color=cols.get(lvl, "grey"))
xs = np.linspace(d.log_gdp.min(), d.log_gdp.max(), 100)
ax.plot(xs, alpha + beta * xs, "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, {analysis_year} (n={n} countries/economies)")
ax.legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=150)

results = {
    "task_id": "F6-dev-01",
    "analysis_year": analysis_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": lo,
    "effect_ci_high": hi,
    "conclusion": conclusion,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": n_input_rows,
    "n_analysed": int(n),
    "random_seed": SEED,
}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(results, fh, indent=2)
print(json.dumps(results, indent=2)); print(json.dumps(sens, indent=2)); print(tab)