"""F1-conf-01: 12-month CPI inflation from two ALFRED vintages.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd

D, O = "data", "out"
os.makedirs(O, exist_ok=True)
f_old, f_new, f_asm = "cpi_vintage_2016-07-15.csv", "cpi_vintage_2026-10-28.csv", "cpi_assembled.csv"
old = pd.read_csv(os.path.join(D, f_old), dtype={"cpi": float})
new = pd.read_csv(os.path.join(D, f_new), dtype={"cpi": float})
asm = pd.read_csv(os.path.join(D, f_asm), dtype={"cpi": float, "vintage": str})
n_input = {f_old: len(old), f_new: len(new), f_asm: len(asm)}
for d in (old, new, asm):
    d["date"] = pd.to_datetime(d["date"])
issues = []

# basic integrity checks on the authoritative vintage files
for name, d in ((f_old, old), (f_new, new)):
    assert not d["date"].duplicated().any(), name
    assert d["date"].is_monotonic_increasing, name
    full = pd.date_range(d["date"].min(), d["date"].max(), freq="MS")
    assert len(full) == len(d), f"{name}: gap in date rows"
    for dt in d.loc[d["cpi"].isna(), "date"]:
        issues.append({"issue": f"{name}: no CPI value reported for {dt:%Y-%m} (empty cell)",
                       "handling": "Treated as missing; not interpolated or imputed. No value exists in the older vintage either "
                                   "(it ends 2016-06), so the index level for this month is missing and g_t for this month "
                                   "(and for t+12, which lies outside the window) cannot be computed."})

# splice: newest vintage where it covers the month (row present with a value), else 2016 vintage
new_ok = new.dropna(subset=["cpi"]).assign(vintage="2026-10-28")
new_missing = set(new.loc[new["cpi"].isna(), "date"])
old_use = old[~old["date"].isin(new["date"])].assign(vintage="2016-07-15")
# months in overlap with missing new value: the most recent vintage in which the month is AVAILABLE could be the old one
fallback = old[old["date"].isin(new_missing)].dropna(subset=["cpi"]).assign(vintage="2016-07-15")
lev = pd.concat([old_use, fallback, new_ok]).sort_values("date").set_index("date")
assert not lev.index.duplicated().any()

ov = old.merge(new, on="date", suffixes=("_old", "_new")).dropna()
dif = ov["cpi_new"] - ov["cpi_old"]
issues.append({"issue": f"Vintages overlap for {len(ov)} months (2014-01..2016-06) and disagree (revisions incl. seasonal-factor "
                        f"updates): new-old ranges {dif.min():.3f} to {dif.max():.3f} index points.",
               "handling": "Per the availability rule the 2026-10-28 value is used for every overlapping month. Consequently "
                           "g_t for 2014-01..2014-12 divides a 2026-vintage level by a 2016-vintage level (2013 months exist "
                           "only in the old extract); this cross-vintage mixing is flagged in the report."})

# audit the hand-assembled file against the authoritative splice (documentation only; not used for estimates)
ref = lev["cpi"].reindex(pd.date_range("1947-01-01", "2026-09-01", freq="MS"))
a = asm.set_index("date")
wrong_vint = a.loc["2014-01-01":"2016-06-01"]
n_wv = int((wrong_vint["vintage"] == "2016-07-15").sum())
if n_wv:
    issues.append({"issue": f"cpi_assembled.csv uses the 2016-07-15 vintage for {n_wv} overlapping months 2014-01..2016-06 "
                            "instead of the most recent (2026-10-28) vintage.",
                   "handling": "Assembled file not used; levels rebuilt from the two vintage files."})
missing_rows = ref.index.difference(a.index)
extra_rows = a.index.difference(ref.index)
if len(missing_rows) or len(extra_rows):
    issues.append({"issue": f"cpi_assembled.csv lacks month(s) {[f'{x:%Y-%m}' for x in missing_rows]} and contains month(s) "
                            f"beyond the sources {[f'{x:%Y-%m}' for x in extra_rows]}: after the dropped month the 2026-vintage "
                            "values are shifted one month later (off-by-one dating).",
                   "handling": "Assembled file not used; levels rebuilt from the two vintage files."})
common = ref.index.intersection(a.index)
mism = (~np.isclose(a.loc[common, "cpi"], ref.loc[common], atol=1e-9) &
        ~(a.loc[common, "cpi"].isna() & ref.loc[common].isna()))
mism |= a.loc[common, "cpi"].isna() != ref.loc[common].isna()
issues.append({"issue": f"cpi_assembled.csv values disagree with the authoritative splice in {int(mism.sum())} of {len(common)} common months.",
               "handling": "Assembled file treated as unreliable and used only for this audit."})

# growth rates
win = pd.date_range("2001-01-01", "2026-09-01", freq="MS")
P = lev["cpi"].reindex(pd.date_range("1947-01-01", "2026-09-01", freq="MS"))
V = lev["vintage"].reindex(P.index)
lag = P.shift(12)
g = 100 * (P / lag - 1)
df = pd.DataFrame({"cpi": P, "cpi_lag12": lag, "vintage_used": V, "growth_12m": g}).loc[win]
n_win = len(df)
step1 = df.dropna(subset=["cpi"])
step2 = step1.dropna(subset=["cpi_lag12"])
an = step2
missing_months = [f"{x:%Y-%m}" for x in df.index[df["growth_12m"].isna()]]

ex = pd.DataFrame([
    {"step": 1, "rule": "Months 2001-01..2026-09 of the spliced series", "n_before": n_win, "n_after": n_win, "n_removed": 0,
     "rationale": "Analysis window defined in TASK.md (one row per calendar month)."},
    {"step": 2, "rule": "Drop months whose index level P_t is missing in all vintages", "n_before": n_win, "n_after": len(step1),
     "n_removed": n_win - len(step1), "rationale": "No interpolation/imputation allowed; 2025-10 has no value in the 2026-10-28 vintage."},
    {"step": 3, "rule": "Drop months whose lagged level P_{t-12} is missing", "n_before": len(step1), "n_after": len(step2),
     "n_removed": len(step1) - len(step2), "rationale": "g_t requires both levels; lag of 2026-10 would be 2025-10 but 2026-10 is outside the window."},
])
ex.to_csv(os.path.join(O, "exclusions.csv"), index=False)

out = an.reset_index().rename(columns={"index": "date"})
out["date"] = out["date"].dt.strftime("%Y-%m-01")
out[["date", "cpi", "cpi_lag12", "vintage_used", "growth_12m"]].to_csv(os.path.join(O, "analytic.csv"), index=False, float_format="%.10g")

gg = an["growth_12m"]
def gm(m):
    v = df.loc[pd.Timestamp(m + "-01"), "growth_12m"]
    return None if pd.isna(v) else round(float(v), 6)
# sensitivity (not a required estimand): 2014 growth using only the 2016-07-15 vintage (single-vintage ratio)
Po = old.set_index("date")["cpi"]
go = 100 * (Po / Po.shift(12) - 1)
sens_2014 = go.loc["2014-01-01":"2014-12-01"].mean()
sens_2014_06 = go.loc["2014-06-01"]
last_level = P.dropna().index.max()
res = {
    "task_id": "F1-conf-01",
    "first_growth_month": f"{an.index.min():%Y-%m}",
    "last_month": f"{last_level:%Y-%m}",
    "n_months_growth": int(len(an)),
    "n_months_missing_growth": int(n_win - len(an)),
    "n_analysed": int(len(an)),
    "mean_growth_12m": round(float(gg.mean()), 6),
    "mean_growth_12m_2001_2013": round(float(gg.loc["2001-01-01":"2013-12-01"].mean()), 6),
    "mean_growth_12m_2014_2026": round(float(gg.loc["2014-01-01":"2026-09-01"].mean()), 6),
    "growth_2014_06": gm("2014-06"), "growth_2016_07": gm("2016-07"), "growth_2020_05": gm("2020-05"),
    "growth_2022_06": gm("2022-06"), "growth_2026_09": gm("2026-09"),
    "max_growth_12m": round(float(gg.max()), 6), "max_growth_month": f"{gg.idxmax():%Y-%m}",
    "min_growth_12m": round(float(gg.min()), 6), "min_growth_month": f"{gg.idxmin():%Y-%m}",
    "feasibility": "complete" if not missing_months else "partial",
    "feasibility_reason": None if not missing_months else
        f"g_t not computable for {missing_months}: the 2026-10-28 vintage reports no CPI value for 2025-10 and imputation is "
        "not allowed. All other window months and all summary statistics are computed (means over available months).",
    "data_issues": issues,
    "n_input_rows": n_input,
    "random_seed": None,
    "sensitivity_mean_growth_2014_mixed_vintage": round(float(gg.loc["2014-01-01":"2014-12-01"].mean()), 6),
    "sensitivity_mean_growth_2014_old_vintage_only": round(float(sens_2014), 6),
    "sensitivity_growth_2014_06_old_vintage_only": round(float(sens_2014_06), 6),
}
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))