"""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
Reads only data/; writes out/results.json, out/forecasts.csv, out/table1.csv, out/figure1.png.
Deterministic (no randomness except the seeded bootstrap below)."""
import json, os
import numpy as np
import pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATA, OUT = "data", "out"
SEED = 20261030
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]
MODELS = ["naive", "snaive", "ar1", "model"]
n_input_rows = {}


def read_series(path, key):
    df = pd.read_csv(path)
    n_input_rows[key] = int(len(df))
    s = df.set_index(pd.PeriodIndex(pd.to_datetime(df.iloc[:, 0]), freq="M")).iloc[:, 1].astype(float)
    return s


def yoy(level):
    lg = np.log(level)
    return 100.0 * (lg - lg.shift(12, freq="M").reindex(lg.index))


def nw_dm(d, h):
    """DM stat with Bartlett long-run variance, truncation lag h-1, no HLN factor."""
    d = np.asarray(d, float); n = len(d); dbar = d.mean(); u = d - dbar
    lrv = np.sum(u * u) / n
    L = h - 1
    for k in range(1, L + 1):
        w = 1.0 - k / (L + 1.0)
        lrv += 2.0 * w * np.sum(u[k:] * u[:-k]) / n
    stat = dbar / np.sqrt(lrv / n)
    p = 2.0 * (1.0 - stats.norm.cdf(abs(stat)))
    return float(stat), float(p), float(lrv)


# ---------------------------------------------------------------- inputs
fr = pd.read_csv(os.path.join(DATA, "first_release.csv"))
n_input_rows["first_release.csv"] = int(len(fr))
actual = pd.Series(fr["yoy_first_release"].values,
                   index=pd.PeriodIndex(pd.to_datetime(fr["target_month"]), freq="M"))
rs_vint = {V: read_series(os.path.join(DATA, "vintages", f"RSAFS_vintage_{V}.csv"),
                          f"vintages/RSAFS_vintage_{V}.csv") for V in ORIGINS}
# PAYEMS / UMCSENT are read only for the audit row counts; they are not inputs to any forecast.
for V in ORIGINS:
    read_series(os.path.join(DATA, "vintages", f"PAYEMS_vintage_{V}.csv"), f"vintages/PAYEMS_vintage_{V}.csv")
read_series(os.path.join(DATA, "UMCSENT.csv"), "UMCSENT.csv")
rs_current = read_series(os.path.join(DATA, "RSAFS_current.csv"), "RSAFS_current.csv")  # estimand (d) only

# ---------------------------------------------------------------- forecasts
rows, ar_info, lastmonths = [], {}, {}
for V in ORIGINS:
    L = rs_vint[V]
    assert L.index.max() < pd.Period(V, "M")  # last published month precedes the origin month
    g = yoy(L).dropna()
    last = L.index.max(); lastmonths[V] = str(last)
    lnL = np.log(L)
    # AR(1): g_m on const, g_{m-1}, 1993-02 <= m <= last
    y = g.loc[pd.Period("1993-02", "M"):last]
    x = g.shift(1, freq="M").reindex(y.index)
    X = np.column_stack([np.ones(len(y)), x.values])
    c, phi = np.linalg.lstsq(X, y.values, rcond=None)[0]
    ar_info[V] = {"const": float(c), "phi": float(phi), "n": int(len(y))}
    for h in HORIZONS:
        T = last + h
        fc = {}
        fc["naive"] = g[last]
        fc["snaive"] = g[T - 12]
        f = g[last]
        for _ in range(h):
            f = c + phi * f
        fc["ar1"] = f
        # Own model: base-effect random walk with drift in log level.
        # ln L_T forecast = ln L_last + h * drift, drift = g_last/12 (avg monthly log growth over last 12 months);
        # YoY forecast = 100*(ln L_T_hat - ln L_{T-12}), with L_{T-12} known in the same vintage.
        drift = g[last] / 100.0 / 12.0
        fc["model"] = 100.0 * (lnL[last] + h * drift - lnL[T - 12])
        for m in MODELS:
            rows.append({"origin_date": V, "target_date": T.to_timestamp().strftime("%Y-%m-%d"),
                         "horizon": h, "forecast": float(fc[m]),
                         "actual_first_release": float(actual[T]) if T in actual.index else np.nan,
                         "model": m, "vintage_used": V})
