"""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 0 for completeness; nothing depends on it)."""
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
np.random.seed(SEED)
TASK_ID = "F6-dev-01"
OUT = "out"
os.makedirs(OUT, exist_ok=True)

# keep_default_na=False so codes such as "NA" are never parsed as missing; only empty cells are NaN
ind = pd.read_csv("data/wb_indicators.csv", keep_default_na=False, na_values=[""],
                  dtype={"iso3": str, "name": str})
ent = pd.read_csv("data/wb_entities.csv", keep_default_na=False, na_values=[""], dtype=str)
n_input_rows = {"wb_indicators.csv": int(len(ind)), "wb_entities.csv": int(len(ent))}
data_issues = []
excl = []

assert not ind.duplicated(["iso3", "year"]).any()
assert set(ind.iso3) <= set(ent.iso3), "indicator codes missing from metadata"

# Step 1: drop aggregates
agg = set(ent.loc[ent.region == "Aggregates", "iso3"])
n0 = len(ind)
c = ind[~ind.iso3.isin(agg)].copy()
excl.append(("1", "drop entities with region == 'Aggregates'", n0, len(c),
             "Analysis population is countries/economies only (TASK rule)"))
data_issues.append({"issue": f"wb_indicators.csv mixes {ind.iso3.isin(agg).groupby(ind.iso3).first().sum()} aggregate entities ({n0-len(c)} rows) with countries/economies",
                    "handling": "Excluded all entities with region == 'Aggregates' in wb_entities.csv"})

# Step 2: analysis year
both = c.gdp_pc_ppp.notna() & c.life_expectancy.notna()
cover = both.groupby(c.year).sum()
eligible = cover[(cover >= 150) & (cover.index >= 2015) & (cover.index <= 2023)]
if eligible.empty:
    raise SystemExit("premise_invalid: no year with >=150 complete countries")
year = int(eligible.index.max())
n1 = len(c)
y = c[c.year == year].copy()
excl.append(("2", f"keep analysis year {year} only", n1, len(y),
             "Latest 2015-2023 year with both indicators for >=150 countries/economies"))

# Step 3: missing indicators
miss = y.gdp_pc_ppp.isna() | y.life_expectancy.isna()
n_excluded_missing = int(miss.sum())
missing_names = sorted(y.loc[miss, "name"])
n2 = len(y)
d = y[~miss].copy()
excl.append(("3", "drop countries with missing GDPpc PPP or life expectancy in analysis year", n2, len(d),
             "No imputation (TASK rule); dropped: " + "; ".join(missing_names)))
data_issues.append({"issue": f"{n_excluded_missing} countries/economies lack GDP per capita PPP (or LE) in {year}: " + ", ".join(missing_names),
                    "handling": "Excluded from the estimate, counted in n_excluded_missing; no imputation (possible selection: many are small territories or conflict/sanctioned economies)"})

# anomaly scan (outside analysis year: documented only)
cs = c.sort_values(["iso3", "year"]).copy()
cs["dle"] = cs.groupby("iso3").life_expectancy.diff()
jumps = cs.loc[cs.dle.abs() > 10, ["iso3", "year", "life_expectancy"]]
if len(jumps):
    data_issues.append({"issue": "Implausible year-to-year life-expectancy swings (>10 years) in: " +
                        ", ".join(f"{r.iso3} {r.year} ({r.life_expectancy:.1f})" for r in jumps.itertuples()),
                        "handling": f"Only the {year} cross-section is used; values in {year} for these entities were inspected and kept as reported (no editing of source data)"})
assert (d.gdp_pc_ppp > 0).all()

pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "rationale"]).assign(
    n_removed=lambda t: t.n_before - t.n_after)[
    ["step", "rule", "n_before", "n_after", "n_removed", "rationale"]].to_csv(f"{OUT}/exclusions.csv", index=False)

