"""T1-forecast: real-time forecasting of US retail sales YoY growth (RSAFS).
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
(and the auxiliary diagnostics file out/state/aux_diagnostics.json).
Deterministic: no random numbers are used (random_seed = null)."""
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"
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"]
MODEL_NAME = "rw_drift_base_effect"
MODEL_DESCRIPTION = (
    "Random walk with drift for the log level of RSAFS, translated into YoY growth so that the known "
    "'base effects' (the monthly changes that drop out of the 12-month window) are used. At origin V, "
    "with x_m = ln L_m from the RSAFS vintage dated V and l = last(V): drift d = (x_l - x_{l-12})/12 "
    "(mean monthly log change over the last 12 published months, i.e. a rolling 12-month estimation window); "
    "forecast of the log level x_{l+h} = x_l + h*d; forecast YoY g_{l+h} = 100*(x_l + h*d - x_{l+h-12}), "
    "where x_{l+h-12} (h<=3) is already published in vintage V. Uses only the RSAFS vintage dated V; "
    "PAYEMS and UMCSENT are not used. The specification (including the 12-month drift window) was fixed "
    "before any forecast error was computed; no parameter was tuned on the evaluation targets.")


def read_series(path):
    d = pd.read_csv(path)
    s = pd.Series(d.iloc[:, 1].values.astype(float), index=pd.to_datetime(d.iloc[:, 0]))
    return s, len(d)


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


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

rows = []
last_months = {}
for V in ORIGINS:
    lv, n = read_series(os.path.join(DATA, "vintages", f"RSAFS_vintage_{V}.csv"))
    n_input_rows[f"vintages/RSAFS_vintage_{V}.csv"] = n
    # PAYEMS vintage is read only for the row count/audit; not used in any forecast.
    _, npay = read_series(os.path.join(DATA, "vintages", f"PAYEMS_vintage_{V}.csv"))
    n_input_rows[f"vintages/PAYEMS_vintage_{V}.csv"] = npay
    assert lv.index.max() <= pd.Timestamp(V), "vintage contains months after origin"
    assert not lv.isna().any()
    last = lv.index.max()
    last_months[V] = last
    g = yoy(lv)
    x = np.log(lv)
    # AR(1): OLS of g_m on const, g_{m-1}, 1993-02 <= m <= last
    gs = g.loc["1993-01-01":last]
    y = gs.loc["1993-02-01":].values
    ylag = gs.shift(1).loc["1993-02-01":].values
    X = np.column_stack([np.ones_like(ylag), ylag])
    a, b = np.linalg.lstsq(X, y, rcond=None)[0]
    d = (x[last] - x[last - pd.DateOffset(months=12)]) / 12.0
    for h in HORIZONS:
        T = last + pd.DateOffset(months=h)
        T12 = T - pd.DateOffset(months=12)
        f = {}
        f["naive"] = g[last]
        f["snaive"] = g[T12]
        fa = g[last]
        for _ in range(h):
            fa = a + b * fa
        f["ar1"] = fa
        f["model"] = 100.0 * (x[last] + h * d - x[T12])
        for m in MODELS:
            rows.append(dict(origin_date=V, target_date=T.strftime("%Y-%m-%d"), horizon=h,
                             forecast=float(f[m]), actual_first_release=float(actual[T]),
                             model=m, vintage_used=V))

fc = pd.DataFrame(rows)
fc = fc.sort_values(["horizon", "model", "origin_date"]).reset_index(drop=True)
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")

res = {"task_id": "T1-forecast", "model_name": MODEL_NAME, "model_description": MODEL_DESCRIPTION,
       "n_origins": len(ORIGINS), "first_origin": ORIGINS[0], "last_origin": ORIGINS[-1]}
tab = []
for m in MODELS:
    for h in HORIZONS:
        e = fc[(fc.model == m) & (fc.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
        tab.append(dict(model=m, horizon=h, rmse=rmse, mae=mae, n=len(e)))
pd.DataFrame(tab).to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")


def dm_test(e_naive, e_model, h):
    dd = e_naive ** 2 - e_model ** 2
    n = len(dd)
    dbar = dd.mean()
    dc = dd - dbar
    lrv = np.sum(dc * dc) / n
    for j in range(1, h):  # Bartlett weights 1 - j/h, truncation lag h-1
        w = 1.0 - j / h
        lrv += 2.0 * w * np.sum(dc[j:] * dc[:-j]) / n
    stat = dbar / np.sqrt(lrv / n)
    p = 2.0 * (1.0 - stats.norm.cdf(abs(stat)))
    return float(stat), float(p), float(lrv), float(dbar)


dm_extra = {}
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}"] = 1.0 - res[f"rmse_model_h{h}"] / res[f"rmse_naive_h{h}"]
    s, p, lrv, dbar = dm_test(en, em, h)
    assert lrv > 0
    res[f"dm_stat_h{h}"], res[f"dm_pvalue_h{h}"] = s, p
    nn = len(en)
    hln = np.sqrt((nn + 1 - 2 * h + h * (h - 1) / nn) / nn)
    dm_extra[h] = (lrv, dbar, s * hln, float(2 * (1 - stats.t.cdf(abs(s * hln), df=nn - 1))),
                   [float(v) for v in (en ** 2 - em ** 2)], int(np.sum(em ** 2 < en ** 2)))
    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}"

