"""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. Deterministic; no randomness used."""
import json, os
import pandas as pd

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

raw = pd.read_csv(IN, dtype={"country_name": str, "iso3": str}, keep_default_na=True)
n_raw = len(raw)
issues, excl = [], []

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

df = raw.copy()
df["iso3"] = df["iso3"].str.strip()
df["year"] = df["year"].astype(int)

# 1. Exact/conflicting duplicates of (iso3, year)
dup = df[df.duplicated(["iso3", "year"], keep=False)]
if len(dup):
    g = dup.groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False)
    conflicting = g[g > 1]
    if len(conflicting):
        raise SystemExit(f"Conflicting duplicate values: {conflicting}")
    for iso, sub in dup.groupby("iso3"):
        names = sorted(sub.country_name.unique())
        issues.append({
            "issue": f"{iso}: {sub.duplicated(['iso3','year']).sum()} duplicated (iso3, year) rows "
                     f"from the two extractions, under country names {names}; values identical in every duplicated year",
            "handling": "Kept one row per (iso3, year) (first occurrence); country identified by iso3; "
                        "removed rows logged in exclusions.csv"})
b = len(df)
df = df.drop_duplicates(["iso3", "year"], keep="first")
step("1_dedupe", "drop duplicate (iso3, year) rows with identical values", b, len(df),
     "each country contributes at most one observation per year (TASK rule); duplicates are assembly artefacts")

# harmonise country name per iso3 (use most recent name in file for display)
name_map = {}
for iso, sub in raw.groupby("iso3"):
    names = list(dict.fromkeys(sub.country_name))
    name_map[iso.strip()] = names[-1]
    if len(names) > 1:
        issues.append({"issue": f"{iso}: multiple country_name spellings {names}",
                       "handling": f"Identified by iso3; analytic.csv uses '{names[-1]}' (last-listed name)"})
df["country_name"] = df["iso3"].map(name_map)

# out-of-window years
b = len(df)
df = df[df.year.between(2000, 2022)]
step("2_year_window", "keep years 2000-2022", b, len(df), "estimand window")

# 3. 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"})
b = len(df)
df = df[df.life_expectancy.notna()]
step("3_drop_missing", "drop rows with empty life_expectancy", b, len(df),
     "missing values are not imputed (TASK rule)")

# rows absent for countries
for iso, sub in df.groupby("iso3"):
    present = set(sub.year)
    absent = [y for y in YEARS if y not in present and not ((miss.iso3 == iso) & (miss.year == y)).any()]
    if absent:
        issues.append({"issue": f"{iso}: no rows at all for years {min(absent)}-{max(absent)} ({len(absent)} years); "
                                f"likely from the narrower earlier extraction",
                       "handling": "Treated as missing (no imputation); country excluded from balanced set; "
                                   "its 2022 value still enters the all-available baseline"})

# plausibility
if not df.life_expectancy.between(20, 100).all():
    raise SystemExit("implausible values")

all_countries = sorted(raw.iso3.str.strip().unique())
counts = df.groupby("iso3").year.nunique()
balanced = sorted(counts[counts == len(YEARS)].index)
incomplete = sorted(set(all_countries) - set(balanced))

b = len(df)
ana = df[df.iso3.isin(balanced)].sort_values(["iso3", "year"])
step("4_balanced", "keep countries with a recorded value in every year 2000-2022", b, len(ana),
     f"balanced-panel estimand; excluded incomplete countries {incomplete}")
ana = ana[["iso3", "country_name", "year", "life_expectancy"]]
assert len(ana) == len(balanced) * len(YEARS)
assert not ana.duplicated(["iso3", "year"]).any()
ana.to_csv(f"{OUT}/analytic.csv", index=False)

ex = pd.DataFrame(excl)
assert ex.n_before.iloc[0] == n_raw and ex.n_after.iloc[-1] == len(ana)
ex.to_csv(f"{OUT}/exclusions.csv", index=False)

m00 = float(ana[ana.year == 2000].life_expectancy.mean())
m22 = float(ana[ana.year == 2022].life_expectancy.mean())
avail22 = df[df.year == 2022].drop_duplicates("iso3")
m22_all = float(avail22.life_expectancy.mean())

# per-year balanced means (supplementary, not required)
per_year = ana.groupby("year").life_expectancy.mean().round(6)
per_year.to_csv(f"{OUT}/balanced_means_by_year.csv")

res = {
    "task_id": "F1-conf-02",
    "n_countries_total": len(all_countries),
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": len(incomplete),
    "balanced_countries": balanced,
    "incomplete_countries": incomplete,
    "n_analysed": int(len(ana)),
    "mean_2000_balanced": round(m00, 6),
    "mean_2022_balanced": round(m22, 6),
    "change_2000_2022_balanced": round(m22 - m00, 6),
    "mean_2022_all_available": round(m22_all, 6),
    "n_countries_2022_available": int(len(avail22)),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["n_countries_balanced", "n_analysed", "mean_2000_balanced",
                                      "mean_2022_balanced", "change_2000_2022_balanced",
                                      "mean_2022_all_available"]}, indent=1))