# Baseline and model
d["lngdp"] = np.log(d.gdp_pc_ppp)
le = d.life_expectancy.to_numpy(); x = d.lngdp.to_numpy(); n = len(d)
baseline_rmse = float(np.sqrt(np.mean((le - le.mean()) ** 2)))
X = sm.add_constant(x)
fit = sm.OLS(le, X).fit(cov_type="HC3")
alpha, beta = map(float, fit.params)
resid = le - fit.fittedvalues
model_rmse = float(np.sqrt(np.mean(resid ** 2)))
se = float(fit.bse[1])
ci_low, ci_high = map(float, fit.conf_int(alpha=0.05)[1])  # t(n-2) critical value
r = float(np.corrcoef(x, le)[0, 1])
r2 = float(fit.rsquared)
# sensitivity: other HC variants and classical SE (reported, not used for conclusion)
sens = {f"sensitivity_se_{k}": float(sm.OLS(le, X).fit(cov_type=k).bse[1]) for k in ["HC0", "HC1", "HC2"]}
sens["sensitivity_se_classical"] = float(sm.OLS(le, X).fit().bse[1])
# sensitivity: drop entities with erratic LE series (flagged above) present in the analysis sample
flag = set(jumps.iso3) & set(d.iso3)
keep = ~d.iso3.isin(flag).to_numpy()
fs = sm.OLS(le[keep], X[keep]).fit(cov_type="HC3")
sens["sensitivity_beta_drop_erratic_LE"] = float(fs.params[1])
sens["sensitivity_n_drop_erratic_LE"] = int(keep.sum())
# sensitivity: max absolute change in beta from leaving out any single country
loo = [float(sm.OLS(np.delete(le, i), np.delete(X, i, axis=0)).fit().params[1]) for i in range(len(le))]
sens["sensitivity_beta_loo_min"] = float(min(loo)); sens["sensitivity_beta_loo_max"] = float(max(loo))

if ci_low > 0: concl = "positive_effect"
elif ci_high < 0: concl = "negative_effect"
else: concl = "no_detectable_effect"

# Table 1
order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
d = d.merge(ent[["iso3", "income_level"]], on="iso3", how="left")
t1 = d.groupby("income_level").agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("lngdp", "mean"),
                                   mean_life_expectancy=("life_expectancy", "mean")).reset_index()
t1["_o"] = t1.income_level.map({k: i for i, k in enumerate(order)}).fillna(99)
t1 = t1.sort_values(["_o", "income_level"]).drop(columns="_o")
t1.round(4).to_csv(f"{OUT}/table1.csv", index=False)

# Figure 1
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 in order:
    s = d[d.income_level == lvl]
    ax.scatter(s.lngdp, s.life_expectancy, s=14, alpha=0.75, color=cols[lvl], label=f"{lvl} (n={len(s)})")
xx = np.linspace(x.min(), x.max(), 100)
ax.plot(xx, alpha + beta * xx, "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 intl $)")
ax.set_ylabel("Life expectancy at birth (years)")
ax.set_title(f"Preston curve, {year} (n={n} countries/economies)")
ax.legend(fontsize=7, loc="lower right")
fig.tight_layout(); fig.savefig(f"{OUT}/figure1.png", dpi=150); plt.close(fig)

res = {
    "task_id": TASK_ID, "feasibility": "complete", "feasibility_reason": None,
    "analysis_year": 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_se_type": "HC3", "effect_ci_low": ci_low, "effect_ci_high": ci_high,
    "conclusion": concl, "data_issues": data_issues, "n_input_rows": n_input_rows,
    "n_analysed": int(n), "random_seed": SEED,
    "n_countries_in_population": int(c.iso3.nunique()),
    "complete_countries_by_year": {str(k): int(v) for k, v in cover.items()},
    **sens,
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(f"year={year} n={n} excluded={n_excluded_missing} beta={beta:.4f} se(HC3)={se:.4f} "
      f"CI=[{ci_low:.4f},{ci_high:.4f}] r={r:.4f} R2={r2:.4f} rmse0={baseline_rmse:.3f} rmse={model_rmse:.3f} -> {concl}")