fc_df = pd.DataFrame(rows)
fc_df["error"] = fc_df["actual_first_release"] - fc_df["forecast"]
fc_df[["origin_date", "target_date", "horizon", "forecast", "actual_first_release", "model", "vintage_used"]] \
    .to_csv(os.path.join(OUT, "forecasts.csv"), index=False, float_format="%.6f")

# ---------------------------------------------------------------- evaluation
res = {"task_id": "T1-forecast",
       "model_name": "BE-RWD (base-effect random walk with drift in log level)",
       "model_description": (
           "Univariate real-time model on the RSAFS vintage dated V only. The log level is forecast as a random walk "
           "with drift: ln L_hat(last+h) = ln L(last) + h*d, with d = g_last/1200, i.e. the average monthly log growth "
           "over the 12 months ending at last(V) in that vintage. The YoY forecast is 100*(ln L_hat(T) - ln L(T-12)), "
           "where L(T-12) is read from the same vintage, so the forecast equals the naive forecast plus a known "
           "base-effect correction (the change in the year-earlier comparison level) and drift. No parameters are "
           "estimated; PAYEMS and UMCSENT are not used. Specification fixed before looking at any evaluation outcome; "
           "no alternative models were tried on the evaluation sample."),
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
table = []
for m in MODELS:
    for h in HORIZONS:
        e = fc_df.query("model==@m and horizon==@h")["error"].values
        rmse, mae = float(np.sqrt(np.mean(e ** 2))), float(np.mean(np.abs(e)))
        res[f"rmse_{m}_h{h}"], res[f"mae_{m}_h{h}"] = rmse, mae
        table.append({"model": m, "horizon": h, "rmse": rmse, "mae": mae, "n": int(len(e))})
pd.DataFrame(table).to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")

extra = {}
rng = np.random.default_rng(SEED)
for h in HORIZONS:
    en = fc_df.query("model=='naive' and horizon==@h").sort_values("origin_date")["error"].values
    em = fc_df.query("model=='model' and horizon==@h").sort_values("origin_date")["error"].values
    res[f"rel_improvement_rmse_h{h}"] = 1.0 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
    d = en ** 2 - em ** 2
    s, p, lrv = nw_dm(d, h)
    res[f"dm_stat_h{h}"], res[f"dm_pvalue_h{h}"] = s, 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: moving-block bootstrap (block length max(h,2)) of RMSE_naive - RMSE_model
    B, bl, n = 5000, max(h, 2), len(d)
    diffs = []
    for _ in range(B):
        starts = rng.integers(0, n - bl + 1, size=int(np.ceil(n / bl)))
        idx = np.concatenate([np.arange(s0, s0 + bl) for s0 in starts])[:n]
        diffs.append(np.sqrt(np.mean(en[idx] ** 2)) - np.sqrt(np.mean(em[idx] ** 2)))
    lo, hi = np.percentile(diffs, [2.5, 97.5])
    res[f"boot_rmse_diff_ci95_lo_h{h}"], res[f"boot_rmse_diff_ci95_hi_h{h}"] = float(lo), float(hi)
    res[f"share_origins_model_smaller_sq_error_h{h}"] = float(np.mean(d > 0))
    extra[f"h{h}"] = {"dm_lrv": lrv, "mean_loss_diff": float(d.mean()),
                      "rmse_diff_naive_minus_model": res[f"rmse_naive_h{h}"] - res[f"rmse_model_h{h}"],
                      "boot_rmse_diff_ci95": [float(lo), float(hi)], "boot_B": B, "boot_block": bl,
                      "share_origins_model_smaller_sq_error": float(np.mean(d > 0))}

# (d) revisions: h=1 targets, current-vintage YoY vs first release
g_cur = yoy(rs_current)
t1 = fc_df.query("model=='naive' and horizon==1")
tm = pd.PeriodIndex(pd.to_datetime(t1["target_date"]), freq="M")
rev = np.abs(g_cur.reindex(tm).values - actual.reindex(tm).values)
res["mean_abs_revision_h1_targets"] = float(np.mean(rev))
extra["revision_by_target"] = {str(k): float(v) for k, v in zip(tm, g_cur.reindex(tm).values - actual.reindex(tm).values)}
extra["ar1_params"] = ar_info
extra["last_month_by_origin"] = lastmonths

res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = [
    {"issue": "Irregular RSAFS publication timing: last(V)=2025-08 at V=2025-10-31 and 2026-01 at V=2026-03-31 (two-month lag); V=2026-04-30 adds two months at once.",
     "handling": "Targets defined as last(V)+h per task rule; the 2025-10-31 and 2026-03-31 origins therefore forecast months that are 2 and 4 months past the last calendar month."},
    {"issue": "h=3 forecast errors of consecutive origins overlap (and some h=1 and h=3 target months coincide), inducing serial correlation in the loss differentials; only 10 origins are available.",
     "handling": "DM long-run variance uses Bartlett kernel with truncation lag h-1 as specified."},
    {"issue": "Large level revision of RSAFS after the 2026-07-31 vintage (current levels ~4-5% below vintage levels); first_release.csv shows the break between targets 2026-07 and 2026-08.",
     "handling": "YoY always computed within one vintage; h=3 actuals for targets 2026-08/2026-09 are post-revision first releases, used as given; estimand (d) reflects this revision."},
    {"issue": "PAYEMS benchmark revision visible in the 2026-02-28 vintage; PAYEMS 2025-10 first appears in the 2025-12-31 vintage.",
     "handling": "PAYEMS not used in any forecast; documented only."},
    {"issue": "UMCSENT has 210 empty values before 1978-01.",
     "handling": "UMCSENT not used in any forecast; documented only."},
]
res["n_input_rows"] = dict(sorted(n_input_rows.items()))
res["n_analysed"] = int(fc_df.query("model=='model'")["error"].notna().sum())
res["random_seed"] = SEED

with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=1)
with open(os.path.join(OUT, "supplementary.json"), "w") as fh:
    json.dump(extra, fh, indent=1)

