"""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})
n_read = len(raw)
for c in ["country_name", "iso3"]:
    raw[c] = raw[c].str.strip()
raw["year"] = raw["year"].astype(int)
d = raw[raw["year"].between(2000, 2022)].copy()
assert len(d) == n_read, "rows outside 2000-2022 found"

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))

# 1. Duplicate (iso3, year) rows
names = d.groupby("iso3")["country_name"].unique()
multi = {k: sorted(v) for k, v in names.items() if len(v) > 1}
dupkey = d.duplicated(["iso3", "year"], keep=False)
dups = d[dupkey]
conflict = dups.groupby(["iso3", "year"])["life_expectancy"].nunique(dropna=False).gt(1).any()
assert not conflict, "conflicting duplicate values - needs manual decision"
# canonical name: the current World Bank spelling ("Turkiye") where alternatives exist
canon = {}
for k, v in multi.items():
    canon[k] = "Turkiye" if "Turkiye" in v else v[-1]
d["country_name"] = d.apply(lambda r: canon.get(r["iso3"], r["country_name"]), axis=1)
b = len(d)
d = d.drop_duplicates(["iso3", "year"], keep="first")
step("Drop duplicate (iso3, year) rows (identical values)", b, len(d),
     "Panel assembled from two extractions; some countries appear twice under different names. iso3 identifies the country; one observation per country-year.")
for k, v in multi.items():
    n = int(dups[dups.iso3 == k].shape[0] // 2)
    issues.append({"issue": f"{k} appears under two country names {v} with {n} duplicated country-years (identical values) — an assembly artifact of the two extractions.",
                   "handling": f"Deduplicated by (iso3, year), keeping one row per year; values identical so no choice of value needed; name harmonised to '{canon[k]}'. Counted as one country."})

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

# 3. Restrict to balanced countries
cov = d.groupby("iso3")["year"].apply(lambda s: set(s))
balanced = sorted([k for k, s in cov.items() if set(YEARS) <= s])
all_iso = sorted(raw["iso3"].unique())
incomplete = [k for k in all_iso if k not in balanced]
for k in incomplete:
    have = sorted(cov.get(k, set()))
    missing_years = sorted(set(YEARS) - set(have))
    if k not in set(miss.iso3):
        issues.append({"issue": f"{k} has no rows for years {missing_years[0]}–{missing_years[-1]} ({len(missing_years)} years missing from the file; the country is present only from 2011)." if missing_years == list(range(missing_years[0], missing_years[-1] + 1)) else f"{k} missing years {missing_years}.",
                       "handling": "Country excluded from the balanced panel; its available rows are used only in the 2022 all-available baseline."})
b = len(d)
an = d[d["iso3"].isin(balanced)].copy()
step("Keep only countries with a recorded value in every year 2000-2022", b, len(an),
     f"Balanced-panel estimand; excluded incomplete countries: {', '.join(incomplete)}.")

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

m00 = an.loc[an.year == 2000, "life_expectancy"].mean()
m22 = an.loc[an.year == 2022, "life_expectancy"].mean()
all22 = d[d.year == 2022]  # deduplicated, non-missing; one row per iso3
assert all22.iso3.is_unique
w = an.pivot(index="iso3", columns="year", values="life_expectancy")
chg = w[2022] - w[2000]
from scipy import stats
t95 = stats.t.ppf(0.975, len(chg) - 1)
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": len(an),
    "mean_2000_balanced": round(float(m00), 4),
    "mean_2022_balanced": round(float(m22), 4),
    "change_2000_2022_balanced": round(float(m22 - m00), 4),
    "mean_2022_all_available": round(float(all22.life_expectancy.mean()), 4),
    "n_countries_2022_all_available": int(all22.iso3.nunique()),
    "descriptive_country_change_sd": round(float(chg.std(ddof=1)), 4),
    "descriptive_country_change_min": {chg.idxmin(): round(float(chg.min()), 3)},
    "descriptive_country_change_max": {chg.idxmax(): round(float(chg.max()), 3)},
    "descriptive_change_t95_halfwidth_across_countries": round(float(t95 * chg.std(ddof=1) / len(chg) ** 0.5), 4),
    "descriptive_mean_2000_all_available": round(float(d.loc[d.year == 2000, "life_expectancy"].mean()), 4),
    "descriptive_n_countries_2000_all_available": int((d.year == 2000).sum()),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": 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))