"""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={"country_name": str, "iso3": str, "year": int})
raw["iso3"] = raw["iso3"].str.strip()
n_read = len(raw)
issues, ledger = [], []

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

n_total = raw["iso3"].nunique()

# 1. duplicate (iso3, year) rows
dup_mask = raw.duplicated(["iso3", "year"], keep=False)
dups = raw[dup_mask]
conflict = dups.groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False).gt(1)
if conflict.any():
    raise SystemExit(f"Conflicting duplicate values: {conflict[conflict].index.tolist()}")
for iso, g in dups.groupby("iso3"):
    names = sorted(g["country_name"].unique())
    issues.append({"issue": f"{iso}: every (iso3, year) for {g['year'].nunique()} years appears twice "
                            f"under different country_name labels {names} (assembly of two extractions); "
                            "values are identical across the copies.",
                   "handling": "Identified countries by iso3 and kept one row per (iso3, year) "
                               "(the later-listed copy, i.e. current API label); values unaffected."})
d = raw.drop_duplicates(["iso3", "year"], keep="last").copy()
step("deduplicate", "keep one row per (iso3, year)", n_read, len(d),
     "Exact-duplicate country-years from merging two extractions (same values, different name label).")
# harmonise name to one label per iso3 (last-listed)
d["country_name"] = d["iso3"].map(d.groupby("iso3")["country_name"].last())

# 2. missing values
miss = d[d["life_expectancy"].isna()]
for _, r in miss.iterrows():
    issues.append({"issue": f"{r.iso3} {r.year}: empty life_expectancy cell (no recorded value).",
                   "handling": "Treated as missing; not imputed. Country therefore not in balanced set."})
before = len(d)
d = d[d["life_expectancy"].notna()]
step("drop_missing", "drop rows with empty life_expectancy", before, len(d),
     "No recorded value; imputation/interpolation is prohibited.")

# 3. restrict to study years (sanity)
before = len(d)
d = d[d["year"].between(2000, 2022)]
step("year_window", "keep years 2000-2022", before, len(d), "Study window.")

# 4. balanced set
cov = d.groupby("iso3")["year"].apply(lambda s: set(s))
balanced = sorted(i for i, s in cov.items() if set(YEARS) <= s)
incomplete = sorted(set(raw["iso3"]) - set(balanced))
for iso in incomplete:
    have = cov.get(iso, set())
    missing_years = sorted(set(YEARS) - have)
    if missing_years and len(missing_years) > 1 and missing_years == list(range(missing_years[0], missing_years[-1] + 1)):
        desc = f"{missing_years[0]}-{missing_years[-1]}"
    else:
        desc = ",".join(map(str, missing_years))
    if iso not in set(miss["iso3"]):
        issues.append({"issue": f"{iso}: no rows at all for years {desc} (series only present "
                                f"{min(have)}-{max(have)}; likely from the narrower extraction).",
                       "handling": "Country excluded from balanced panel; still counted in the 2022 all-available baseline."})
before = len(d)
d = d[d["iso3"].isin(balanced)]
step("balanced_panel", "keep countries with a recorded value in every year 2000-2022", before, len(d),
     f"Excluded incomplete countries: {', '.join(incomplete)}.")

analytic = d[["iso3", "country_name", "year", "life_expectancy"]].sort_values(["iso3", "year"])
assert len(analytic) == len(balanced) * len(YEARS)
assert not analytic.duplicated(["iso3", "year"]).any()
analytic.to_csv(f"{OUT}/analytic.csv", index=False)
pd.DataFrame(ledger).to_csv(f"{OUT}/exclusions.csv", index=False)

m = analytic.groupby("year")["life_expectancy"].mean()
avail22 = raw.drop_duplicates(["iso3", "year"])
avail22 = avail22[(avail22["year"] == 2022) & avail22["life_expectancy"].notna()]
w = analytic.pivot(index="iso3", columns="year", values="life_expectancy")
chg = w[2022] - w[2000]
res = {
    "task_id": "F1-conf-02",
    "n_countries_total": int(n_total),
    "n_countries_balanced": len(balanced),
    "n_countries_incomplete": len(incomplete),
    "balanced_countries": balanced,
    "n_analysed": int(len(analytic)),
    "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(avail22["life_expectancy"].mean()), 4),
    "n_countries_2022_available": int(avail22["iso3"].nunique()),
    "descriptive_extra": {
        "country_change_min": round(float(chg.min()), 3), "country_change_min_iso3": chg.idxmin(),
        "country_change_max": round(float(chg.max()), 3), "country_change_max_iso3": chg.idxmax(),
        "country_change_sd": round(float(chg.std()), 4),
        "country_change_se_illustrative": round(float(chg.std() / len(chg) ** 0.5), 4),
        "mean_2022_balanced_minus_all_available": round(float(m[2022] - avail22["life_expectancy"].mean()), 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(pd.DataFrame(ledger).to_string())