# ---------------------------------------------------------------- figure
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5), sharey=True)
for ax, h in zip(axes, HORIZONS):
    sub = fc_df.query("horizon==@h")
    for m, sty in [("naive", "s--"), ("model", "o-"), ("ar1", "^:")]:
        s = sub.query("model==@m").sort_values("target_date")
        ax.plot(pd.to_datetime(s["target_date"]), s["forecast"], sty, label=m, alpha=.8)
    s = sub.query("model=='naive'").sort_values("target_date")
    ax.plot(pd.to_datetime(s["target_date"]), s["actual_first_release"], "k-", lw=2.2, label="first release")
    ax.set_title(f"h = {h}: RSAFS YoY growth (%), real-time forecasts")
    ax.set_xlabel("target month"); ax.grid(alpha=.3); ax.tick_params(axis="x", rotation=45)
axes[0].set_ylabel("percent"); axes[0].legend()
fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)

print(pd.DataFrame(table).to_string(index=False))
for h in HORIZONS:
    print(f"h={h}: rel_impr={res[f'rel_improvement_rmse_h{h}']:.4f} DM={res[f'dm_stat_h{h}']:.4f} p={res[f'dm_pvalue_h{h}']:.4f} {res[f'conclusion_code_h{h}']}")
print("mean_abs_revision_h1_targets", round(res["mean_abs_revision_h1_targets"], 4))