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

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

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

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

# Step 1: duplicated (iso3, year) rows
dkey = raw[raw.duplicated(["iso3", "year"], keep=False)]
d = raw
if len(dkey):
    conflicts = dkey.groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False)
    if (conflicts > 1).any():
        raise SystemExit("Conflicting duplicate values found: " + str(conflicts[conflicts > 1]))
    for iso, g in dkey.groupby("iso3"):
        names = sorted(g["country_name"].unique())
        issues.append({"issue": f"{iso}: {g[['iso3','year']].drop_duplicates().shape[0]} country-years appear twice "
                                f"under names {names} with identical life_expectancy values (assembly duplicate of two extractions)",
                       "handling": "kept one row per (iso3, year) (last occurrence, name 'Turkiye' for TUR); "
                                   "country identified by iso3 as required"})
    d = raw.drop_duplicates(["iso3", "year"], keep="last")
log(1, "keep one row per (iso3, year); drop exact-value duplicates", n_raw, len(d),
    "duplicated country-years from two extractions (TUR as 'Turkey' and 'Turkiye'), identical values")

# Step 2: missing values
miss = d[d["life_expectancy"].isna()]
for _, r in miss.iterrows():
    issues.append({"issue": f"{r.iso3} {int(r.year)}: empty life_expectancy cell",
                   "handling": "treated as missing (no imputation); country therefore not in balanced set"})
d2 = d.dropna(subset=["life_expectancy"])
log(2, "drop rows with empty life_expectancy", len(d), len(d2), "no recorded value; imputation not allowed")

# Step 3: restrict to balanced countries
d2 = d2[d2["year"].isin(YEARS)]
cover = d2.groupby("iso3")["year"].nunique()
all_iso = sorted(d["iso3"].unique())
balanced = sorted([c for c in all_iso if cover.get(c, 0) == len(YEARS)])
incomplete = [c for c in all_iso if c not in balanced]
for c in incomplete:
    yrs = sorted(d2.loc[d2.iso3 == c, "year"])
    if c in miss.iso3.values:
        continue
    issues.append({"issue": f"{c}: recorded only for {min(yrs)}-{max(yrs)} ({len(yrs)} of {len(YEARS)} years); rows for 2000-{min(yrs)-1} absent from file",
                   "handling": "excluded from balanced set; included in 2022 all-available baseline"})
an = d2[d2["iso3"].isin(balanced)].sort_values(["iso3", "year"])
log(3, "keep only countries with a recorded value in every year 2000-2022", len(d2), len(an),
    f"balanced-panel definition; excluded {incomplete}")

an = an[["iso3", "country_name", "year", "life_expectancy"]]
assert len(an) == len(balanced) * len(YEARS) and not an.duplicated(["iso3", "year"]).any()

m = an.groupby("year")["life_expectancy"].mean()
m2000, m2022 = float(m[2000]), float(m[2022])
base = d2[d2.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(m2000, 6),
    "mean_2022_balanced": round(m2022, 6),
    "change_2000_2022_balanced": round(m2022 - m2000, 6),
    "mean_2022_all_available": round(float(base), 6),
    "n_countries_2022_all_available": int(d2[d2.year == 2022]["iso3"].nunique()),
    "incomplete_countries": incomplete,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
}
os.makedirs(OUT, exist_ok=True)
an.to_csv(f"{OUT}/analytic.csv", index=False)
pd.DataFrame(excl).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(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))