"""F1-conf-02: balanced-panel mean life expectancy 2000-2022.
Reads only data/life_expectancy_panel.csv; writes out/results.json, out/analytic.csv, out/exclusions.csv.
Deterministic; no randomness used."""
import json, os
import pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
IN = os.path.join(ROOT, "data", "life_expectancy_panel.csv")
OUT = os.path.join(ROOT, "out")
YEARS = list(range(2000, 2023))

raw = pd.read_csv(IN, dtype={"iso3": str, "country_name": str},
                  keep_default_na=False, na_values=[""])
raw["iso3"] = raw["iso3"].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
d = raw.copy()
dup_mask = d.duplicated(["iso3", "year"], keep=False)
dups = d[dup_mask]
for iso, g in dups.groupby("iso3"):
    conflicting = g.groupby("year")["life_expectancy"].nunique(dropna=False).gt(1).sum()
    if conflicting:
        raise RuntimeError(f"Conflicting duplicate values for {iso}; manual resolution needed")
    names = sorted(g["country_name"].unique())
    issues.append({
        "issue": f"{iso} appears twice for each of {g['year'].nunique()} years under names {names} "
                 f"(assembly of two extractions); duplicate values are identical",
        "handling": "Kept one row per (iso3, year) (identity is iso3); dropped the duplicate rows; "
                    "country_name harmonised to the last-listed spelling in the file"})
# keep the last occurrence (later, wider extraction) -> name 'Turkiye'
d = d.drop_duplicates(["iso3", "year"], keep="last")
# harmonise country_name per iso3 to the name of its last row
last_name = d.groupby("iso3")["country_name"].last()
d["country_name"] = d["iso3"].map(last_name)
step("1_dedupe", "drop duplicate (iso3, year) rows; keep one per country-year",
     n_raw, len(d), "each country contributes at most one observation per year; duplicates had identical values")

# 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",
                   "handling": "Treated as missing (no imputation); country excluded from the balanced set"})
n0 = len(d)
d = d[d["life_expectancy"].notna()]
step("2_drop_missing", "drop rows with empty life_expectancy", n0, len(d),
     "no imputation allowed; a country-year without a recorded value is missing")

# 3. Year range check
n0 = len(d)
d = d[d["year"].between(2000, 2022)]
step("3_year_range", "keep years 2000-2022", n0, len(d), "study window")

# 4. Balanced panel
yrs = d.groupby("iso3")["year"].apply(lambda s: set(s))
balanced = sorted([i for i, s in yrs.items() if set(YEARS) <= s])
all_iso = sorted(raw["iso3"].unique())
incomplete = [i for i in all_iso if i not in balanced]
for i in incomplete:
    have = sorted(yrs.get(i, set()))
    if have and have == list(range(have[0], have[-1] + 1)) and have[0] > 2000:
        issues.append({"issue": f"{i}: rows only for {have[0]}-{have[-1]} (years 2000-{have[0]-1} absent from file)",
                       "handling": "Not imputed; country excluded from the balanced set (still counted in 2022 all-available baseline)"})
n0 = len(d)
a = d[d["iso3"].isin(balanced)].sort_values(["iso3", "year"]).reset_index(drop=True)
step("4_balanced_panel", "keep countries with a recorded value in every year 2000-2022", n0, len(a),
     f"balanced-panel estimand; excluded countries: {', '.join(incomplete)}")
assert len(a) == len(balanced) * len(YEARS)
assert not a.duplicated(["iso3", "year"]).any()

a[["iso3", "country_name", "year", "life_expectancy"]].to_csv(os.path.join(OUT, "analytic.csv"), index=False)
ex = pd.DataFrame(excl)
ex.to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
assert ex["n_removed"].sum() == n_raw - len(a)

m00 = float(a.loc[a.year == 2000, "life_expectancy"].mean())
m22 = float(a.loc[a.year == 2022, "life_expectancy"].mean())
avail22 = d[d.year == 2022].drop_duplicates("iso3")  # already deduplicated
m22_all = float(avail22["life_expectancy"].mean())

chg = (a[a.year == 2022].set_index("iso3")["life_expectancy"]
       - a[a.year == 2000].set_index("iso3")["life_expectancy"])
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(a)),
    "mean_2000_balanced": round(m00, 6),
    "mean_2022_balanced": round(m22, 6),
    "change_2000_2022_balanced": round(m22 - m00, 6),
    "mean_2022_all_available": round(m22_all, 6),
    "n_countries_2022_available": int(avail22["iso3"].nunique()),
    "incomplete_countries": incomplete,
    "gap_2022_all_available_minus_balanced": round(m22_all - m22, 6),
    "per_country_change_min": round(float(chg.min()), 6),
    "per_country_change_min_iso3": str(chg.idxmin()),
    "per_country_change_max": round(float(chg.max()), 6),
    "per_country_change_max_iso3": str(chg.idxmax()),
    "per_country_change_sd": round(float(chg.std(ddof=1)), 6),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"life_expectancy_panel.csv": int(n_raw)},
    "random_seed": None,
}
with open(os.path.join(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))