"""F1-conf-02: balanced-panel mean life expectancy 2000-2022. Reads data/ only."""
import json, hashlib
from pathlib import Path
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "data" / "life_expectancy_panel.csv"
OUT = ROOT / "out"
YEARS = list(range(2000, 2023))

raw = pd.read_csv(SRC, dtype={"iso3": str, "country_name": str})
raw["iso3"] = raw["iso3"].str.strip()
raw["country_name"] = raw["country_name"].str.strip()
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))

# --- 1. duplicate (iso3, year) rows ---
dup_mask = raw.duplicated(["iso3", "year"], keep=False)
dups = raw[dup_mask]
for iso, g in dups.groupby("iso3"):
    # verify duplicates carry identical values (otherwise conflict -> cannot choose)
    conflict = g.groupby("year")["life_expectancy"].nunique(dropna=False).gt(1).any()
    if conflict:
        raise SystemExit(f"Conflicting duplicate values for {iso}; not resolvable without outside info")
    names = sorted(g["country_name"].unique())
    issues.append({"issue": f"iso3 {iso} appears under {len(names)} country_name spellings {names} "
                            f"with duplicated (iso3, year) rows for {g['year'].nunique()} years "
                            f"({len(g)} rows); values are identical across duplicates",
                   "handling": "kept one row per (iso3, year) (country identified by iso3); "
                               "duplicate rows removed; country counted once"})
# canonical name: last spelling alphabetically ('Turkiye' over 'Turkey'), current WDI name
canon = raw.groupby("iso3")["country_name"].agg(lambda s: sorted(s.unique())[-1])
df = raw.drop_duplicates(["iso3", "year"], keep="first").copy()
df["country_name"] = df["iso3"].map(canon)
step("1_dedup_iso3_year", "keep one row per (iso3, year)", n_raw, len(df),
     "same country (iso3) assembled twice under two name spellings with identical values")

# --- 2. incomplete countries ---
n_total = df["iso3"].nunique()
obs = df.dropna(subset=["life_expectancy"])
obs = obs[obs["year"].isin(YEARS)]
cnt = obs.groupby("iso3")["year"].nunique().reindex(sorted(df["iso3"].unique()), fill_value=0)
balanced = sorted(cnt[cnt == len(YEARS)].index)
incomplete = sorted(set(df["iso3"]) - set(balanced))
for iso in incomplete:
    yrs = sorted(set(obs.loc[obs.iso3 == iso, "year"]))
    missing = [y for y in YEARS if y not in yrs]
    blank = sorted(df.loc[(df.iso3 == iso) & df.life_expectancy.isna(), "year"].tolist())
    desc = (f"{iso}: no recorded value for {len(missing)} of 23 years "
            f"({missing[0]}-{missing[-1]}" + (")" if len(missing) == missing[-1]-missing[0]+1 else f"; years {missing})"))
    if blank:
        desc += f"; empty life_expectancy cell in year(s) {blank}"
    issues.append({"issue": desc, "handling": "treated as missing (no imputation); country excluded from balanced set"})

before = len(df)
part_missing = df[df.iso3.isin(incomplete) & (df.groupby("iso3")["year"].transform("size") < len(YEARS))]
# split exclusion into two transparent steps
short = [i for i in incomplete if (df.iso3 == i).sum() < len(YEARS)]
gaps = [i for i in incomplete if i not in short]
df2 = df[~df.iso3.isin(short)]
step("2_drop_countries_missing_rows", f"drop countries without a row for every year 2000-2022 ({', '.join(short)})",
     before, len(df2), "not observed in every year; balanced-panel estimand requires all 23 years")
df3 = df2[~df2.iso3.isin(gaps)]
step("3_drop_countries_with_empty_value", f"drop countries with an empty life_expectancy cell in some year ({', '.join(gaps)})",
     len(df2), len(df3), "missing value may not be imputed; country not in balanced set")

ana = df3[df3.year.isin(YEARS)].sort_values(["iso3", "year"])[["iso3", "country_name", "year", "life_expectancy"]]
assert ana.life_expectancy.notna().all() and len(ana) == len(balanced) * len(YEARS)
assert not ana.duplicated(["iso3", "year"]).any()
if len(ana) != len(df3):
    step("4_restrict_years", "keep years 2000-2022", len(df3), len(ana), "outside study window")

m2000 = float(ana.loc[ana.year == 2000, "life_expectancy"].mean())
m2022 = float(ana.loc[ana.year == 2022, "life_expectancy"].mean())
avail22 = df[(df.year == 2022) & df.life_expectancy.notna()]
assert not avail22.iso3.duplicated().any()
m22all = float(avail22["life_expectancy"].mean())
# descriptive dispersion of per-country changes (fixed set of countries, not a sample)
w = ana.pivot(index="iso3", columns="year", values="life_expectancy")
chg = (w[2022] - w[2000])
# sensitivity: naive 2022 mean if the duplicated TUR rows were NOT removed (wrong; shows impact)
m22_raw_dups = float(raw.loc[raw.year == 2022, "life_expectancy"].dropna().mean())

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(ana)),
    "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(m22all, 6),
    "n_countries_2022_available": int(avail22.iso3.nunique()),
    "change_per_country_sd": round(float(chg.std(ddof=1)), 6),
    "change_per_country_min": round(float(chg.min()), 6),
    "change_per_country_min_iso3": str(chg.idxmin()),
    "change_per_country_max": round(float(chg.max()), 6),
    "change_per_country_max_iso3": str(chg.idxmax()),
    "sensitivity_mean_2022_naive_with_duplicate_rows": round(m22_raw_dups, 6),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
    "input_sha256": hashlib.sha256(SRC.read_bytes()).hexdigest(),
}
OUT.mkdir(exist_ok=True)
ana.to_csv(OUT / "analytic.csv", index=False, float_format="%.3f")
pd.DataFrame(excl).to_csv(OUT / "exclusions.csv", index=False)
(OUT / "results.json").write_text(json.dumps(res, indent=2, ensure_ascii=False) + "\n")
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"]}))