"""T1-forecast: real-time forecasting of US retail-sales YoY growth from ALFRED vintages.

Run from the 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.
"""
import json
import os
import glob

import numpy as np
import pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATA = "data"
OUT = "out"
SEED = 20261030
B_BOOT = 9999
BLOCK = 3
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"]


def read_series(path):
    d = pd.read_csv(path, parse_dates=[0], index_col=0)
    s = d.iloc[:, 0].astype(float)
    s.index = s.index.to_period("M")
    return s


def n_rows(path):
    with open(path) as f:
        return sum(1 for _ in f) - 1


def yoy(level):
    return 100.0 * (np.log(level) - np.log(level.shift(12)))


def forecasts_for_origin(V):
    """All forecasts at origin V use ONLY the RSAFS vintage dated V."""
    L = read_series(os.path.join(DATA, "vintages", f"RSAFS_vintage_{V}.csv"))
    assert L.index.is_monotonic_increasing and L.index.is_unique
    last = L.index[-1]
    g = yoy(L)
    lnL = np.log(L)
    out = {}
    # AR(1) on g over 1993-02 <= m <= last (consecutive pairs; g starts 1993-01)
    gs = g.loc[pd.Period("1993-01", "M"):last]
    y = gs.loc[pd.Period("1993-02", "M"):].values
    x = gs.shift(1).loc[pd.Period("1993-02", "M"):].values
    X = np.column_stack([np.ones_like(x), x])
    c, phi = np.linalg.lstsq(X, y, rcond=None)[0]
    # model: base-effect drift. log level path extrapolated with trailing-12m mean monthly log change
    mu = (lnL.loc[last] - lnL.loc[last - 12]) / 12.0
    for h in HORIZONS:
        T = last + h
        f_naive = g.loc[last]
        f_snaive = g.loc[T - 12]
        f = g.loc[last]
        for _ in range(h):
            f = c + phi * f
        f_ar1 = f
        f_model = 100.0 * (lnL.loc[last] + h * mu - lnL.loc[T - 12])
        out[h] = dict(target=T, naive=f_naive, snaive=f_snaive, ar1=f_ar1, model=f_model)
    return last, out, (c, phi)


def dm_test(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 * dc) / n
    for k in range(1, h):  # truncation lag h-1, Bartlett weights 1 - k/h
        w = 1.0 - k / h
        lrv += 2.0 * w * np.sum(dc[k:] * dc[:-k]) / n
    stat = dbar / np.sqrt(lrv / n)
    p = 2.0 * (1.0 - stats.norm.cdf(abs(stat)))
    return float(stat), float(p), float(dbar), float(lrv)


def block_boot_rmse_diff(e_naive, e_model, rng):
    n = len(e_naive)
    nb = int(np.ceil(n / BLOCK))
    diffs = np.empty(B_BOOT)
    for b in range(B_BOOT):
        starts = rng.integers(0, n, nb)
        idx = np.concatenate([(s + np.arange(BLOCK)) % n for s in starts])[:n]
        diffs[b] = np.sqrt(np.mean(e_model[idx] ** 2)) - np.sqrt(np.mean(e_naive[idx] ** 2))
    return float(np.percentile(diffs, 2.5)), float(np.percentile(diffs, 97.5)), float(np.mean(diffs < 0))


