"""Replication entry point for F6-dev-01 (Preston curve, WDI).
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 = 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

OUT = "out"
os.makedirs(OUT, exist_ok=True)
TASK_ID = "F6-dev-01"
MIN_COUNTRIES = 150

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))}

assert not ind.duplicated(["iso3", "year"]).any()
assert not ent.duplicated("iso3").any()
assert set(ind.iso3) <= set(ent.iso3)

ex = []
def log_ex(step, rule, nb, na, why):
    ex.append(dict(step=step, rule=rule, n_before=nb, n_after=na, n_removed=nb - na, rationale=why))

# Step 1: remove aggregates (entity-year rows)
df = ind.merge(ent[["iso3", "region", "income_level"]], on="iso3", how="left")
n0 = len(df)
df = df[df.region != "Aggregates"].copy()
log_ex(1, "drop region == 'Aggregates'", n0, len(df), "analysis population is countries/economies only (TASK info rule)")

# Analysis year: latest with both indicators for >= 150 countries
complete = df.gdp_pc_ppp.notna() & df.life_expectancy.notna()
counts = df[complete].groupby("year").size()
eligible = counts[counts >= MIN_COUNTRIES]
if eligible.empty:
    raise SystemExit("No year with >=150 complete countries: premise invalid")
year = int(eligible.index.max())

# Step 2: keep analysis year
n1 = len(df)
d = df[df.year == year].copy()
log_ex(2, f"keep year == {year}", n1, len(d), "latest year 2015-2023 with both indicators for >=150 countries/economies")
n_year = len(d)

# Step 3: drop missing indicator
d = d[d.gdp_pc_ppp.notna() & d.life_expectancy.notna()].copy()
n_excl_missing = n_year - len(d)
log_ex(3, "drop rows with missing gdp_pc_ppp or life_expectancy", n_year, len(d), "no imputation (TASK info rule)")
missing_iso = sorted(set(df[df.year == year].iso3) - set(d.iso3))
pd.DataFrame(ex).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)

assert (d.gdp_pc_ppp > 0).all()
d["log_gdp"] = np.log(d.gdp_pc_ppp)
y = d.life_expectancy.to_numpy(); x = d.log_gdp.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 (primary); t(n-2) critical value for 95% CI
X = sm.add_constant(x)
fit = sm.OLS(y, X).fit(cov_type="HC3")
beta, alpha = float(fit.params[1]), float(fit.params[0])
se = float(fit.bse[1])
from scipy import stats
tcrit = float(stats.t.ppf(0.975, n - 2))
ci_low, ci_high = beta - tcrit * se, beta + tcrit * se
resid = fit.resid
model_rmse = float(np.sqrt(np.mean(resid ** 2)))
r = float(np.corrcoef(x, y)[0, 1])
r2 = float(fit.rsquared)

# Sensitivity: other HC variants and classical SE
se_alt = {f"sensitivity_se_{k.lower()}": float(sm.OLS(y, X).fit(cov_type=k).bse[1]) for k in ("HC0", "HC1", "HC2")}
se_alt["sensitivity_se_classical"] = float(sm.OLS(y, X).fit().bse[1])
# Sensitivity: drop the 5 largest |residual| observations (influence check, descriptive only)
order = np.argsort(-np.abs(resid))
keep = np.ones(n, bool); keep[order[:5]] = False
fit_trim = sm.OLS(y[keep], X[keep]).fit(cov_type="HC3")
top_resid = d.iloc[order[:5]].iso3.tolist()

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

# Table 1
lvl_order = ["Low income", "Lower middle income", "Upper middle income", "High income"]
t1 = (d.groupby("income_level")
        .agg(n_countries=("iso3", "size"), mean_log_gdp_pc=("log_gdp", "mean"), mean_life_expectancy=("life_expectancy", "mean"))
        .reset_index())
t1["_o"] = t1.income_level.map({k: i for i, k in enumerate(lvl_order)}).fillna(99)
t1 = t1.sort_values(["_o", "income_level"]).drop(columns="_o")
t1.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.4f")

# Figure 1
fig, ax = plt.subplots(figsize=(7, 5))
colors = {"Low income": "tab:red", "Lower middle income": "tab:orange", "Upper middle income": "tab:green", "High income": "tab:blue"}
for lvl in lvl_order + sorted(set(d.income_level) - set(lvl_order)):
    s = d[d.income_level == lvl]
    if len(s):
        ax.scatter(s.log_gdp, s.life_expectancy, s=14, alpha=0.75, color=colors.get(lvl, "gray"), 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 int'l $)")
ax.set_ylabel("Life expectancy at birth (years)")
ax.set_title(f"Preston curve, {year} ({n} countries/economies, WDI)")
ax.legend(fontsize=7, loc="lower right")
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=150)
plt.close(fig)

caf = d[d.iso3 == "CAF"]
data_issues = [
    {"issue": "wb_indicators.csv contains regional/income-group aggregates (79 entities with region 'Aggregates')",
     "handling": "Excluded via wb_entities.csv region == 'Aggregates' before choosing the analysis year."},
    {"issue": f"{n_excl_missing} countries/economies lack GDP per capita PPP in {year} (" + ", ".join(missing_iso) + "); missingness is non-random (small territories, conflict/sanctioned states)",
     "handling": "Excluded without imputation as required; counted in n_excluded_missing; possible selection bias discussed in report."},
    {"issue": "Central African Republic life expectancy is erratic across 2015-2022 (e.g. 18.8 in 2022, 31.5 in 2019)",
     "handling": f"Only the analysis-year value is used ({float(caf.life_expectancy.iloc[0]) if len(caf) else 'n/a'} in {year}), which is plausible; kept. Influence assessed by the drop-5-largest-residuals sensitivity."},
    {"issue": "West Bank and Gaza life expectancy falls ~11.5 years in 2023 (conflict)",
     "handling": "Kept as reported; real data feature, covered by the residual-trimming sensitivity if influential."},
    {"issue": "Recent-year WDI values may be modelled estimates/projections rather than measurements",
     "handling": "Not verifiable from the provided files; noted as a limitation."},
]

res = {
    "task_id": TASK_ID,
    "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_countries_in_year": int(n_year),
    "n_excluded_missing": int(n_excl_missing),
    "complete_countries_by_year": {str(k): int(v) for k, v in counts.items()},
    "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": float(ci_low),
    "effect_ci_high": float(ci_high),
    "ci_critical_value_t": tcrit,
    "effect_t_stat": beta / se,
    "conclusion": conclusion,
    **se_alt,
    "sensitivity_trim5_effect": float(fit_trim.params[1]),
    "sensitivity_trim5_se_hc3": float(fit_trim.bse[1]),
    "sensitivity_trim5_dropped_iso3": top_resid,
    "mean_life_expectancy": float(y.mean()),
}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=2)
print(f"year={year} n={n} excluded_missing={n_excl_missing} beta={beta:.4f} se_HC3={se:.4f} "
      f"CI=[{ci_low:.4f},{ci_high:.4f}] R2={r2:.4f} rmse_base={baseline_rmse:.4f} rmse_ols={model_rmse:.4f} -> {conclusion}")