scripts/plot_regime_vs_range.py

"""Do the two live forecasts tell the same story? Regime vs price range, out of sample.

Both forecasts are about the size of the next day's move, not its direction:
  * the VAE+Transformer regime model gives P(next 24h more volatile than the
    trailing 90-day median), from 33 market features over the last week;
  * the price-range model gives a 24h 80% interval, from HAR volatility, the
    hour-of-week pattern and the Fed/CPI calendar.

They are built differently, so agreement is evidence that both track the same
underlying state. Both are the live versions with Deribit DVOL inputs. Period: 5 April -
25 September 2026, inside the regime model's hold-out (trained on data before
4 April 2026), with the range model refitted quarterly on the two years
before each quarter.

Usage:
    python scripts/plot_regime_vs_range.py
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "scripts"))
from app.forecasting.price_range import fit, forecast, hourly            # noqa: E402
from news_topic_event_study import GRID, INK, INK_2, SURFACE             # noqa: E402
from src.pipeline import VAETransformer                                   # noqa: E402
from src.volatility import add_volatility_target, build_volatility_features  # noqa: E402

# The live regime model (with DVOL) holds out everything from 4 Apr 2026 15:00.
START, END = pd.Timestamp("2026-04-05", tz="UTC"), pd.Timestamp("2026-09-25", tz="UTC")
VOL = ROOT / "btc_deployment_models" / "volatility"
BLUE, ORANGE, RED = "#2a78d6", "#eb6834", "#c24a12"


def regime_scores(frame: pd.DataFrame) -> pd.Series:
    meta = json.loads((VOL / "volatility_btc_24h_vaetransformer.json").read_text())
    ck = torch.load(VOL / meta["weights"], map_location="cpu", weights_only=False)
    models = []
    for state in ck["state_dicts"]:
        m = VAETransformer(n_market=int(ck["n_market"]), n_sent=0, latent=int(ck.get("latent", 24)),
                           d_model=int(ck["d_model"]), nhead=int(ck.get("n_heads", 4)),
                           num_layers=int(ck.get("n_layers", 2)), dropout=float(ck.get("dropout", 0.15)))
        m.load_state_dict(state, strict=True)
        m.eval()
        models.append(m)
    X = (frame[meta["features"]].to_numpy(float) - np.array(meta["scaler_mean"])) / np.array(meta["scaler_scale"])
    seq = int(meta["sequence_length"])
    rows = np.where((frame["timestamp"] >= START) & (frame["timestamp"] <= END))[0]
    rows = rows[rows >= seq - 1]
    # contiguous, finite windows only, as the live predictor requires
    stamps = frame["timestamp"].to_numpy()
    ok = [(stamps[i] - stamps[i - seq + 1]) == np.timedelta64(seq - 1, "h")
          and np.isfinite(X[i - seq + 1:i + 1]).all() for i in rows]
    rows = rows[np.array(ok)]
    out = []
    with torch.no_grad():
        for b in range(0, len(rows), 256):
            idx = rows[b:b + 256]
            win = torch.tensor(np.stack([X[i - seq + 1:i + 1] for i in idx]), dtype=torch.float32)
            out.append(np.mean([torch.sigmoid(m(win)[0]).numpy() for m in models], axis=0))
    return pd.Series(np.concatenate(out), index=frame["timestamp"].iloc[rows].to_numpy()), meta["decision_threshold"]


def range_forecasts(close: pd.Series, dvol: pd.Series) -> pd.DataFrame:
    """24h 80% range issued at every hour, refitted at each quarter start."""
    rows = []
    for q_start in pd.date_range(START, END, freq="QS", tz="UTC").union([START]):
        q_end = min(q_start + pd.offsets.QuarterBegin(1, startingMonth=1), END)
        params = fit(close.loc[:q_start], now=q_start, dvol=dvol.loc[:q_start])
        for t in close.loc[q_start:q_end].index:
            hist = close.loc[:t].dropna().iloc[-800:]
            row = forecast(hist, params, dvol=dvol)["horizons"][2]   # 24h; reads DVOL at t only
            lo, hi = row["ranges"]["80"]
            rows.append((t, lo, hi, float(hist.iloc[-1])))
    f = pd.DataFrame(rows, columns=["t", "lo", "hi", "price"]).drop_duplicates("t").set_index("t")
    f["width_pct"] = (f["hi"] / f["lo"] - 1) * 100
    return f


CACHE = ROOT / "outputs" / "regime_vs_range_hourly.csv"


def main() -> None:
    if "--redraw" in sys.argv and CACHE.exists():
        # Wording-only changes: redraw from the hours computed last time.
        both = pd.read_csv(CACHE, index_col=0)
        both.index = pd.to_datetime(both.index, utc=True)
        close = both.pop("close_series").dropna()
        both = both.dropna(subset=["p_high"])
        st = json.loads((ROOT / "outputs" / "regime_vs_range.json").read_text())
        draw(both, close, st["decision_threshold"], st["range_80_held"], st["regime_accuracy"],
             st["pearson_p_high_vs_width"], st["spearman_p_high_vs_width"],
             st["median_width_pct"]["predicted_high"], st["median_width_pct"]["predicted_low"],
             st["median_realised_abs_move_pct"]["predicted_high"],
             st["median_realised_abs_move_pct"]["predicted_low"])
        return
    btc = pd.read_csv(ROOT / "outputs" / "volatility" / "btc_1h.csv")
    btc["timestamp"] = pd.to_datetime(btc["timestamp"], utc=True)
    btc = btc.drop_duplicates("timestamp").sort_values("timestamp").reset_index(drop=True)
    dv = pd.read_csv(ROOT / "outputs" / "volatility" / "btc_dvol_1h.csv")
    dv["timestamp"] = pd.to_datetime(dv["timestamp"], utc=True)
    dvol = dv.set_index("timestamp")["close"]
    frame, _ = build_volatility_features(btc, dvol=dvol)
    frame = add_volatility_target(frame, 24)
    close = hourly(btc.set_index("timestamp")["close"])

    p_high, cut = regime_scores(frame)
    rng = range_forecasts(close, dvol)
    truth = frame.set_index("timestamp")["target_volatility_regime"]
    both = rng.join(p_high.rename("p_high"), how="inner").join(truth.rename("actual"), how="left")
    both["price_24h_later"] = close.shift(-24).reindex(both.index)
    both = both.dropna(subset=["p_high"])
    scored = both.dropna(subset=["price_24h_later"])
    inside = ((scored["price_24h_later"] >= scored["lo"]) & (scored["price_24h_later"] <= scored["hi"])).mean()
    pred_high = both["p_high"] >= cut
    corr = float(np.corrcoef(both["p_high"], both["width_pct"])[0, 1])
    rank_corr = float(both[["p_high", "width_pct"]].rank().corr().iloc[0, 1])
    w_hi, w_lo = both.loc[pred_high, "width_pct"].median(), both.loc[~pred_high, "width_pct"].median()
    a = both.dropna(subset=["actual"])
    acc = float(((a["p_high"] >= cut).astype(int) == a["actual"].astype(int)).mean())
    move = (np.log(scored["price_24h_later"] / scored["price"]).abs() * 100)
    mv_hi, mv_lo = move[scored["p_high"] >= cut].median(), move[scored["p_high"] < cut].median()
    print(f"{len(both):,} hours; range held {inside:.1%}; regime accuracy {acc:.1%} (cut {cut:.2f})")
    print(f"agreement: Pearson r {corr:.2f}, Spearman {rank_corr:.2f}; median 24h width "
          f"{w_hi:.2f}% when HIGH vs {w_lo:.2f}% when LOW; realised |24h move| {mv_hi:.2f}% vs {mv_lo:.2f}%")
    (ROOT / "outputs" / "regime_vs_range.json").write_text(json.dumps({
        "period": [str(START.date()), str(END.date())], "hours": int(len(both)),
        "range_80_held": float(inside), "regime_accuracy": acc, "decision_threshold": cut,
        "pearson_p_high_vs_width": corr, "spearman_p_high_vs_width": rank_corr,
        "median_width_pct": {"predicted_high": float(w_hi), "predicted_low": float(w_lo)},
        "median_realised_abs_move_pct": {"predicted_high": float(mv_hi), "predicted_low": float(mv_lo)}},
        indent=2))
    span = close.loc[both.index[0]:both.index[-1] + pd.Timedelta(hours=24)]
    both.join(span.rename("close_series"), how="outer").to_csv(CACHE)
    draw(both, close, cut, inside, acc, corr, rank_corr, w_hi, w_lo, mv_hi, mv_lo)


def draw(both, close, cut, inside, acc, corr, rank_corr, w_hi, w_lo, mv_hi, mv_lo) -> None:
    plt.rcParams.update({"font.family": "DejaVu Sans", "text.color": INK, "axes.labelcolor": INK_2})
    fig = plt.figure(figsize=(15, 10.5), facecolor=SURFACE)
    gs = fig.add_gridspec(3, 2, width_ratios=[3.3, 1], height_ratios=[1.5, 1, 1],
                          left=0.06, right=0.98, top=0.85, bottom=0.07, hspace=0.18, wspace=0.18)
    a = fig.add_subplot(gs[0, 0])
    b = fig.add_subplot(gs[1, 0], sharex=a)
    c = fig.add_subplot(gs[2, 0], sharex=a)
    s = fig.add_subplot(gs[:, 1])
    for ax in (a, b, c, s):
        ax.set_facecolor(SURFACE)
        for sp in ("top", "right"):
            ax.spines[sp].set_visible(False)
        for sp in ("left", "bottom"):
            ax.spines[sp].set_color(GRID)
        ax.tick_params(colors=INK_2, labelsize=8.5, length=0)
        ax.grid(color=GRID, linewidth=0.8, axis="y")
        ax.set_axisbelow(True)

    # actual HIGH hours, shaded on every time panel
    actual = both["actual"].fillna(0).astype(bool)
    edges = actual.astype(int).diff().fillna(0)
    starts = list(both.index[edges == 1]) + ([both.index[0]] if actual.iloc[0] else [])
    ends = list(both.index[edges == -1]) + ([both.index[-1]] if actual.iloc[-1] else [])
    for ax in (a, b, c):
        for s0, e0 in zip(sorted(starts), sorted(ends)):
            ax.axvspan(s0, e0, color="#fbe3d6", linewidth=0, zorder=0)

    # A: price and the 24h range, plotted at the hour it is ABOUT (issued 24h earlier)
    target = both.index + pd.Timedelta(hours=24)
    a.fill_between(target, both["lo"], both["hi"], color=BLUE, alpha=0.22, linewidth=0,
                   label="24h 80% range, forecast a day earlier")
    window = close.loc[both.index[0] + pd.Timedelta(hours=24):both.index[-1] + pd.Timedelta(hours=24)]
    a.plot(window.index, window, color=INK, linewidth=0.9, label="BTC price")
    a.set_ylabel("BTC / USDT", fontsize=8.5)
    a.yaxis.set_major_formatter(matplotlib.ticker.FuncFormatter(lambda v, _: f"${v / 1000:.0f}k"))
    a.legend(loc="upper left", frameon=False, fontsize=8.5)
    a.set_title(f"Prediction 1 · Where will the price be in 24 hours? The real price stayed inside "
                f"the 80% range {inside:.0%} of the time",
                loc="left", fontsize=11, fontweight="bold")

    # B: range width
    b.plot(both.index, both["width_pct"], color=BLUE, linewidth=1)
    b.set_ylabel("24h 80% range width, %", fontsize=8.5)
    b.set_title("Prediction 1 · How wide that 24-hour price range is", loc="left",
                fontsize=11, fontweight="bold")

    # C: regime probability
    c.plot(both.index, both["p_high"], color=ORANGE, linewidth=1)
    c.axhline(cut, color=INK, linewidth=1, linestyle=(0, (4, 2)))
    c.text(both.index[0], cut + 0.03, f" above {cut:.2f} = volatile", fontsize=8, color=INK_2)
    c.set_ylim(0, 1)
    c.set_ylabel("probability next 24h\nis more volatile", fontsize=8.5)
    c.set_title(f"Prediction 2 · Will the next 24 hours be more volatile than normal? "
                f"Correct {acc:.0%} of the hours", loc="left",
                fontsize=11, fontweight="bold")
    c.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    for ax in (a, b):
        ax.tick_params(labelbottom=False)

    # side: do they agree?
    hi_mask = both["p_high"] >= cut
    data = [both.loc[~hi_mask, "width_pct"], both.loc[hi_mask, "width_pct"]]
    bp = s.boxplot(data, positions=[0, 1], widths=0.55, patch_artist=True, showfliers=False,
                   medianprops={"color": INK, "linewidth": 1.6})
    for patch, col in zip(bp["boxes"], ("#c4dcf6", "#f6c9b1")):
        patch.set_facecolor(col)
        patch.set_edgecolor(INK_2)
    s.set_xticks([0, 1])
    s.set_xticklabels([f"Prediction 2 says\ncalm ({(~hi_mask).sum():,} h)",
                       f"Prediction 2 says\nvolatile ({hi_mask.sum():,} h)"], fontsize=9)
    s.set_ylabel("Prediction 1: 24h 80% range width, %", fontsize=8.5)
    s.set_title("Do the two predictions agree?", loc="left", fontsize=11, fontweight="bold")
    s.set_ylim(top=both["width_pct"].quantile(0.999) * 1.45)
    s.text(0.5, 0.98, f"median width {w_lo:.1f}% vs {w_hi:.1f}%\nrank correlation {rank_corr:.2f}\n\n"
                      f"real 24h move afterwards:\n{mv_lo:.1f}% vs {mv_hi:.1f}% (median)",
           transform=s.transAxes, ha="center", va="top", fontsize=9, color=INK)

    fig.suptitle("The two predictions move together: when a volatile day is expected, "
                 "the price range widens", x=0.06, ha="left", fontsize=14, fontweight="bold", y=0.975)
    fig.text(0.06, 0.905,
             "5 April – 25 September 2026, every hour, on data neither prediction was trained on.  "
             "Prediction 1: the 24-hour price range.  Prediction 2: calm or volatile next 24 hours.\n"
             "Pale orange: hours whose next 24 hours actually turned out more volatile than normal.",
             fontsize=9.2, color=INK_2, ha="left")
    fig.savefig(ROOT / "outputs" / "figures" / "fig19_regime_vs_range.png", dpi=150, facecolor=SURFACE)
    plt.close(fig)
    print("written: outputs/figures/fig19_regime_vs_range.png")


if __name__ == "__main__":
    main()