def main():
    os.makedirs(OUT, exist_ok=True)
    fr = pd.read_csv(os.path.join(DATA, "first_release.csv"), parse_dates=["target_month"])
    fr["T"] = fr["target_month"].dt.to_period("M")
    actual = dict(zip(fr["T"], fr["yoy_first_release"]))

    rows = []
    lasts = {}
    arcoef = {}
    for V in ORIGINS:
        last, fc, cp = forecasts_for_origin(V)
        lasts[V] = last
        arcoef[V] = cp
        for h in HORIZONS:
            T = fc[h]["target"]
            if T not in actual:
                raise RuntimeError(f"missing first-release actual for {T}")
            for m in MODELS:
                rows.append(dict(origin_date=V, target_date=T.to_timestamp().strftime("%Y-%m-%d"),
                                 horizon=h, forecast=round(float(fc[h][m]), 10),
                                 actual_first_release=float(actual[T]), model=m, vintage_used=V))
    F = pd.DataFrame(rows)
    F["error"] = F["actual_first_release"] - F["forecast"]
    F.drop(columns="error").to_csv(os.path.join(OUT, "forecasts.csv"), index=False)

    res = {"task_id": "T1-forecast",
           "model_name": "base-effect drift (BED)",
           "model_description": (
               "Forecast of the log level of RSAFS at T = last(V)+h from the RSAFS vintage dated V only: "
               "ln L_hat_T = ln L_last + h*mu, where mu = (ln L_last - ln L_{last-12})/12 is the trailing-12-month "
               "mean monthly log change in the same vintage (i.e. the latest YoY growth spread evenly over months). "
               "The YoY forecast is g_hat_T = 100*(ln L_hat_T - ln L_{T-12}), where L_{T-12} (always <= last(V) since h<=3) "
               "is the base-period level as published in vintage V. Equivalently: naive forecast plus the known "
               "base effect (drop-out of the monthly log changes of months last-11..last-12+h) plus h months of drift. "
               "No parameters are estimated on the evaluation sample; the specification was fixed before computing any "
               "out-of-sample error and no variant was compared. PAYEMS and UMCSENT are not used; RSAFS_current and "
               "first_release.csv never enter any forecast."),
           "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
    t1 = []
    for m in MODELS:
        for h in HORIZONS:
            e = F[(F.model == m) & (F.horizon == h)]["error"].values
            rmse = float(np.sqrt(np.mean(e ** 2)))
            mae = float(np.mean(np.abs(e)))
            res[f"rmse_{m}_h{h}"] = rmse
            res[f"mae_{m}_h{h}"] = mae
            t1.append(dict(model=m, horizon=h, rmse=rmse, mae=mae, n=len(e)))
    pd.DataFrame(t1).to_csv(os.path.join(OUT, "table1.csv"), index=False)

    rng = np.random.default_rng(SEED)
    extra = {}
    for h in HORIZONS:
        en = F[(F.model == "naive") & (F.horizon == h)].sort_values("origin_date")["error"].values
        em = F[(F.model == "model") & (F.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}"]
        stat, p, dbar, lrv = dm_test(en, em, h)
        res[f"dm_stat_h{h}"] = stat
        res[f"dm_pvalue_h{h}"] = p
        better = "model" if res[f"rmse_model_h{h}"] < res[f"rmse_naive_h{h}"] else "naive"
        sig = "significant" if p < 0.05 else "not_significant"
        res[f"conclusion_code_h{h}"] = f"{better}_better_{sig}"
        lo, hi, share = block_boot_rmse_diff(en, em, rng)
        extra[f"dm_mean_loss_diff_h{h}"] = dbar
        extra[f"dm_lrv_h{h}"] = lrv
        extra[f"boot_rmse_diff_lo_h{h}"] = lo
        extra[f"boot_rmse_diff_hi_h{h}"] = hi
        extra[f"boot_share_model_better_h{h}"] = share
        extra[f"n_model_wins_h{h}"] = int(np.sum(em ** 2 < en ** 2))

    # estimand (d): revision of h=1 targets, current vintage vs first release
    cur = read_series(os.path.join(DATA, "RSAFS_current.csv"))
    gcur = yoy(cur)
    t_h1 = [pd.Period(t, "M") for t in F[(F.model == "naive") & (F.horizon == 1)]["target_date"].str[:7]]
    revs = np.array([abs(gcur.loc[T] - actual[T]) for T in t_h1])
    res["mean_abs_revision_h1_targets"] = float(revs.mean())
    extra["mean_signed_revision_h1_targets"] = float(np.mean([gcur.loc[T] - actual[T] for T in t_h1]))
    extra["max_abs_revision_h1_targets"] = float(revs.max())

    res["feasibility"] = "complete"
    res["feasibility_reason"] = None
    res["data_issues"] = [
        {"issue": "Irregular RSAFS publication timing: vintages 2025-10-31 to 2026-03-31 end two months before the vintage date; the 2026-04-30 vintage adds 2026-02 and 2026-03 at once.",
         "handling": "Targets defined as last(V)+h from the last month actually present in each vintage, as TASK.md prescribes; target months are therefore unevenly spaced (h=1 skips 2026-03, h=3 skips 2026-05)."},
        {"issue": "Recent months are revised from vintage to vintage (e.g. 2026-05 level 763705 in the 2026-06-30 vintage, 766876 in 2026-07-31).",
         "handling": "Each forecast uses the values as published in its own vintage; actuals are first releases from first_release.csv."},
        {"issue": "Large downward level revision (about 3.5-4%) between the 2026-08-31 and 2026-09-30 vintages, visible in first_release.csv (2026-08 base-year level 700218 vs 731700 in shipped vintages) and in RSAFS_current.",
         "handling": "YoY growth is always computed within a single vintage so the level shift mostly cancels; not corrected. It affects the h=3 actuals of the last two origins (2026-08, 2026-09) and estimand (d)."},
        {"issue": "PAYEMS publication also delayed (2025-10-31 vintage ends 2025-08; 2025-12-31 adds two months).",
         "handling": "PAYEMS not used by the model; recorded only."},
        {"issue": "UMCSENT has 210 empty fields, all before 1978 (quarterly survey era); file is current vintage.",
         "handling": "UMCSENT not used by the model; recorded only."},
        {"issue": "RSAFS_current.csv and first_release.csv contain post-origin information.",
         "handling": "RSAFS_current used only for estimand (d); first_release only as evaluation actual."},
    ]
    files = sorted(glob.glob(os.path.join(DATA, "vintages", "*.csv"))) + [
        os.path.join(DATA, f) for f in ("RSAFS_current.csv", "UMCSENT.csv", "first_release.csv")]
    # PAYEMS and UMCSENT files are read only for the audit row counts (they do not enter forecasts)
    res["n_input_rows"] = {os.path.relpath(f, DATA): n_rows(f) for f in files}
    res["n_analysed"] = int(len(F[F.model == "model"]))
    res["random_seed"] = SEED
    res.update(extra)
    res["block_bootstrap"] = {"B": B_BOOT, "block_length": BLOCK, "type": "circular moving-block, paired errors"}
    res["last_month_by_origin"] = {V: str(lasts[V]) for V in ORIGINS}
    res["ar1_coef_by_origin"] = {V: {"const": float(arcoef[V][0]), "phi": float(arcoef[V][1])} for V in ORIGINS}
    res["abs_revision_by_h1_target"] = {str(T): float(r) for T, r in zip(t_h1, revs)}

    with open(os.path.join(OUT, "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 = F[F.horizon == h]
        a = sub[sub.model == "naive"].sort_values("target_date")
        x = pd.to_datetime(a["target_date"])
        ax.plot(x, a["actual_first_release"], "k-o", label="first-release actual")
        for m, st in (("naive", "b--s"), ("model", "r-.^")):
            s = sub[sub.model == m].sort_values("target_date")
            ax.plot(pd.to_datetime(s["target_date"]), s["forecast"], st,
                    label=f"{m} (RMSE {res[f'rmse_{m}_h{h}']:.2f})" if m == "naive" else
                    f"BED model (RMSE {res[f'rmse_{m}_h{h}']:.2f})")
        ax.set_title(f"h = {h}")
        ax.set_ylabel("YoY growth, % (100 ln)")
        ax.grid(alpha=0.3)
        ax.legend(fontsize=8)
        ax.tick_params(axis="x", rotation=45)
    fig.suptitle("RSAFS YoY growth: real-time forecasts vs first release (10 origins, 2025-10-31 to 2026-07-31)")
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)
    plt.close(fig)

    for h in HORIZONS:
        print(f"h={h}: " + ", ".join(f"{m} RMSE={res[f'rmse_{m}_h{h}']:.4f}" for m in MODELS)
              + f"; DM={res[f'dm_stat_h{h}']:.3f} p={res[f'dm_pvalue_h{h}']:.3f} -> {res[f'conclusion_code_h{h}']}")
    print(f"mean abs revision (h=1 targets) = {res['mean_abs_revision_h1_targets']:.4f}")


if __name__ == "__main__":
    main()