"""F1-conf-01: 12-month CPIAUCSL growth from two ALFRED vintages.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Authoritative inputs are the two vintage files; the assembled file is read only for an audit."""
import json, os
import numpy as np
import pandas as pd

D, O = "data", "out"
OLD, NEW, ASM = "cpi_vintage_2016-07-15.csv", "cpi_vintage_2026-10-28.csv", "cpi_assembled.csv"
W0, W1 = pd.Timestamp("2001-01-01"), pd.Timestamp("2026-09-01")

def load(f):
    d = pd.read_csv(os.path.join(D, f), dtype={"date": str})
    d["date"] = pd.to_datetime(d["date"])
    return d

old, new, asm = load(OLD), load(NEW), load(ASM)
n_rows = {OLD: len(old), NEW: len(new), ASM: len(asm)}
issues = []

# basic integrity checks on the vintage files
for name, d in ((OLD, old), (NEW, new)):
    assert not d["date"].duplicated().any(), name
    full = pd.date_range(d["date"].min(), d["date"].max(), freq="MS")
    assert len(full) == len(d), f"gap in date sequence of {name}"
    for dt in d.loc[d["cpi"].isna(), "date"]:
        issues.append({"issue": f"{name}: no CPI value published for {dt:%Y-%m} (empty cell)",
                       "handling": "Left missing, not imputed; 12-month growth for that month (and for the month 12 later, if in the window) reported as missing."})

# overlap / revisions between vintages
ov = old.merge(new, on="date", suffixes=("_old", "_new")).dropna()
nrev = int((ov["cpi_old"] != ov["cpi_new"]).sum())
issues.append({"issue": f"Vintages overlap for {ov['date'].min():%Y-%m}..{ov['date'].max():%Y-%m} ({len(ov)} months); {nrev} of them differ (revisions, e.g. seasonal-factor updates; max abs diff {np.abs(ov.cpi_old-ov.cpi_new).max():.3f} index points).",
               "handling": "Per the availability rule the 2026-10-28 vintage is used for every month it covers (2014-01 onward, including the overlap); 2016-07-15 vintage only before 2014-01."})
issues.append({"issue": "Growth rates for 2014-01..2014-12 divide a 2026-vintage level by a 2013 level from the 2016-07-15 vintage (mixed vintages), because the 2026 extract starts at 2014-01.",
               "handling": "Computed as specified by the rule (most recent available vintage per month); flagged as a limitation since 2013 levels may have been revised since 2016."})

# combine: newest vintage where it covers the month (including months it covers with an empty cell -> missing, not back-filled)
o = old.assign(vintage_used="2016-07-15").set_index("date")
n = new.assign(vintage_used="2026-10-28").set_index("date")
comb = pd.concat([o[o.index < n.index.min()], n]).sort_index()
newer_missing = n.index[n["cpi"].isna()]
for dt in newer_missing:
    if dt in o.index:
        issues.append({"issue": f"{dt:%Y-%m} missing in 2026 vintage but present in 2016 vintage", "handling": "Not back-filled from the older vintage."})

# audit of the hand-assembled file
a = asm.set_index("date")
exp = comb.reindex(a.index)
val_mis = a.index[~((a["cpi"] == exp["cpi"]) | (a["cpi"].isna() & exp["cpi"].isna()))]
vin_mis = a.index[a["vintage"].astype(str) != exp["vintage_used"].fillna("none")]
missing_dates = comb.index.difference(a.index)
extra_dates = a.index.difference(comb.index)
ov_old = a.loc[(a.index >= n.index.min()) & (a.index <= o.index.max())]
issues.append({"issue": f"cpi_assembled.csv contains assembly errors: overlap months {ov_old.index.min():%Y-%m}..{ov_old.index.max():%Y-%m} taken from the older 2016 vintage; months absent: {[f'{x:%Y-%m}' for x in missing_dates]}; spurious months not in any source: {[f'{x:%Y-%m}' for x in extra_dates]}; from 2016-08 onward values are shifted one month later than in the source (so its empty cell sits at 2025-11 instead of 2025-10); in total {len(val_mis)} rows have a value and {len(vin_mis)} rows a vintage label that disagree with the authoritative vintage files.",
               "handling": "The assembled file is not used for any computation; the series is rebuilt directly from the two vintage files (authoritative per TASK.md)."})

