"""T1-forecast: real-time forecasting of US retail-sales YoY growth from ALFRED vintages.
Run from workspace root:  ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np, pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATA, OUT = "data", "out"
ORIGINS = ["2025-10-31", "2025-11-30", "2025-12-31", "2026-01-31", "2026-02-28",
           "2026-03-31", "2026-04-30", "2026-05-31", "2026-06-30", "2026-07-31"]
HORIZONS = [1, 3]
DRIFT_WINDOW = 36  # months of m/m log changes used for the drift (fixed a priori)
SEED = 20261030
n_rows = {}

def read(path):
    d = pd.read_csv(os.path.join(DATA, path), parse_dates=["observation_date"])
    n_rows[path] = int(len(d))
    s = d.set_index("observation_date").iloc[:, 0].astype(float)
    s.index.freq = None
    return s

def yoy(level):
    return 100 * np.log(level / level.shift(12))

fr = pd.read_csv(os.path.join(DATA, "first_release.csv"), parse_dates=["target_month"])
n_rows["first_release.csv"] = int(len(fr))
actual = fr.set_index("target_month")["yoy_first_release"]

def add_months(ts, k):
    return ts + pd.DateOffset(months=k)

rows, audit = [], []
for V in ORIGINS:
    L = read(f"vintages/RSAFS_vintage_{V}.csv")
    P = read(f"vintages/PAYEMS_vintage_{V}.csv")  # read for audit; not used by the model
    assert L.isna().sum() == 0
    g = yoy(L)
    last = L.index[-1]
    vdate = pd.Timestamp(V)
    lag = (vdate.year - last.year) * 12 + vdate.month - last.month
    audit.append({"origin": V, "last_month": str(last.date()), "months_behind": int(lag),
                  "payems_last": str(P.index[-1].date())})
    # AR(1) on g, pairs with 1993-02 <= m <= last
    gg = g.dropna()
    y = gg.loc["1993-02-01":last]
    x = gg.shift(1).loc[y.index]
    X = np.column_stack([np.ones(len(y)), x.values])
    beta = np.linalg.lstsq(X, y.values, rcond=None)[0]
    # model: drift-in-levels with exact base effects
    dl = 100 * np.log(L).diff().dropna()
    mu = dl.iloc[-DRIFT_WINDOW:].mean()
    for h in HORIZONS:
        T = add_months(last, h)
        f_naive = g.loc[last]
        f_snaive = g.loc[add_months(T, -12)]
        f = g.loc[last]
        for _ in range(h):
            f = beta[0] + beta[1] * f
        f_ar1 = f
        # g_T = 100 ln(L_last/L_{T-12}) + h*mu  (L_{T-12} is in the vintage)
        f_model = 100 * np.log(L.loc[last] / L.loc[add_months(T, -12)]) + h * mu
        for name, val in [("naive", f_naive), ("snaive", f_snaive), ("ar1", f_ar1), ("model", f_model)]:
            rows.append({"origin_date": V, "target_date": T.strftime("%Y-%m-%d"), "horizon": h,
                         "forecast": float(val), "actual_first_release": float(actual.loc[T]),
                         "model": name, "vintage_used": V})

fc = pd.DataFrame(rows)
fc["error"] = fc["actual_first_release"] - fc["forecast"]
fc.drop(columns="error").to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")

tab = []
for m in ["naive", "snaive", "ar1", "model"]:
    for h in HORIZONS:
        e = fc[(fc.model == m) & (fc.horizon == h)]["error"].values
        tab.append({"model": m, "horizon": h, "rmse": float(np.sqrt(np.mean(e**2))),
                    "mae": float(np.mean(np.abs(e))), "n": int(len(e))})
tab = pd.DataFrame(tab)
tab.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

def dm(e_naive, e_model, h):
    d = e_naive**2 - e_model**2
    n = len(d); dbar = d.mean(); dc = d - dbar
    lrv = np.sum(dc**2) / n
    for j in range(1, h):
        w = 1 - j / h
        lrv += 2 * w * np.sum(dc[j:] * dc[:-j]) / n
    stat = dbar / np.sqrt(lrv / n)
    return float(stat), float(2 * (1 - stats.norm.cdf(abs(stat))))

res = {"task_id": "T1-forecast",
       "model_name": "Drift-in-levels with exact base effects (RW-drift36)",
       "model_description": (
           "The YoY growth of target month T = last(V)+h is written exactly as "
           "g_T = 100*ln(L_last/L_{T-12}) + sum_{k=1..h} 100*dln L_{last+k}. The first term uses only "
           "levels published in the RSAFS vintage dated V (L_{T-12} is always in the vintage), so the "
           "'base effect' of months dropping out of the 12-month window is known exactly at the origin. "
           "The unknown future monthly log changes are each forecast by their mean over the last 36 "
           "monthly log changes available in the same vintage (a random walk with drift in log levels; "
           "window fixed a priori, not tuned on evaluation targets). Re-estimated at every origin from the "
           "vintage dated V only; PAYEMS and UMCSENT are not used (no information leakage possible); "
           "RSAFS_current is not used."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
for _, r in tab.iterrows():
    res[f"rmse_{r.model}_h{r.horizon}"] = r.rmse
    res[f"mae_{r.model}_h{r.horizon}"] = r.mae

rng = np.random.default_rng(SEED)
boot = {}
for h in HORIZONS:
    en = fc[(fc.model == "naive") & (fc.horizon == h)].sort_values("origin_date")["error"].values
    em = fc[(fc.model == "model") & (fc.horizon == h)].sort_values("origin_date")["error"].values
    res[f"rel_improvement_rmse_h{h}"] = float(1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"])
    s, p = dm(en, em, h)
    res[f"dm_stat_h{h}"], res[f"dm_pvalue_h{h}"] = s, p
    better = res[f"rmse_model_h{h}"] < res[f"rmse_naive_h{h}"]
    sig = p < 0.05
    res[f"conclusion_code_h{h}"] = ("model_better_" if better else "naive_better_") + ("significant" if sig else "not_significant")
    # circular moving-block bootstrap (block length max(h,2)) of RMSE_naive - RMSE_model
    n = len(en); b = max(h, 2); B = 9999; diffs = np.empty(B)
    for i in range(B):
        starts = rng.integers(0, n, int(np.ceil(n / b)))
        idx = np.concatenate([(np.arange(s0, s0 + b) % n) for s0 in starts])[:n]
        diffs[i] = np.sqrt(np.mean(en[idx]**2)) - np.sqrt(np.mean(em[idx]**2))
    boot[h] = {"rmse_diff_naive_minus_model": float(np.sqrt(np.mean(en**2)) - np.sqrt(np.mean(em**2))),
               "ci95": [float(np.percentile(diffs, 2.5)), float(np.percentile(diffs, 97.5))],
               "share_positive": float(np.mean(diffs > 0))}

# (d) revisions
cur = read("RSAFS_current.csv")
gcur = yoy(cur)
t1 = fc[(fc.model == "naive") & (fc.horizon == 1)]["target_date"].map(pd.Timestamp)
revs = [(str(T.date()), float(gcur.loc[T]), float(actual.loc[T]), float(gcur.loc[T] - actual.loc[T])) for T in t1]
res["mean_abs_revision_h1_targets"] = float(np.mean([abs(r[3]) for r in revs]))

um = read("UMCSENT.csv")  # read for audit only; not used by the model
res["feasibility"] = "complete"
res["feasibility_reason"] = None
lag_txt = "; ".join(f"{a['origin']}->last {a['last_month'][:7]}" for a in audit if a['months_behind'] > 1)
res["data_issues"] = [
    {"issue": "Irregular RSAFS publication timing (normal: the month-end vintage ends one calendar month before V). Delayed origins (vintage ends two months before V): " + lag_txt + ". The 2026-04-30 vintage added two months (2026-02 and 2026-03) at once; from then on timing is normal.",
     "handling": "Targets defined relative to last(V) as required (T = last(V)+h), never relative to the calendar; for delayed origins the h=1 target is a month that would normally already be published. Consequently target month 2026-03 is never an h=1 target and 2026-05 is never an h=3 target."},
    {"issue": "Two target months (2026-02, 2026-03) were first released in the same month-end vintage (2026-04-30); 2025-09 first appeared only in the 2025-11-30 vintage (publication disruption).",
     "handling": "Actuals taken as given from first_release.csv (first month-end appearance); no adjustment."},
    {"issue": "Annual/benchmark revision of RSAFS between the 2026-08-31 and 2026-09-30 vintages lowered the level of the whole series by roughly 4% (e.g. 2025-08: 731700 in the 2026-07-31 vintage vs 700407 in the current file); first_release levels for 2026-08 onward are on the new basis.",
     "handling": "All comparisons use log YoY growth computed within a single vintage, which is largely invariant to proportional level rebasing; documented as affecting estimand (d)."},
    {"issue": "PAYEMS vintage 2026-02-28 contains the annual benchmark revision (level about 1 million lower than in the 2026-01-31 vintage).",
     "handling": "PAYEMS not used by the model; noted only."},
    {"issue": "UMCSENT has 210 empty values, all before 1978 (quarterly/irregular survey period).",
     "handling": "UMCSENT not used by the model; no imputation needed."},
    {"issue": "RSAFS 2026-07 first release (763602, old basis) vs current value 729915 (new basis): level comparisons across vintages are not meaningful.",
     "handling": "Revision measure (d) computed on growth rates only."},
]
res["n_input_rows"] = n_rows
res["n_analysed"] = int(len(fc[fc.model == "model"]))
res["random_seed"] = SEED
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
sens = {}
for W in [12, 24, 36, 60, 120]:
    for h in HORIZONS:
        errs = []
        for V in ORIGINS:
            L = pd.read_csv(os.path.join(DATA, f"vintages/RSAFS_vintage_{V}.csv"), parse_dates=[0]).set_index("observation_date").iloc[:, 0]
            last = L.index[-1]; T = add_months(last, h)
            mu = (100 * np.log(L).diff().dropna()).iloc[-W:].mean()
            errs.append(actual.loc[T] - (100 * np.log(L.loc[last] / L.loc[add_months(T, -12)]) + h * mu))
        sens[f"W{W}_h{h}_rmse"] = float(np.sqrt(np.mean(np.square(errs))))
with open(os.path.join(OUT, "supplement.json"), "w") as f:
    json.dump({"publication_audit": audit, "bootstrap": boot, "drift_window_sensitivity_posthoc": sens, "revisions_h1": revs,
               "ar1_note": "see forecasts.csv"}, f, indent=2)

fig, axes = plt.subplots(2, 1, figsize=(8, 7), sharex=True)
for ax, h in zip(axes, HORIZONS):
    sub = fc[fc.horizon == h]
    a = sub[sub.model == "naive"].sort_values("target_date")
    t = pd.to_datetime(a.target_date)
    ax.plot(t, a.actual_first_release, "k-o", label="first-release actual")
    for m, st in [("naive", "C0--s"), ("ar1", "C2:^"), ("model", "C3-d")]:
        s2 = sub[sub.model == m].sort_values("target_date")
        ax.plot(pd.to_datetime(s2.target_date), s2.forecast, st, label=m)
    ax.set_title(f"h = {h}"); ax.set_ylabel("YoY growth, %"); ax.grid(alpha=.3)
axes[0].legend(fontsize=8)
fig.suptitle("US retail & food services sales YoY growth: real-time forecasts vs first release")
fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)
print(json.dumps({k: v for k, v in res.items() if k not in ("data_issues", "model_description")}, indent=1))