# (d) revision: current-vintage growth vs first release, h=1 targets (RSAFS_current used ONLY here)
cur, ncur = read_series(os.path.join(DATA, "RSAFS_current.csv"))
n_input_rows["RSAFS_current.csv"] = ncur
gcur = yoy(cur)
t1 = fc[(fc.model == "naive") & (fc.horizon == 1)]["target_date"].map(pd.Timestamp)
rev = pd.DataFrame({"target": t1.values, "g_current": gcur.reindex(t1.values).values,
                    "first_release": actual.reindex(t1.values).values})
rev["abs_rev"] = (rev.g_current - rev.first_release).abs()
res["mean_abs_revision_h1_targets"] = float(rev["abs_rev"].mean())

res["feasibility"] = "complete"
res["feasibility_reason"] = None
lag = {V: (pd.Timestamp(V).to_period("M") - last_months[V].to_period("M")).n for V in ORIGINS}
res["data_issues"] = [
    {"issue": "Irregular RSAFS publication timing: last(V) is 2 months before the origin month for origins "
              + ", ".join(V for V in ORIGINS if lag[V] == 2)
              + " and 1 month before it for the others; the 2026-04-30 vintage added two months at once.",
     "handling": "Targets defined as last(V)+h exactly as TASK.md requires; each forecast uses only the vintage dated V (vintage_used = origin_date)."},
    {"issue": "First-release dating is at month-end granularity: 2025-09 first appears in the 2025-11-30 vintage, and 2026-02 and 2026-03 both first appear in the 2026-04-30 vintage (so the 2026-02 'first release' need not be the advance estimate).",
     "handling": "first_release.csv used as provided (task definition of the actual); documented in the report."},
    {"issue": "Large level revision: RSAFS levels in the current vintage are about 4.7-4.9% below those of the 2026-07-31 vintage for 2025-2026 months (first_release rows for 2026-08/09 already use the revised level base).",
     "handling": "YoY growth is always computed within a single vintage (levels never mixed across vintages); the revision enters only estimand (d)."},
    {"issue": "UMCSENT has 210 empty values, all before 1978 (quarterly survey era).",
     "handling": "UMCSENT is not used by the model; no imputation needed."},
    {"issue": "RSAFS_current.csv contains information not available at any forecast origin.",
     "handling": "Read only in the revision block for estimand (d), after all forecasts are produced."},
]
n_input_rows["UMCSENT.csv"] = len(pd.read_csv(os.path.join(DATA, "UMCSENT.csv")))  # audited, not used in forecasts
res["n_input_rows"] = dict(sorted(n_input_rows.items()))
res["n_analysed"] = int(len(fc[fc.model == "model"]))
res["random_seed"] = None

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

# auxiliary (non-required) diagnostics cited in the report
aux = {"dm_long_run_variance": {f"h{h}": dm_extra[h][0] for h in HORIZONS},
       "dm_mean_loss_diff": {f"h{h}": dm_extra[h][1] for h in HORIZONS},
       "sensitivity_dm_hln_stat": {f"h{h}": dm_extra[h][2] for h in HORIZONS},
       "sensitivity_dm_hln_pvalue_t_nminus1": {f"h{h}": dm_extra[h][3] for h in HORIZONS},
       "loss_diff_by_origin": {f"h{h}": dm_extra[h][4] for h in HORIZONS},
       "n_origins_model_sq_error_below_naive": {f"h{h}": dm_extra[h][5] for h in HORIZONS},
       "revision_by_target": rev.assign(target=rev.target.dt.strftime("%Y-%m")).round(6).to_dict("records"),
       "last_month_by_origin": {V: last_months[V].strftime("%Y-%m") for V in ORIGINS}}
os.makedirs(os.path.join(OUT, "state"), exist_ok=True)
with open(os.path.join(OUT, "state", "aux_diagnostics.json"), "w") as fh:
    json.dump(aux, fh, indent=2)

# Figure
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")
    ax.plot(pd.to_datetime(a_.target_date), a_.actual_first_release, "k-o", label="first release (actual)")
    for m, st in [("naive", "C0--s"), ("model", "C3-^"), ("ar1", "C2:x")]:
        s_ = sub[sub.model == m].sort_values("target_date")
        ax.plot(pd.to_datetime(s_.target_date), s_.forecast, st, label=m)
    ax.set_title(f"h = {h}: RSAFS YoY growth (%), by target month")
    ax.set_ylabel("percent"); ax.grid(alpha=.3)
axes[0].legend(fontsize=8)
fig.tight_layout()
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=120)

print(pd.DataFrame(tab).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}']:.3f} "
          f"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"], 6))