"""Replication entry point for F6-dev-01 (Preston curve).
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"
os.makedirs(OUT, exist_ok=True)
read = lambda f: pd.read_csv(os.path.join(DATA, f), keep_default_na=False, na_values=[""])
ind, ent = read("wb_indicators.csv"), read("wb_entities.csv")
n_input_rows = {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))}
issues, excl = [], []

def log_excl(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() and not ent.iso3.duplicated().any()
assert set(ind.iso3) <= set(ent.iso3)
assert (ind.gdp_pc_ppp.dropna() > 0).all()

# Step 1: drop aggregates
m = ind.merge(ent[["iso3", "region", "income_level"]], on="iso3", how="left")
n0 = len(m)
n_agg_ent = int(m.loc[m.region == "Aggregates", "iso3"].nunique())
c = m[m.region != "Aggregates"].copy()
log_excl(1, "drop entity-years whose region == 'Aggregates' in wb_entities.csv", n0, len(c),
         "information rule: countries/economies only")
issues.append({"issue": f"wb_indicators.csv contains {n_agg_ent} regional/income/other aggregate entities "
               f"({n0 - len(c)} entity-year rows)", "handling": "excluded via region == 'Aggregates' (rule)"})

# Step 2: analysis year selection
cov = c.assign(both=c.gdp_pc_ppp.notna() & c.life_expectancy.notna()).groupby("year").both.sum()
elig = cov[cov >= 150]
analysis_year = int(elig.index.max())
n1 = len(c)
y = c[c.year == analysis_year].copy()
log_excl(2, f"keep analysis year {analysis_year} (latest year with >=150 complete countries)", n1, len(y),
         "information rule: single cross-section in analysis year")

# Step 3: drop missing
miss = y[y.gdp_pc_ppp.isna() | y.life_expectancy.isna()]
n_excluded_missing = int(len(miss))
n2 = len(y)
d = y.dropna(subset=["gdp_pc_ppp", "life_expectancy"]).copy()
log_excl(3, "drop countries with gdp_pc_ppp or life_expectancy missing in analysis year", n2, len(d),
         "information rule: no imputation")
issues.append({"issue": f"{n_excluded_missing} countries/economies lack an indicator in {analysis_year} "
               f"(gdp missing: {int(miss.gdp_pc_ppp.isna().sum())}, LE missing: {int(miss.life_expectancy.isna().sum())}): "
               + ", ".join(sorted(miss.iso3)),
               "handling": "excluded, counted in n_excluded_missing; possible selection bias discussed in report"})
issues.append({"issue": "Central African Republic life expectancy series is implausibly volatile in 2015-2022 "
               "(e.g. 18.8 years in 2022, 31.5 in 2019); 2023 value 57.4 is plausible",
               "handling": "affects non-analysis years only; 2023 value kept; noted as data-quality caveat"})
issues.append({"issue": "West Bank and Gaza life expectancy drops from 76.7 (2022) to 65.2 (2023), consistent with "
               "conflict mortality but a large outlier relative to its income",
               "handling": "kept (genuine reported value, no rule permits dropping it); a leave-5-most-influential-out (Cook distance) sensitivity is reported"})

# Estimation
d["ln_gdp"] = np.log(d.gdp_pc_ppp)
X = sm.add_constant(d.ln_gdp.values)
yv = d.life_expectancy.values
fit = sm.OLS(yv, X).fit(cov_type="HC3")
alpha, beta = float(fit.params[0]), float(fit.params[1])
se = float(fit.bse[1])
z = 1.959963984540054
ci_lo, ci_hi = beta - z * se, beta + z * se
resid = yv - (alpha + beta * d.ln_gdp.values)
baseline_rmse = float(np.sqrt(np.mean((yv - yv.mean()) ** 2)))
model_rmse = float(np.sqrt(np.mean(resid ** 2)))
r = float(np.corrcoef(d.ln_gdp, yv)[0, 1])
conclusion = "positive_effect" if ci_lo > 0 else ("negative_effect" if ci_hi < 0 else "no_detectable_effect")

# Supplementary (not estimand): other robust SEs, sensitivity dropping the 5 most influential obs (Cook's D)
fit_ols = sm.OLS(yv, X).fit()
se_alt = {k: float(sm.OLS(yv, X).fit(cov_type=k).bse[1]) for k in ["HC0", "HC1", "HC2", "HC3"]}
cooks = fit_ols.get_influence().cooks_distance[0]
top = np.argsort(cooks)[::-1][:5]
keep = np.setdiff1d(np.arange(len(d)), top)
fit_s = sm.OLS(yv[keep], X[keep]).fit(cov_type="HC3")
d["resid"] = resid

results = {
    "task_id": "F6-dev-01",
    "feasibility": "complete",
    "feasibility_reason": None,
    "analysis_year": analysis_year,
    "n_excluded_missing": n_excluded_missing,
    "n_analysed": int(len(d)),
    "baseline_metric": baseline_rmse,
    "model_metric": model_rmse,
    "pearson_r": r,
    "r_squared": float(fit.rsquared),
    "effect_estimate": beta,
    "intercept": alpha,
    "effect_se": se,
    "effect_ci_low": float(ci_lo),
    "effect_ci_high": float(ci_hi),
    "conclusion": conclusion,
    "data_issues": issues,
    "n_input_rows": n_input_rows,
    "random_seed": None,
    "se_type": "HC3",
    "complete_countries_by_year": {str(k): int(v) for k, v in cov.items()},
    "n_countries_in_analysis_year": int(n2),
    "supp_effect_se_by_type": se_alt,
    "supp_effect_se_classical": float(fit_ols.bse[1]),
    "supp_mean_life_expectancy": float(yv.mean()),
    "supp_sens_drop5_cooks_iso3": [str(s) for s in d.iso3.values[top]],
    "supp_sens_drop5_effect_estimate": float(fit_s.params[1]),
    "supp_sens_drop5_effect_se": float(fit_s.bse[1]),
    "supp_sens_drop5_r_squared": float(fit_s.rsquared),
    "supp_largest_neg_residual_iso3": str(d.iso3.values[np.argmin(resid)]),
    "supp_largest_neg_residual": float(resid.min()),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

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

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=("ln_gdp", "mean"),
                                   mean_life_expectancy=("life_expectancy", "mean")).reset_index())
t["o"] = t.income_level.map({k: i for i, k in enumerate(order)}).fillna(99)
t = t.sort_values(["o", "income_level"]).drop(columns="o")
t.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.4f")

fig, ax = plt.subplots(figsize=(7, 5))
cols = dict(zip(order, ["#d7191c", "#fdae61", "#abd9e9", "#2c7bb6"]))
for lvl in order:
    s = d[d.income_level == lvl]
    ax.scatter(s.ln_gdp, s.life_expectancy, s=18, c=cols[lvl], label=f"{lvl} (n={len(s)})", edgecolor="k", lw=0.3)
xs = np.linspace(d.ln_gdp.min(), d.ln_gdp.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 intl $)")
ax.set_ylabel("Life expectancy at birth (years)")
ax.set_title(f"Preston curve, {analysis_year} (n={len(d)} countries/economies)")
ax.legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120, metadata={"Software": None})
plt.close(fig)

print(f"year={analysis_year} n={len(d)} excluded_missing={n_excluded_missing} beta={beta:.4f} "
      f"HC3 se={se:.4f} CI=[{ci_lo:.4f},{ci_hi:.4f}] R2={fit.rsquared:.4f} conclusion={conclusion}")