"""F1-conf-02: balanced-panel mean life expectancy 2000-2022.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/analytic.csv, out/exclusions.csv. No randomness."""
import json
from pathlib import Path
import pandas as pd

DATA = Path("data/life_expectancy_panel.csv")
OUT = Path("out")
YEARS = list(range(2000, 2023))

raw = pd.read_csv(DATA, dtype={"country_name": str, "iso3": str})
n_raw = len(raw)
ledger, issues = [], []

def step(name, rule, before, after, why):
    ledger.append(dict(step=name, rule=rule, n_before=before, n_after=after,
                       n_removed=before - after, rationale=why))

df = raw[raw["year"].between(2000, 2022)].copy()
step("1_year_window", "keep rows with 2000<=year<=2022", n_raw, len(df),
     "Study window; no rows outside it in this file.")

# Duplicate (iso3, year) rows
dupmask = df.duplicated(["iso3", "year"], keep=False)
if dupmask.any():
    dups = df[dupmask]
    nval = dups.groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False)
    if (nval > 1).any():
        raise SystemExit("Conflicting duplicate values; cannot resolve without outside information.")
    for iso, g in dups.groupby("iso3"):
        names = sorted(g["country_name"].unique())
        issues.append({"issue": f"{iso}: {g[['iso3','year']].drop_duplicates().shape[0]} country-years appear twice "
                                f"under names {names} with identical values (assembly duplicate).",
                       "handling": "Kept one row per (iso3, year); values identical so no number depends on the choice; "
                                   "label 'Turkiye' (current World Bank name) retained where applicable."})
    # prefer 'Turkiye' label deterministically, otherwise the last occurrence
    df["_pref"] = (df["country_name"] == "Turkiye").astype(int)
    df = df.sort_values(["iso3", "year", "_pref"]).drop_duplicates(["iso3", "year"], keep="last").drop(columns="_pref")
before = len(raw[raw["year"].between(2000, 2022)])
step("2_dedupe_iso3_year", "one row per (iso3, year); drop exact duplicates", before, len(df),
     "Each country contributes at most one observation per year (TUR listed as both Turkey and Turkiye).")

# Missing values
miss = df[df["life_expectancy"].isna()]
for _, r in miss.iterrows():
    issues.append({"issue": f"{r.iso3} {r.year}: empty life_expectancy cell.",
                   "handling": "Treated as missing (no imputation); country excluded from balanced set."})
before = len(df)
df = df[df["life_expectancy"].notna()]
step("3_drop_missing_value", "drop country-years with empty life_expectancy", before, len(df),
     "No imputation allowed; missing country-year has no value.")

# Coverage
all_iso = sorted(raw["iso3"].unique())
cov = df.groupby("iso3")["year"].nunique().reindex(all_iso, fill_value=0)
balanced = sorted(cov[cov == len(YEARS)].index)
incomplete = sorted(set(all_iso) - set(balanced))
for iso in incomplete:
    yrs = sorted(df.loc[df.iso3 == iso, "year"])
    absent = sorted(set(YEARS) - set(raw.loc[raw.iso3 == iso, "year"]))
    if absent:
        issues.append({"issue": f"{iso}: no rows for years {absent[0]}-{absent[-1]} ({len(absent)} years absent; "
                                f"recorded {yrs[0]}-{yrs[-1]}).",
                       "handling": "Not imputed or looked up; country excluded from balanced set (still in 2022 baseline)."})
before = len(df)
analytic = df[df["iso3"].isin(balanced)].sort_values(["iso3", "year"])[["iso3", "country_name", "year", "life_expectancy"]]
step("4_balanced_panel", "keep countries with a recorded value in all 23 years 2000-2022", before, len(analytic),
     f"Balanced-panel estimand; excluded incomplete countries: {incomplete}.")

assert len(analytic) == len(balanced) * len(YEARS)
assert not analytic.duplicated(["iso3", "year"]).any()

m2000 = analytic.loc[analytic.year == 2000, "life_expectancy"].mean()
m2022 = analytic.loc[analytic.year == 2022, "life_expectancy"].mean()
avail22 = df[df.year == 2022]  # deduped, non-missing: one row per iso3
assert avail22["iso3"].is_unique
m2022_all = avail22["life_expectancy"].mean()

results = {
    "task_id": "F1-conf-02",
    "n_countries_total": len(all_iso),
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": len(incomplete),
    "balanced_countries": balanced,
    "incomplete_countries": incomplete,
    "n_analysed": int(len(analytic)),
    "mean_2000_balanced": round(float(m2000), 6),
    "mean_2022_balanced": round(float(m2022), 6),
    "change_2000_2022_balanced": round(float(m2022 - m2000), 6),
    "mean_2022_all_available": round(float(m2022_all), 6),
    "n_countries_2022_available": int(avail22["iso3"].nunique()),
    "baseline_minus_balanced_2022": round(float(m2022_all - m2022), 6),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
}
OUT.mkdir(exist_ok=True)
analytic.to_csv(OUT / "analytic.csv", index=False)
pd.DataFrame(ledger).to_csv(OUT / "exclusions.csv", index=False)
with open(OUT / "results.json", "w") as f:
    json.dump(results, f, indent=2, ensure_ascii=False)
print(json.dumps({k: v for k, v in results.items() if k != "data_issues"}, indent=1))