"""F6-dev-01: Preston curve, replication entry point.
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, pandas as pd
import statsmodels.api as sm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "out"
os.makedirs(OUT, exist_ok=True)
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)
n_input_rows = {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))}

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

assert not ind.duplicated(["iso3", "year"]).any()
unmatched = sorted(set(ind.iso3) - set(ent.iso3))
assert not unmatched, unmatched
df = ind.merge(ent[["iso3", "region", "income_level"]], on="iso3", how="left")

n0 = len(df)
df = df[df.region != "Aggregates"].copy()
log(1, "drop entities with region == 'Aggregates'", n0, len(df),
    "analysis population is countries/economies only (TASK rule)")

df["complete"] = df.gdp_pc_ppp.notna() & df.life_expectancy.notna()
counts = df.groupby("year").complete.sum()
eligible = counts[counts >= 150]
if eligible.empty:
    raise SystemExit("no year with >=150 complete countries")
year = int(eligible.index.max())

n1 = len(df)
d = df[df.year == year].copy()
log(2, f"keep analysis year {year} (most recent year with >=150 complete countries)", n1, len(d),
    "estimand is a single cross-section in the analysis year")
n2 = len(d)
d = d[d.complete].copy()
n_excl_missing = n2 - len(d)
log(3, "drop countries with missing gdp_pc_ppp or life_expectancy", n2, len(d),
    "no imputation (TASK rule)")
assert (d.gdp_pc_ppp > 0).all()
pd.DataFrame(excl).to_csv(f"{OUT}/exclusions.csv", index=False)

d["log_gdp_pc"] = np.log(d.gdp_pc_ppp)
y, x = d.life_expectancy.to_numpy(), d.log_gdp_pc.to_numpy()
n = len(d)

# Baseline: mean-only prediction
baseline_rmse = float(np.sqrt(np.mean((y - y.mean()) ** 2)))

# OLS with HC3 robust SE
X = sm.add_constant(x)
fit = sm.OLS(y, X).fit(cov_type="HC3")
alpha, beta = map(float, fit.params)
se = float(fit.bse[1])
ci_lo, ci_hi = map(float, fit.conf_int(alpha=0.05)[1])  # normal quantile with HC3
model_rmse = float(np.sqrt(np.mean(fit.resid ** 2)))
r = float(np.corrcoef(x, y)[0, 1])
r2 = float(fit.rsquared)

if ci_lo > 0: concl = "positive_effect"
elif ci_hi < 0: concl = "negative_effect"
else: concl = "no_detectable_effect"

missing_list = sorted(df[(df.year == year) & ~df.complete].iso3)
caf = df[(df.iso3 == "CAF")].set_index("year").life_expectancy
data_issues = [
    {"issue": f"wb_indicators.csv mixes countries with {int((ind.merge(ent, on='iso3').region == 'Aggregates').sum() // 9)} aggregate entities (regions/income groups)",
     "handling": "Excluded every entity whose region is 'Aggregates' in wb_entities.csv"},
    {"issue": f"{n_excl_missing} countries/economies lack gdp_pc_ppp in {year}: {', '.join(missing_list)} (small territories plus conflict/sanctioned economies; missingness likely not random)",
     "handling": "Excluded and counted in n_excluded_missing; no imputation"},
    {"issue": "Central African Republic life-expectancy series is implausibly volatile in earlier years (e.g. "
              f"{caf[2019]:.1f} in 2019, {caf[2022]:.1f} in 2022, {caf[2023]:.1f} in 2023)",
     "handling": f"Its {year} value is plausible and kept (no task rule allows dropping it); flagged as a data-quality concern"},
    {"issue": "Life expectancy is reported for every entity-year including tiny territories, indicating modelled (UN WPP) estimates; latest-year values may be provisional",
     "handling": "Used as published; noted as a limitation"},
]

res = {
    "task_id": "F6-dev-01",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": int(n),
    "random_seed": None,
    "analysis_year": year,
    "n_excluded_missing": int(n_excl_missing),
    "baseline_metric": baseline_rmse,
    "model_metric": model_rmse,
    "pearson_r": r,
    "r_squared": r2,
    "effect_estimate": beta,
    "intercept": alpha,
    "effect_se": se,
    "effect_se_type": "HC3",
    "effect_ci_low": ci_lo,
    "effect_ci_high": ci_hi,
    "conclusion": concl,
    "n_countries_total": int(df.iso3.nunique()),
    "complete_counts_by_year": {str(k): int(v) for k, v in counts.items()},
    "effect_se_classical": float(sm.OLS(y, X).fit().bse[1]),
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)

# Table 1
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["_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.round(4).to_csv(f"{OUT}/table1.csv", index=False)

# Figure 1
fig, ax = plt.subplots(figsize=(7, 5))
colors = dict(zip(order, ["tab:red", "tab:orange", "tab:green", "tab:blue"]))
for lvl, g in d.groupby("income_level"):
    ax.scatter(g.log_gdp_pc, g.life_expectancy, s=14, alpha=0.75, label=lvl, color=colors.get(lvl, "grey"))
xs = np.linspace(x.min(), x.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. $)")
ax.set_ylabel("Life expectancy at birth (years)")
ax.set_title(f"Preston curve, {year} (n = {n} countries/economies)")
ax.legend(fontsize=8, loc="lower right")
fig.tight_layout()
fig.savefig(f"{OUT}/figure1.png", dpi=150)

print(f"year={year} n={n} excluded_missing={n_excl_missing} beta={beta:.3f} se(HC3)={se:.3f} "
      f"CI=[{ci_lo:.3f},{ci_hi:.3f}] R2={r2:.3f} RMSE base={baseline_rmse:.3f} model={model_rmse:.3f} -> {concl}")