"""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
Regenerates out/results.json, out/forecasts.csv, out/table1.csv, out/figure1.png from data/ only."""
import json, os
import numpy as np
import pandas as pd
from scipy.stats import norm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

D, O = "data", "out"
os.makedirs(O, exist_ok=True)
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]
SEED = 20261030
n_rows = {}


def read(rel):
    df = pd.read_csv(os.path.join(D, rel), parse_dates=["observation_date"]).set_index("observation_date")
    n_rows[rel] = len(df)
    s = df.iloc[:, 0].astype(float)
    assert s.index.is_monotonic_increasing and not s.index.duplicated().any()
    return s


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


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

data_issues = []
rows = []
last_months = {}
for V in ORIGINS:
    L = read(f"vintages/RSAFS_vintage_{V}.csv")
    P = read(f"vintages/PAYEMS_vintage_{V}.csv")  # read for audit only; not used by the chosen model
    assert L.isna().sum() == 0 and (L.index.to_series().diff().dt.days.dropna() <= 31).all()
    g = yoy(L)
    last = L.index[-1]
    last_months[V] = last
    # AR(1) on g_m, 1993-02 <= m <= last
    y = g.loc["1993-02-01":last]
    x = g.shift(1).loc["1993-02-01":last]
    X = np.column_stack([np.ones(len(y)), x.values])
    a, b = np.linalg.lstsq(X, y.values, rcond=None)[0]
    lnL = np.log(L)
    mu = (lnL.loc[last] - lnL.loc[last - pd.DateOffset(months=12)]) / 12  # mean monthly log growth, last 12 months
    for h in HORIZONS:
        T = last + pd.DateOffset(months=h)
        f_naive = g.loc[last]
        f_snaive = g.loc[T - pd.DateOffset(months=12)]
        f = g.loc[last]
        for _ in range(h):
            f = a + b * f
        f_ar1 = f
        # Model: base-effect-aware drift. Forecast the level by extrapolating the last published level with the
        # trailing-12-month mean monthly log growth, then divide by the (already published) base-period level.
        f_model = 100 * (lnL.loc[last] + h * mu - lnL.loc[T - pd.DateOffset(months=12)])
        for m, fc in [("naive", f_naive), ("snaive", f_snaive), ("ar1", f_ar1), ("model", f_model)]:
            rows.append(dict(origin_date=V, target_date=T.strftime("%Y-%m-%d"), horizon=h, forecast=round(float(fc), 6),
                             actual_first_release=float(actual.loc[T]), model=m, vintage_used=V))

fc = pd.DataFrame(rows)
fc.to_csv(os.path.join(O, "forecasts.csv"), index=False)
fc["err"] = fc.actual_first_release - fc.forecast

# publication-timing audit
for V in ORIGINS:
    lag = (pd.Timestamp(V).year - last_months[V].year) * 12 + pd.Timestamp(V).month - last_months[V].month
    if lag != 1:
        data_issues.append({"issue": f"Irregular publication timing: RSAFS vintage {V} ends at {last_months[V]:%Y-%m} ({lag} months before the origin month).",
                            "handling": "Per task rules the target is defined relative to last(V), the last month present in the vintage, not to the calendar; forecasts use only that vintage."})
lm = pd.Series(last_months)
added = (lm.dt.year * 12 + lm.dt.month).diff()
for V, k in added.items():
    if k == 2:
        data_issues.append({"issue": f"Vintage {V} adds two new months at once (catch-up after delayed release); the month {lm[V] - pd.DateOffset(months=1):%Y-%m} is never an h=1 target and its first release is in the same vintage as the following month.",
                            "handling": "No adjustment; targets follow last(V)+h. Horizon-1 targets therefore skip one month and h=3 targets skip one month, so the 10 origins are not equally spaced in target time."})

table = []
for m in ["naive", "snaive", "ar1", "model"]:
    for h in HORIZONS:
        e = fc[(fc.model == m) & (fc.horizon == h)].err.values
        table.append(dict(model=m, horizon=h, rmse=float(np.sqrt(np.mean(e ** 2))), mae=float(np.mean(np.abs(e))), n=len(e)))
tab = pd.DataFrame(table)
tab.to_csv(os.path.join(O, "table1.csv"), index=False)


def nw_lrv(d, L):
    d = d - d.mean()
    n = len(d)
    v = np.sum(d * d) / n
    for l in range(1, L + 1):
        v += 2 * (1 - l / (L + 1)) * np.sum(d[l:] * d[:-l]) / n
    return v


res = dict(task_id="T1-forecast",
           model_name="Base-effect drift (trailing-12m mean growth level extrapolation)",
           model_description=(
               "At origin V, using only the RSAFS vintage dated V: let last = last month in the vintage and "
               "mu = (ln L_last - ln L_{last-12})/12, the mean monthly log growth over the last 12 published months. "
               "The level of target month T = last+h is forecast as ln L_T = ln L_last + h*mu, and the YoY forecast is "
               "100*(ln L_last + h*mu - ln L_{T-12}), where L_{T-12} is the base-period level already published in the same vintage. "
               "Equivalently: naive forecast + h*g_last/12 - 100*(ln L_{T-12} - ln L_{last-12}), i.e. it corrects the naive forecast "
               "for known base effects (the months dropping out of the 12-month window). No parameters are estimated; "
               "PAYEMS and UMCSENT are not used (specification fixed a priori, before any evaluation against actuals). "
               "All inputs are from vintage V, so the availability rule holds trivially."),
           n_origins=len(ORIGINS), first_origin=ORIGINS[0], last_origin=ORIGINS[-1])
