"""Replication entry point for F6-dev-01 (Preston curve). Run from workspace root."""
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 = None  # no randomness used
os.makedirs("out", exist_ok=True)
# keep_default_na=False so that e.g. Namibia's ISO2-like strings / "NA" are never read as missing
rd = dict(keep_default_na=False, na_values=[""])
ind = pd.read_csv("data/wb_indicators.csv", **rd)
ent = pd.read_csv("data/wb_entities.csv", **rd)
issues = []

dup = int(ind.duplicated(["iso3", "year"]).sum())
issues.append({"issue": f"wb_indicators.csv mixes countries with regional/income aggregates ({ent.region.eq('Aggregates').sum()} aggregate codes in metadata)",
               "handling": "Dropped every iso3 whose region is 'Aggregates' in wb_entities.csv"})
unmatched = sorted(set(ind.iso3) - set(ent.iso3))
issues.append({"issue": f"{len(unmatched)} iso3 codes in indicators not found in entity metadata; {dup} duplicate iso3-year rows",
               "handling": "Checked; nothing to handle" if not unmatched and dup == 0 else "Unmatched codes excluded"})
issues.append({"issue": "Literal country codes/names such as 'NA' could be parsed as missing by default CSV readers",
               "handling": "Read CSVs with only empty strings treated as missing"})

countries = set(ent.loc[ent.region != "Aggregates", "iso3"])
c = ind[ind.iso3.isin(countries)].copy()
both = c.gdp_pc_ppp.notna() & c.life_expectancy.notna()
counts = both.groupby(c.year).sum()
year = int(max(y for y, n in counts.items() if n >= 150))

caf = c[(c.iso3 == "CAF") & (c.life_expectancy < 35)]
if len(caf):
    issues.append({"issue": "Implausible life-expectancy values for Central African Republic in non-analysis years ("
                   + ", ".join(f"{int(r.year)}: {r.life_expectancy:.1f}" for r in caf.itertuples()) + ")",
                   "handling": f"Not in analysis year {year}; left unchanged, flagged only"})

y = c[c.year == year].merge(ent[["iso3", "income_level", "region"]], on="iso3")
miss = y[y.gdp_pc_ppp.isna() | y.life_expectancy.isna()]
issues.append({"issue": f"{len(miss)} countries/economies in {year} lack an indicator (all GDP pc PPP missing: "
               + ", ".join(miss.iso3) + "); missingness concentrated in small territories and conflict/sanctioned economies",
               "handling": "Excluded listwise, no imputation; counted in n_excluded_missing"})
a = y.drop(miss.index).copy()
nonpos = int((a.gdp_pc_ppp <= 0).sum())
issues.append({"issue": f"{nonpos} non-positive GDP values (log undefined)", "handling": "Checked; none found" if nonpos == 0 else "Excluded"})
a = a[a.gdp_pc_ppp > 0]
a["lngdp"] = np.log(a.gdp_pc_ppp)

X = sm.add_constant(a.lngdp.values)
fit = sm.OLS(a.life_expectancy.values, X).fit(cov_type="HC3")
fit0 = sm.OLS(a.life_expectancy.values, X).fit()
le = a.life_expectancy.values
base = float(np.sqrt(np.mean((le - le.mean()) ** 2)))
rmse = float(np.sqrt(np.mean(fit.resid ** 2)))
r = float(np.corrcoef(a.lngdp, le)[0, 1])
beta, se = float(fit.params[1]), float(fit.bse[1])
lo, hi = [float(v) for v in fit.conf_int(0.05)[1]]
concl = "positive_effect" if lo > 0 else "negative_effect" if hi < 0 else "no_detectable_effect"

# sensitivity (reported in report, not in required keys)
sens = {}
for hc in ["HC0", "HC1", "HC2", "HC3"]:
    f = sm.OLS(le, X).fit(cov_type=hc); sens[hc] = float(f.bse[1])
sens["classical_se"] = float(fit0.bse[1])
yrs = {}
for yy in sorted(c.year.unique()):
    s = c[(c.year == yy) & c.gdp_pc_ppp.notna() & c.life_expectancy.notna()]
    f = sm.OLS(s.life_expectancy.values, sm.add_constant(np.log(s.gdp_pc_ppp.values))).fit(cov_type="HC3")
    yrs[int(yy)] = {"n": len(s), "beta": float(f.params[1]), "se": float(f.bse[1])}
infl = fit0.get_influence(); cd = infl.cooks_distance[0]
top = a.assign(cook=cd, resid=fit0.resid).sort_values("cook", ascending=False).head(5)
trim = a[cd <= 4 / len(a)]
ft = sm.OLS(trim.life_expectancy.values, sm.add_constant(trim.lngdp.values)).fit(cov_type="HC3")
quad = sm.OLS(le, np.column_stack([np.ones(len(a)), a.lngdp, a.lngdp ** 2])).fit(cov_type="HC3")
sensitivity = {"se_by_estimator": sens, "slope_by_year": yrs,
               "top5_cooks": top[["iso3", "name", "gdp_pc_ppp", "life_expectancy", "resid", "cook"]].to_dict("records"),
               "trimmed_cooks_4_over_n": {"n": len(trim), "beta": float(ft.params[1]), "se": float(ft.bse[1])},
               "quadratic_in_log": {"coef": [float(v) for v in quad.params], "se": [float(v) for v in quad.bse]}}
with open("out/sensitivity.json", "w") as fh:
    json.dump(sensitivity, fh, indent=2)

tab = (a.groupby("income_level").agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("lngdp", "mean"),
                                     mean_life_expectancy=("life_expectancy", "mean")).reset_index())
order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
tab["k"] = tab.income_level.map({k: i for i, k in enumerate(order)}).fillna(99)
tab.sort_values("k").drop(columns="k").round(4).to_csv("out/table1.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
for lvl in order:
    s = a[a.income_level == lvl]
    ax.scatter(s.lngdp, s.life_expectancy, s=14, alpha=.75, label=f"{lvl} (n={len(s)})")
xs = np.linspace(a.lngdp.min(), a.lngdp.max(), 100)
ax.plot(xs, fit.params[0] + beta * xs, "k-", lw=2, label=f"OLS: LE = {fit.params[0]:.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={len(a)} countries/economies)")
ax.legend(fontsize=8); fig.tight_layout(); fig.savefig("out/figure1.png", dpi=150)

res = {"task_id": "F6-dev-01", "analysis_year": year, "n_excluded_missing": int(len(miss)),
       "baseline_metric": base, "model_metric": rmse, "pearson_r": r, "r_squared": float(fit.rsquared),
       "effect_estimate": beta, "intercept": float(fit.params[0]), "effect_se": se,
       "effect_ci_low": lo, "effect_ci_high": hi, "conclusion": concl,
       "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,
       "n_input_rows": {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))},
       "n_analysed": int(len(a)), "random_seed": SEED}
with open("out/results.json", "w") as fh:
    json.dump(res, fh, indent=2)
print(json.dumps(res, indent=2)); print(json.dumps(sensitivity, indent=1)); print(tab)