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

DATA = "data/life_expectancy_panel.csv"
OUT = "out"
YEARS = list(range(2000, 2023))
os.makedirs(OUT, exist_ok=True)

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

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

# --- Step 1: restrict to 2000-2022 (scope check) ---
d0 = raw
d1 = d0[d0["year"].isin(YEARS)].copy()
step("1_year_scope", "keep rows with year in 2000..2022", d0, d1, "Task scope; all rows already in range")

# --- Step 2: resolve duplicate (iso3, year) rows ---
dup_mask = d1.duplicated(["iso3", "year"], keep=False)
dups = d1[dup_mask]
conflicts = []
for (iso, yr), g in dups.groupby(["iso3", "year"]):
    vals = g["life_expectancy"].dropna().unique()
    if len(vals) > 1:
        conflicts.append((iso, yr, list(vals)))
if len(conflicts):
    raise SystemExit(f"Conflicting duplicate values, need a decision: {conflicts}")
names_by_iso = d1.groupby("iso3")["country_name"].unique()
multi_name = {k: sorted(v) for k, v in names_by_iso.items() if len(v) > 1}
# canonical name: current World Bank spelling for TUR; otherwise the unique name
canonical = {"TUR": "Turkiye"}
d1["country_name"] = [canonical.get(i, n) if i in multi_name else n
                      for i, n in zip(d1["iso3"], d1["country_name"])]
d2 = d1.sort_values(["iso3", "year"]).drop_duplicates(["iso3", "year"], keep="first")
step("2_dedupe_iso3_year", "one row per (iso3, year); duplicates had identical values",
     d1, d2, "Country identified by iso3; each country counted at most once per year")
if len(dups):
    for iso, nm in multi_name.items():
        n_dup = int(dups[dups.iso3 == iso].shape[0] - dups[dups.iso3 == iso][["iso3", "year"]].drop_duplicates().shape[0])
        issues.append({"issue": f"{iso} appears under two country names {nm} (two extractions), giving {n_dup} duplicated (iso3, year) rows with identical life_expectancy values",
                       "handling": f"Kept one row per (iso3, year) (values identical, so no conflict); country_name harmonised to '{canonical.get(iso)}'; country counted once"})

# --- Step 3: drop countries lacking a row for some year ---
yrs_present = d2.groupby("iso3")["year"].nunique()
short = sorted(yrs_present[yrs_present < len(YEARS)].index)
d3 = d2[~d2["iso3"].isin(short)]
step("3_drop_countries_missing_year_rows", "drop countries without a row for every year 2000..2022",
     d2, d3, f"Not in balanced panel (no imputation allowed): {', '.join(short)}")
for iso in short:
    yy = d2[d2.iso3 == iso]["year"]
    issues.append({"issue": f"{iso} has rows only for {yy.min()}-{yy.max()} ({yy.nunique()} of 23 years)",
                   "handling": "Excluded from the balanced set; no imputation. Included in the 2022 all-available baseline"})

# --- Step 4: drop countries with an empty value in any year ---
nmiss = d3.groupby("iso3")["life_expectancy"].apply(lambda s: s.isna().sum())
withna = sorted(nmiss[nmiss > 0].index)
d4 = d3[~d3["iso3"].isin(withna)]
step("4_drop_countries_with_empty_value", "drop countries with an empty life_expectancy cell in any year",
     d3, d4, f"Missing country-year => not balanced (no interpolation): {', '.join(withna)}")
for iso in withna:
    yy = d3[(d3.iso3 == iso) & d3.life_expectancy.isna()]["year"].tolist()
    issues.append({"issue": f"{iso} has an empty life_expectancy cell for year(s) {yy}",
                   "handling": "Treated as missing (no interpolation); country excluded from the balanced set"})

analytic = d4[["iso3", "country_name", "year", "life_expectancy"]].sort_values(["iso3", "year"]).reset_index(drop=True)
balanced = sorted(analytic["iso3"].unique())
assert analytic.groupby("iso3")["year"].nunique().eq(len(YEARS)).all()
assert analytic["life_expectancy"].notna().all()
assert len(analytic) == len(balanced) * len(YEARS)

ym = analytic.groupby("year")["life_expectancy"].mean()
m2000, m2022 = float(ym.loc[2000]), float(ym.loc[2022])

# baseline: all distinct countries with a recorded 2022 value (after dedupe, each counted once)
b22 = d2[(d2.year == 2022) & d2.life_expectancy.notna()]
assert b22["iso3"].is_unique
base2022 = float(b22["life_expectancy"].mean())

n_total = int(d2["iso3"].nunique())
res = {
    "task_id": "F1-conf-02",
    "n_countries_total": n_total,
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": n_total - len(balanced),
    "balanced_countries": balanced,
    "n_analysed": int(len(analytic)),
    "mean_2000_balanced": round(m2000, 6),
    "mean_2022_balanced": round(m2022, 6),
    "change_2000_2022_balanced": round(m2022 - m2000, 6),
    "mean_2022_all_available": round(base2022, 6),
    "n_countries_2022_all_available": int(b22["iso3"].nunique()),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
    "balanced_yearly_means": {str(k): round(float(v), 6) for k, v in ym.items()},
}
analytic.to_csv(f"{OUT}/analytic.csv", index=False)
pd.DataFrame(ledger).to_csv(f"{OUT}/exclusions.csv", index=False)
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2, ensure_ascii=False)
print(f"balanced countries={len(balanced)}, n_analysed={len(analytic)}, "
      f"mean2000={m2000:.4f}, mean2022={m2022:.4f}, change={m2022-m2000:.4f}, baseline2022={base2022:.4f}")