# growth
s = comb["cpi"]
full = pd.date_range(s.index.min(), s.index.max(), freq="MS")
s = s.reindex(full)
lag = s.shift(12)
g = 100 * (s / lag - 1)
win = pd.date_range(W0, W1, freq="MS")
tab = pd.DataFrame({"cpi": s.reindex(win), "cpi_lag12": lag.reindex(win),
                    "vintage_used": comb["vintage_used"].reindex(win), "growth_12m": g.reindex(win)})
ok = tab["growth_12m"].notna()
an = tab[ok]
miss = tab.index[~ok]
out = an.reset_index().rename(columns={"index": "date"})
out["date"] = out["date"].dt.strftime("%Y-%m-01")
out["growth_12m"] = out["growth_12m"].round(10)
out.to_csv(os.path.join(O, "analytic.csv"), index=False)

no_level = tab.index[tab["cpi"].isna()]
no_lag = tab.index[tab["cpi"].notna() & tab["cpi_lag12"].isna()]
steps = [
    ("1", "Months 2001-01..2026-09 in the analysis window (one row per calendar month)", len(win), len(win), 0,
     "Starting frame: every month of the stated window, levels from most recent vintage covering the month."),
    ("2", "Drop months whose own index level P_t is missing in the source", len(win), len(win) - len(no_level), len(no_level),
     "No imputation allowed; missing: " + ", ".join(f"{x:%Y-%m}" for x in no_level)),
    ("3", "Drop months whose lagged level P_{t-12} is missing", len(win) - len(no_level), len(win) - len(no_level) - len(no_lag), len(no_lag),
     "No imputation allowed; missing: " + (", ".join(f"{x:%Y-%m}" for x in no_lag) or "none")),
]
pd.DataFrame(steps, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(os.path.join(O, "exclusions.csv"), index=False)
assert steps[-1][3] == len(an)

gg = an["growth_12m"]
def gm(d):
    v = g.get(pd.Timestamp(d))
    return None if v is None or pd.isna(v) else float(v)
last_level = s.dropna().index.max()
res = {
    "task_id": "F1-conf-01",
    "first_growth_month": f"{an.index.min():%Y-%m}",
    "last_month": f"{min(last_level, W1):%Y-%m}",
    "n_months_growth": int(len(an)),
    "n_months_missing_growth": int(len(miss)),
    "n_analysed": int(len(an)),
    "mean_growth_12m": float(gg.mean()),
    "mean_growth_12m_2001_2013": float(gg[:"2013-12-01"].mean()),
    "mean_growth_12m_2014_2026": float(gg["2014-01-01":].mean()),
    "growth_2014_06": gm("2014-06-01"), "growth_2016_07": gm("2016-07-01"),
    "growth_2020_05": gm("2020-05-01"), "growth_2022_06": gm("2022-06-01"),
    "growth_2026_09": gm("2026-09-01"),
    "max_growth_12m": float(gg.max()), "max_growth_month": f"{gg.idxmax():%Y-%m}",
    "min_growth_12m": float(gg.min()), "min_growth_month": f"{gg.idxmin():%Y-%m}",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": n_rows,
    "random_seed": None,
    "missing_growth_months": [f"{x:%Y-%m}" for x in miss],
    "n_2001_2013": int(len(gg[:"2013-12-01"])), "n_2014_2026": int(len(gg["2014-01-01":])),
}
# sensitivity (not the estimand): old-vintage-only growth in overlap, and naive use of the assembled file
so = old.set_index("date")["cpi"]; go = 100 * (so / so.shift(12) - 1)
sa = a["cpi"].reindex(pd.date_range(a.index.min(), a.index.max(), freq="MS")); ga = (100 * (sa / sa.shift(12) - 1)).reindex(win).dropna()
res["sensitivity"] = {
    "growth_2014_06_vintage2016": float(go["2014-06-01"]),
    "mean_growth_2014_01_2016_06_vintage2016": float(go["2014-01-01":"2016-06-01"].mean()),
    "mean_growth_2014_01_2016_06_used": float(gg["2014-01-01":"2016-06-01"].mean()),
    "assembled_file_naive_mean_growth_12m": float(ga.mean()),
    "assembled_file_naive_n_growth": int(len(ga)),
    "assembled_file_naive_growth_2026_09": float(ga.get(pd.Timestamp("2026-09-01"), np.nan)),
}
with open(os.path.join(O, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k != "data_issues"}, indent=2))
for i in issues: print("-", i["issue"])