"""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"
YEARS = list(range(2000, 2023))
os.makedirs("out", exist_ok=True)

raw = pd.read_csv(IN, dtype={"country_name": str, "iso3": str})
for c in ["country_name", "iso3"]:
    raw[c] = raw[c].str.strip()
n_read = len(raw)
issues, ledger = [], []

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

# --- integrity checks
assert raw["year"].between(2000, 2022).all()
names = raw.groupby("iso3")["country_name"].unique()
multi = names[names.apply(len) > 1]
for iso, nm in multi.items():
    issues.append({"issue": f"iso3 {iso} appears under several country names {sorted(nm)} (assembly of two extractions)",
                   "handling": "Countries identified by iso3 only; a single name kept (the last listed in the file, i.e. the later spelling)."})

# --- step 1: duplicate (iso3, year) rows
d = raw.copy()
dupmask = d.duplicated(["iso3", "year"], keep=False)
if dupmask.any():
    g = d[dupmask].groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False)
    conflicting = g[g > 1]
    if len(conflicting):
        raise SystemExit(f"Conflicting duplicate values: {conflicting}")
    n_dup_rows = int(d.duplicated(["iso3", "year"]).sum())
    issues.append({"issue": f"{n_dup_rows} duplicate (iso3, year) rows for {sorted(d[dupmask].iso3.unique())}; values identical across duplicates",
                   "handling": "Kept one row per (iso3, year) (last occurrence); since values are identical the choice does not affect results."})
before = len(d)
d = d.drop_duplicates(["iso3", "year"], keep="last").copy()
step("Remove duplicate (iso3, year) rows", before, len(d),
     "Each country contributes at most one observation per year; duplicates had identical values.")
# harmonise name per iso3 (last listed)
d["country_name"] = d["iso3"].map(raw.groupby("iso3")["country_name"].last())

# --- step 2: rows with missing value
miss = d[d["life_expectancy"].isna()]
for _, r in miss.iterrows():
    issues.append({"issue": f"Missing life_expectancy for {r.iso3} {r.year} (empty cell)",
                   "handling": "Treated as missing (no imputation); country therefore excluded from the balanced set."})
before = len(d)
d = d.dropna(subset=["life_expectancy"])
step("Remove country-years with no recorded life_expectancy", before, len(d),
     "Empty cell = no recorded value; no imputation allowed.")

# --- coverage gaps (absent rows)
cov = d.groupby("iso3")["year"].apply(lambda s: sorted(set(s)))
balanced = sorted([i for i, ys in cov.items() if ys == YEARS])
all_iso = sorted(raw["iso3"].unique())
for iso in all_iso:
    if iso in balanced:
        continue
    ys = cov.get(iso, [])
    missing_years = [y for y in YEARS if y not in ys]
    absent = [y for y in missing_years if not ((raw.iso3 == iso) & (raw.year == y)).any()]
    if absent:
        issues.append({"issue": f"{iso} has no rows for years {absent[0]}-{absent[-1]} ({len(absent)} years absent from file)",
                       "handling": "Not imputed; country excluded from the balanced set (still counted in the 2022 all-available baseline if 2022 recorded)."})

# --- step 3: restrict to balanced countries
before = len(d)
excluded = [i for i in all_iso if i not in balanced]
an = d[d["iso3"].isin(balanced)].sort_values(["iso3", "year"])[["iso3", "country_name", "year", "life_expectancy"]]
step(f"Keep only countries with a recorded value in every year 2000-2022 (drops {', '.join(excluded)})",
     before, len(an), "Balanced-panel estimand: fixed country set across all years.")
assert len(an) == len(balanced) * len(YEARS)
assert not an.duplicated(["iso3", "year"]).any()

an.to_csv("out/analytic.csv", index=False)
pd.DataFrame(ledger).to_csv("out/exclusions.csv", index=False)

m = an.groupby("year")["life_expectancy"].mean()
dedup = d  # one row per iso3-year, non-missing
m22_all = dedup[dedup.year == 2022]["life_expectancy"].mean()
n22_all = int((dedup.year == 2022).sum())

res = {
    "task_id": "F1-conf-02",
    "n_countries_total": len(all_iso),
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": len(all_iso) - len(balanced),
    "balanced_countries": balanced,
    "n_analysed": int(len(an)),
    "mean_2000_balanced": round(float(m[2000]), 6),
    "mean_2022_balanced": round(float(m[2022]), 6),
    "change_2000_2022_balanced": round(float(m[2022] - m[2000]), 6),
    "mean_2022_all_available": round(float(m22_all), 6),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_read)},
    "random_seed": None,
}
with open("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("n countries with 2022 value:", n22_all)
print(pd.DataFrame(ledger).to_string())

# --- supplementary descriptives (not part of results.json schema)
import numpy as np
from scipy import stats
w = an.pivot(index="iso3", columns="year", values="life_expectancy")
ch = w[2022] - w[2000]
se = ch.std(ddof=1) / np.sqrt(len(ch))
tcrit = stats.t.ppf(0.975, len(ch) - 1)
supp = {
    "country_change_2000_2022": {k: round(float(v), 3) for k, v in ch.sort_values().items()},
    "sd_country_change": round(float(ch.std(ddof=1)), 4),
    "se_mean_change_superpopulation": round(float(se), 4),
    "ci95_mean_change_superpopulation": [round(float(ch.mean() - tcrit * se), 4), round(float(ch.mean() + tcrit * se), 4)],
    "n_countries_2022_all_available": n22_all,
    "mean_2000_all_available": round(float(d[d.year == 2000]["life_expectancy"].mean()), 6),
    "n_countries_2000_all_available": int((d.year == 2000).sum()),
    "mean_2022_raw_rows_with_duplicates": round(float(raw[raw.year == 2022]["life_expectancy"].mean()), 6),
}
with open("out/supplementary.json", "w") as f:
    json.dump(supp, f, indent=2)
print(json.dumps(supp, indent=2))