for r in table:
    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").err.values
    em = fc[(fc.model == "model") & (fc.horizon == h)].sort_values("origin_date").err.values
    res[f"rel_improvement_rmse_h{h}"] = 1 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
    d = en ** 2 - em ** 2
    n = len(d)
    lrv = nw_lrv(d, h - 1)
    dm = d.mean() / np.sqrt(lrv / n)
    p = 2 * (1 - norm.cdf(abs(dm)))
    res[f"dm_stat_h{h}"] = float(dm)
    res[f"dm_pvalue_h{h}"] = float(p)
    better = "model" if res[f"rmse_model_h{h}"] < res[f"rmse_naive_h{h}"] else "naive"
    res[f"conclusion_code_h{h}"] = f"{better}_better_{'significant' if p < 0.05 else 'not_significant'}"
    # supplementary: HLN-adjusted DM with t(n-1), and moving-block bootstrap (block 3) of RMSE_naive - RMSE_model
    k = np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
    hln = dm * k
    B, bl = 10000, 3
    diffs = []
    for _ in range(B):
        starts = rng.integers(0, n - bl + 1, size=int(np.ceil(n / bl)))
        idx = np.concatenate([np.arange(s, s + bl) for s in starts])[:n]
        diffs.append(np.sqrt(np.mean(en[idx] ** 2)) - np.sqrt(np.mean(em[idx] ** 2)))
    diffs = np.array(diffs)
    from scipy.stats import t as tdist
    boot[h] = dict(hln_stat=float(hln), hln_p=float(2 * (1 - tdist.cdf(abs(hln), n - 1))),
                   rmse_diff=float(res[f"rmse_naive_h{h}"] - res[f"rmse_model_h{h}"]),
                   boot_ci95=[float(np.percentile(diffs, 2.5)), float(np.percentile(diffs, 97.5))],
                   wins_model=int(np.sum(np.abs(em) < np.abs(en))), n=n)

# (d) revision of h=1 targets
cur = read("RSAFS_current.csv")
um = read("UMCSENT.csv")  # audited only (not used by the model)
n_um_missing = int(um.isna().sum())
gcur = yoy(cur)
t1 = pd.to_datetime(fc[(fc.horizon == 1) & (fc.model == "naive")].target_date)
revs = np.abs(gcur.loc[t1].values - actual.loc[t1].values)
res["mean_abs_revision_h1_targets"] = float(np.mean(revs))
lvl_rev = 100 * (cur.loc[fr.target_month].values / fr.level_first_release.values - 1)
data_issues.append({"issue": f"Large level revision: current RSAFS levels for 2024-06..2026-07 are {-lvl_rev[:-2].max():.1f}% to {-lvl_rev[:-2].min():.1f}% below their first-release levels (a benchmark/annual revision that first appears in the 2026-09-30 month-end vintage shifted the whole level; the 2026-08 and 2026-09 first releases are already on the new level); YoY growth is revised by up to {np.max(np.abs(gcur.loc[fr.target_month].values - fr.yoy_first_release.values)):.2f} pp.",
                    "handling": "Forecasts are evaluated against first-release YoY (computed within one vintage, so level shifts cancel to first order). The current file is used only for estimand (d)."})
data_issues.append({"issue": "Real-time vintages differ from each other through routine revisions of the last 2-3 months and through annual revisions; values of a given month differ across origins.",
                    "handling": "Each origin uses only its own vintage file; no splicing of vintages."})
data_issues.append({"issue": f"UMCSENT has empty values for {n_um_missing} months before 1978 (quarterly survey); PAYEMS vintages also show irregular timing (2025-10-31 ends at 2025-08; 2025-12-31 adds two months at once) due to delayed BLS releases.",
                    "handling": "Neither predictor enters the chosen model; files were read, audited and row-counted only (PAYEMS, UMCSENT)."})
res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = data_issues
res["n_input_rows"] = n_rows
res["n_analysed"] = int(len(fc[fc.model == "model"]))
res["random_seed"] = SEED
res["supplementary"] = {"per_horizon": {f"h{h}": boot[h] for h in HORIZONS},
                        "abs_revision_h1_targets": {t.strftime("%Y-%m"): float(r) for t, r in zip(t1, revs)},
                        "ar1_note": "AR(1) re-estimated at every origin on that vintage only."}
with open(os.path.join(O, "results.json"), "w") as f:
    json.dump(res, f, indent=2)

# figure
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), sharey=True)
for ax, h in zip(axes, HORIZONS):
    sub = fc[fc.horizon == h]
    a = actual.loc["2025-06-01":"2026-09-01"]
    ax.plot(a.index, a.values, "k-o", ms=4, label="first-release actual")
    for m, st in [("naive", "s--"), ("model", "^-"), ("ar1", "x:")]:
        s = sub[sub.model == m]
        ax.plot(pd.to_datetime(s.target_date), s.forecast, st, label=m)
    ax.set_title(f"h = {h}")
    ax.set_ylabel("YoY growth of RSAFS, %")
    ax.grid(alpha=.3)
    ax.tick_params(axis="x", rotation=45)
axes[0].legend()
fig.suptitle("Real-time forecasts of US retail sales YoY growth vs first release")
fig.tight_layout()
fig.savefig(os.path.join(O, "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))