"""F1-conf-02: balanced-panel mean life expectancy 2000-2022. Run from workspace root."""
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={"iso3": str, "country_name": str})
n_read = len(raw)
issues, ledger = [], []

def step(name, rule, before, after, why):
    ledger.append(dict(step=len(ledger) + 1, rule=f"{name}: {rule}", n_before=before,
                       n_after=after, n_removed=before - after, rationale=why))

d = raw.copy()
# 1. Rows outside 2000-2022 (none expected)
b = len(d); d = d[d.year.isin(YEARS)]
step("year_window", "keep rows with year in 2000-2022", b, len(d), "Study period.")

# 2. Duplicate (iso3, year) rows from assembly of two extractions
dup = d[d.duplicated(["iso3", "year"], keep=False)]
if len(dup):
    conflicts = dup.groupby(["iso3", "year"]).life_expectancy.nunique(dropna=False).gt(1).sum()
    names = {k: sorted(v.unique()) for k, v in dup.groupby("iso3").country_name}
    if conflicts:
        raise SystemExit(f"Conflicting duplicate values: {conflicts}")
    issues.append({"issue": f"{len(dup) - dup[['iso3','year']].drop_duplicates().shape[0]} duplicate (iso3, year) rows "
                   f"from the two-extraction assembly: {names} (same iso3, different country_name spelling); "
                   "life_expectancy values identical in every duplicated pair.",
                   "handling": "Identified countries by iso3; kept one row per (iso3, year) (first occurrence), "
                   "using the latest-spelling country name as the label for the country."})
# canonical name: last-listed spelling per iso3 (later, wider extraction)
canon = d.groupby("iso3").country_name.last()
b = len(d); d = d.drop_duplicates(["iso3", "year"], keep="first").copy()
d["country_name"] = d.iso3.map(canon)
step("dedupe", "one row per (iso3, year)", b, len(d), "Each country counted once per year; duplicates were identical.")

# 3. Missing values
miss = d[d.life_expectancy.isna()]
if len(miss):
    issues.append({"issue": "Empty life_expectancy cell(s): " + ", ".join(f"{r.iso3} {r.year}" for r in miss.itertuples()),
                   "handling": "Treated as missing (no imputation); country excluded from the balanced set."})
b = len(d); d = d[d.life_expectancy.notna()]
step("drop_missing", "drop country-years with empty life_expectancy", b, len(d), "No imputation allowed.")

# 4. Incomplete coverage
cov = d.groupby("iso3").year.nunique()
all_iso = sorted(raw.iso3.unique())
balanced = sorted(cov[cov == len(YEARS)].index)
incomplete = [c for c in all_iso if c not in balanced]
for c in incomplete:
    yrs = sorted(d[d.iso3 == c].year)
    absent = [y for y in YEARS if y not in yrs]
    if c not in set(miss.iso3):
        issues.append({"issue": f"{c} has no row at all for {len(absent)} years ({min(absent)}-{max(absent)}); "
                       f"only {len(yrs)} of 23 years present (likely only in the narrower extraction).",
                       "handling": "Not imputed; country excluded from the balanced set (included in the 2022 all-available baseline)."})
b = len(d); an = d[d.iso3.isin(balanced)].sort_values(["iso3", "year"])
step("balanced_panel", "keep countries with a recorded value in every year 2000-2022", b, len(an),
     f"Balanced panel; removed countries: {', '.join(incomplete)}.")

an[["iso3", "country_name", "year", "life_expectancy"]].to_csv(f"{OUT}/analytic.csv", index=False)
pd.DataFrame(ledger).to_csv(f"{OUT}/exclusions.csv", index=False)

m = an.groupby("year").life_expectancy.mean()
all22 = d[d.year == 2022].drop_duplicates("iso3").life_expectancy.mean()
res = {
    "task_id": "F1-conf-02",
    "n_countries_total": len(all_iso),
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": len(incomplete),
    "balanced_countries": balanced,
    "n_analysed": int(len(an)),
    "mean_2000_balanced": round(float(m[2000]), 4),
    "mean_2022_balanced": round(float(m[2022]), 4),
    "change_2000_2022_balanced": round(float(m[2022] - m[2000]), 4),
    "mean_2022_all_available": round(float(all22), 4),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_read)},
    "random_seed": None,
}
with open(f"{OUT}/results.json", "w") as f:
    json.dump(res, f, indent=2, ensure_ascii=False)
print(json.dumps(res, indent=2, ensure_ascii=False))
print(m.round(3).to_string())

# Descriptive sensitivity checks (printed only; not part of results.json)
naive_dup = raw[(raw.year == 2022) & raw.life_expectancy.notna()].life_expectancy.mean()
with_che = d[d.iso3.isin(balanced + ["CHE"])]
print("SENS naive 2022 mean if TUR duplicate rows were not removed:", round(float(naive_dup), 4))
print("SENS 15+CHE mean 2000/2022/change:",
      [round(float(with_che[with_che.year == y].life_expectancy.mean()), 4) for y in (2000, 2022)],
      round(float(with_che[with_che.year == 2022].life_expectancy.mean() - with_che[with_che.year == 2000].life_expectancy.mean()), 4))
print("SENS n countries with 2022 value:", d[d.year == 2022].iso3.nunique())
print("SENS per-country change range (balanced):",
      (an[an.year == 2022].set_index("iso3").life_expectancy - an[an.year == 2000].set_index("iso3").life_expectancy).agg(["min", "idxmin", "max", "idxmax"]).to_dict())