source_code/scripts/forecasting_hippo.py

"""
Forecasting HiPPO (online, statistics-only) + interpretability plot (3 histories x 2 horizons)

What this script does
---------------------
1) Builds THREE HiPPO systems per horizon:
   - S1: recent past memory on window length H (Legendre/HiPPO, ZOH discretized).
   - S2: "past-before-that" memory, driven by the lagged value reconstructed from S1.
   - S3: identical dynamics to S2, but driven by the true input x(t) (so S3 encodes the recent past
         in the same coordinates as S2, allowing us to learn the shift operator).
2) Learns a linear map T (RRR bottleneck d) ONLINE using only 1st/2nd order stats:
   - X := S2(t), Y := S3(t)  (predict recent-window state from preceding-window state)
3) Input signal: mixture of independent 1D RBF GP draws with different lengthscales (FFT/circulant approx).
   Weights are normalized so Var[x_t] ~= 1.
4) Runs TWO forecasters on the same x(t):
   - "short horizon" H_short
   - "long horizon"  H_long
5) After training, makes a 3x2 plot:
   rows = {true history, HiPPO reconstructed history, HiPPO predictive history}
   cols = {short horizon, long horizon}
   where the plotted history is the last 2H time units ([-2H, 0]):
     - True: actual x[t-2H+1 ... t]
     - HiPPO reconstructed: stitch (S2 recon on [-2H,-H]) + (S3 recon on [-H,0])
     - HiPPO predictive: stitch (S2 recon on [-2H,-H]) + (T_d S2 recon on [-H,0])

Notes
-----
- "Exact ZOH" here means we discretize continuous-time HiPPO with expm(A*dt) and the exact ZOH b_d.
- The RRR "metric in coefficient space" is Q = T^T T (or T^T G T with a non-identity metric).
  In this experiment we keep the future OP coefficients orthonormal under the implicit continuous uniform
  measure induced by HiPPO evaluation; the downstream interpretability uses T only to generate predictive
  histories.

"""

from dataclasses import dataclass
from typing import Sequence, Tuple, Dict

import jax
import jax.numpy as jnp
import jax.random as jr
from jax.scipy.linalg import expm
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

from mpm import get_system_params, get_output_vector


# ----------------------------
# Config
# ----------------------------

@dataclass
class Config:
    # HiPPO / forecasting
    measure: str = "legt"
    n: int = 64
    bottleneck_d: int = 12
    ridge: float = 1e-4

    # Two horizons (in "time units"; dt=1 step)
    horizon_short: float = 4.0
    horizon_long: float = 32.0
    horizon_23: float = 32.0 # shared history timescale for systems 2 and 3

    # Simulation
    T_train: int = 20000
    T_total: int = T_train + int(max(horizon_short, horizon_long)) + 10

    burnin: int = 0
    dt: float = 1.0

    # GP mixture (RBF)
    gp_lengthscales: Tuple[float, ...] = (0.1, 3.0, 16.0, 32.0, 64.0)
    gp_weights: Tuple[float, ...] | None = (0.5, 2, 5, 5, 5)
    gp_circulant_pad: int = 0  # if >0, uses length (T+pad) for FFT embedding

    # Plot
    plot_nt: int = 400
    out_png: str = "forecasting_hippo_interpretability.png"
    seed: int = 0


# ----------------------------
# GP mixture sampling (RBF) on integer grid via FFT/circulant approx
# ----------------------------

def _rbf_cov_first_row(n: int, ell: float) -> jnp.ndarray:
    """First row of Toeplitz covariance [C(0), C(1), ..., C(n-1)] for RBF on Z."""
    k = jnp.arange(n)
    return jnp.exp(-(k * k) / (2.0 * ell * ell))


def sample_rbf_gp_fft(key: jax.Array, T: int, ell: float, pad: int = 0) -> jnp.ndarray:
    """
    Approx sample from stationary GP on {0,...,T-1} with RBF covariance using circulant embedding.
    Returns approx zero-mean unit-variance series (variance ~ 1).
    """
    # Use embedding length N >= 2*(T+pad)+1 to fit [c0..c_{T+pad}, 0..0, c_{T+pad}..c1]
    N = int(2 * (T + pad) + 1)

    # Build Toeplitz first-row values up to lag (T+pad)
    c_toe = _rbf_cov_first_row(T + pad + 1, ell)  # length (T+pad+1): lags 0..T+pad

    # Construct length-N circulant row:
    # [c(0), c(1), ..., c(T+pad), 0, ..., 0, c(T+pad), ..., c(1)]
    left = c_toe                              # length L = T+pad+1
    right = c_toe[1:][::-1]                   # length L-1 = T+pad
    mid_len = N - (left.shape[0] + right.shape[0])  # should be >= 0 (often 0)
    mid = jnp.zeros((mid_len,), dtype=left.dtype)

    c_circ = jnp.concatenate([left, mid, right], axis=0)  # length N

    lam = jnp.real(jnp.fft.fft(c_circ))
    lam = jnp.maximum(lam, 0.0)

    z = jr.normal(key, (N,))
    zf = jnp.fft.fft(z)
    x = jnp.real(jnp.fft.ifft(jnp.sqrt(lam) * zf))
    x = x[:T]

    x = x - jnp.mean(x)
    x = x / (jnp.std(x) + 1e-8)
    return x


def sample_gp_mixture(key: jax.Array, cfg: Config) -> jnp.ndarray:
    """Mixture of independent RBF GPs with weights normalized so marginal variance ~ 1."""
    ells = list(cfg.gp_lengthscales)
    m = len(ells)
    if cfg.gp_weights is None:
        w = jnp.ones((m,))
    else:
        w = jnp.array(cfg.gp_weights, dtype=jnp.float32)
        if w.shape != (m,):
            raise ValueError("gp_weights must match gp_lengthscales length.")
    # Normalize so sum_i w_i^2 = 1 (independent unit-var components => total var ~ sum w_i^2)
    w = w / jnp.sqrt(jnp.sum(w * w) + 1e-12)

    keys = jr.split(key, m)
    xs = []
    for i, ell in enumerate(ells):
        xs.append(sample_rbf_gp_fft(keys[i], cfg.T_total, float(ell), pad=cfg.gp_circulant_pad))
    X = jnp.stack(xs, axis=0)  # (m, T)
    x = jnp.tensordot(w, X, axes=([0], [0]))  # (T,)
    # final normalize (numerical)
    x = x - jnp.mean(x)
    x = x / (jnp.std(x) + 1e-8)
    return x


# ----------------------------
# HiPPO discretization (exact ZOH)
# ----------------------------

def discretize_hippo_zoh(A: jnp.ndarray, b: jnp.ndarray, dt: float, horizon: float) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Continuous-time HiPPO: dS/dt = (A/h) S + (b/h) u(t)  (scaling by horizon h)
    Discretize with exact ZOH at step dt:
        S_{t+dt} = Ad S_t + bd u_t
    """
    Ah = A / horizon
    bh = b / horizon
    Ad = expm(Ah * dt)
    # bd = \int_0^dt exp(Ah * t) bh dt = Ah^{-1}(Ad - I) bh
    I = jnp.eye(A.shape[0], dtype=A.dtype)
    # Solve Ah X = (Ad - I) bh  => X = Ah^{-1}(Ad-I)bh
    bd = jnp.linalg.solve(Ah, (Ad - I) @ bh)
    return Ad, bd


# ----------------------------
# Reduced-rank regression from streaming covariances
# ----------------------------

def _sym_sqrt_and_invsqrt(S: jnp.ndarray, ridge: float) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Return (S^{1/2}, S^{-1/2}) for PSD S via eigendecomposition with ridge."""
    # Add ridge to stabilize
    S = 0.5 * (S + S.T) + ridge * jnp.eye(S.shape[0], dtype=S.dtype)
    evals, evecs = jnp.linalg.eigh(S)
    evals = jnp.maximum(evals, 1e-12)
    sqrt = (evecs * jnp.sqrt(evals)) @ evecs.T
    invsqrt = (evecs * (1.0 / jnp.sqrt(evals))) @ evecs.T
    return sqrt, invsqrt


def rrr_map_and_projector_from_covs(
    Sigma_xx: jnp.ndarray,
    Sigma_yy: jnp.ndarray,
    Sigma_yx: jnp.ndarray,
    d: int,
    ridge: float,
):
    """
    y_hat = W_d x  (rank-d RRR)
    also returns P_x: rank-d projector on x-space that preserves the bottleneck subspace.

    Uses whitening:
      C = \Sigma_yy^{-1/2} \Sigma_yx \Sigma_xx^{-1/2} = U diag(s) V^T
      W_d = \Sigma_yy^{1/2} U_d diag(s_d) V_d^T \Sigma_xx^{-1/2}
      P_x = \Sigma_xx^{1/2} V_d V_d^T \Sigma_xx^{-1/2}
    """
    Sy_sqrt, Sy_invsqrt = _sym_sqrt_and_invsqrt(Sigma_yy, ridge)
    Sx_sqrt, Sx_invsqrt = _sym_sqrt_and_invsqrt(Sigma_xx, ridge)

    C = Sy_invsqrt @ Sigma_yx @ Sx_invsqrt
    U, s, Vt = jnp.linalg.svd(C, full_matrices=False)
    U_d = U[:, :d]
    s_d = s[:d]
    V_d = Vt[:d, :].T  # (nx, d)

    W_d = Sy_sqrt @ (U_d * s_d) @ V_d.T @ Sx_invsqrt
    P_x = Sx_sqrt @ (V_d @ V_d.T) @ Sx_invsqrt
    return W_d, P_x



# ----------------------------
# Forecasting HiPPO rollout (two horizons in parallel)
# ----------------------------

def run_two_forecasters(cfg: Config) -> Dict[str, Dict[str, jnp.ndarray]]:
    """
    Returns dict with keys {"short","long"} each containing:
      - "S1","S2","S3": final states
      - "Sigma_xx","Sigma_yy","Sigma_yx": covariances accumulated (x=S2, y=S3)
      - "T_d": learned RRR map (rank d)
      - "eval_func": evaluation function for reconstruction
      - "M": output vector for lag readout (oldest value)
      - "horizon": horizon float
      - "x": full input series (shared)
    """
    key = jr.PRNGKey(cfg.seed)
    key_x, key_init = jr.split(key, 2)
    x = sample_gp_mixture(key_x, cfg)

    (A, b), eval_func, _ = get_system_params(cfg.measure, cfg.n)
    M = get_output_vector(cfg.measure, cfg.n)  # used as lagged readout (oldest endpoint)
    # Two horizons
    horizons = {"short": float(cfg.horizon_short), "long": float(cfg.horizon_long)}

    # Precompute discretizations per horizon

    # Shared dynamics for systems 2 and 3 (same for short/long)
    Ad23, bd23 = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_23)

    # Task-specific dynamics for system 1
    Ad1_short, bd1_short = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_short)
    Ad1_long,  bd1_long  = discretize_hippo_zoh(A, b, cfg.dt, cfg.horizon_long)

    discs = dict(
        short=(Ad1_short, bd1_short),
        long=(Ad1_long, bd1_long),
    )

    def init_pack():
        return dict(
            S1=jnp.zeros((cfg.n,)),
            S2=jnp.zeros((cfg.n,)),
            S3=jnp.zeros((cfg.n,)),
            # streaming covs for RRR: x=S2, y=S3
            Sigma_xx=jnp.zeros((cfg.n, cfg.n)),
            Sigma_yy=jnp.zeros((cfg.n, cfg.n)),
            Sigma_yx=jnp.zeros((cfg.n, cfg.n)),
            count=jnp.array(0.0),
        )

    packs = {k: init_pack() for k in horizons.keys()}

    def step_one(pack, u_t, Ad1, bd1):

        # Update S1 with true input
        S1 = Ad1 @ pack["S1"] + bd1 * u_t

        # Lagged scalar from oldest endpoint of S1 (approx x(t-H))
        u_lag = jnp.inner(S1, M)

        # Update S2 with lagged input (encodes window before the recent one)
        S2 = Ad23 @ pack["S2"] + bd23 * u_lag

        # Update S3 with true input (same dynamics as S2, but on recent window)
        S3 = Ad23 @ pack["S3"] + bd23 * u_t

        count = pack["count"]
        Sigma_xx = pack["Sigma_xx"]
        Sigma_yy = pack["Sigma_yy"]
        Sigma_yx = pack["Sigma_yx"]

        return S1, S2, S3, Sigma_xx, Sigma_yy, Sigma_yx, count

    for t in tqdm(range(cfg.T_train)):
        u_t = x[t]
        for name, H in horizons.items():
            Ad, bd = discs[name]
            pack = packs[name]
            S1, S2, S3, Sigma_xx, Sigma_yy, Sigma_yx, count = step_one(pack, u_t, Ad, bd)

            if t >= cfg.burnin:
                # online second-order stats
                Sigma_xx = Sigma_xx + jnp.outer(S2, S2)
                Sigma_yy = Sigma_yy + jnp.outer(S1, S1)
                Sigma_yx = Sigma_yx + jnp.outer(S1, S2)
                count = count + 1.0

            packs[name] = dict(
                S1=S1, S2=S2, S3=S3,
                Sigma_xx=Sigma_xx, Sigma_yy=Sigma_yy, Sigma_yx=Sigma_yx, count=count
            )

    out = {}
    for name, H in horizons.items():
        pack = packs[name]
        count = jnp.maximum(pack["count"], 1.0)
        Sigma_xx = pack["Sigma_xx"] / count
        Sigma_yy = pack["Sigma_yy"] / count
        Sigma_yx = pack["Sigma_yx"] / count

        T_d, P_x = rrr_map_and_projector_from_covs(Sigma_xx, Sigma_yy, Sigma_yx, cfg.bottleneck_d, cfg.ridge)

        out[name] = dict(
            horizon=jnp.array(H),
            S1=pack["S1"],
            S2=pack["S2"],
            S3=pack["S3"],
            Sigma_xx=Sigma_xx,
            Sigma_yy=Sigma_yy,
            Sigma_yx=Sigma_yx,
            T_d=T_d,
            P_x=P_x,
            eval_func=eval_func,
            M=M,
            x=x,
        )
    return out


# ----------------------------
# Reconstruction helpers + plotting
# ----------------------------

def eval_matrix_from_eval_func(eval_func, u: jnp.ndarray, n: int) -> jnp.ndarray:
    """
    Build L (len(u) x n) such that for any state S (n,),
    eval_func(u, S) == L @ S
    """
    I = jnp.eye(n)
    # columns: eval_func(u, e_j)
    cols = [eval_func(1-u, I[j]) for j in range(n)]
    L = jnp.stack(cols, axis=1)  # (len(u), n)
    return L


def rbf_kernel(delta: jnp.ndarray, ell: float) -> jnp.ndarray:
    return jnp.exp(-(delta * delta) / (2.0 * ell * ell))

def mixture_kernel(delta: jnp.ndarray, ells: jnp.ndarray, w: jnp.ndarray) -> jnp.ndarray:
    # total covariance C_tot(delta) = sum (w_m^2 * exp(-delta^2 / (2 ell_m^2)))
    # (weights w assumed already normalized so sum w^2 = 1)
    out = jnp.zeros_like(delta, dtype=jnp.float32)
    for ell, wi in zip(list(ells), list(w)):
        out = out + (wi * wi) * rbf_kernel(delta, float(ell))
    return out

def gp_posterior_mean_mixture(
    tau_query: jnp.ndarray,           # (Q,) future times (relative, continuous), e.g. in [0, H]
    t_obs: jnp.ndarray,               # (N,) observed times (relative), e.g. [-N,...,-1]
    y_obs: jnp.ndarray,               # (N,) observed values x(t_obs)
    ells: jnp.ndarray,                # (M,)
    w: jnp.ndarray,                   # (M,) with sum w^2 = 1
    ridge: float = 1e-6,
) -> jnp.ndarray:
    """
    GP posterior mean for a zero-mean stationary GP with C_tot induced by an RBF mixture.

    \mu(\tau) = K(\tau, t_obs) [K(t_obs, t_obs) + ridge I]^{-1} y_obs
    """
    # Kxx
    D_xx = t_obs[:, None] - t_obs[None, :]
    Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0])
    # Kqx
    D_qx = tau_query[:, None] - t_obs[None, :]
    Kqx = mixture_kernel(D_qx, ells, w)

    alpha = jnp.linalg.solve(Kxx, y_obs)
    return Kqx @ alpha


def gp_posterior_mean_std_from_alpha(
    tau_query: jnp.ndarray,          # (Q,)
    t_obs: jnp.ndarray,              # (N,)
    Kxx: jnp.ndarray,                # (N,N) already includes ridge
    alpha: jnp.ndarray,              # (N,) = solve(Kxx, y_obs)
    ells: jnp.ndarray,               # (M,)
    w: jnp.ndarray,                  # (M,) sum w^2 = 1
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    D_qx = tau_query[:, None] - t_obs[None, :]
    Kqx = mixture_kernel(D_qx, ells, w)         # (Q,N)
    mean = Kqx @ alpha                          # (Q,)

    # Solve Kxx^{-1} Kxq via Cholesky
    L = jnp.linalg.cholesky(Kxx)                # (N,N)
    # v = L^{-1} Kxq, where Kxq = Kqx^T
    v = jax.scipy.linalg.solve_triangular(L, Kqx.T, lower=True)   # (N,Q)
    # var = kqq - ||v||^2
    kqq = mixture_kernel(jnp.zeros_like(tau_query), ells, w)      # (Q,) = C_tot(0)=1
    var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0))          # (Q,)
    std = jnp.sqrt(var + 1e-12)
    return mean, std


def reconstruct_window(eval_func, state: jnp.ndarray, nt: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Return (u_grid in [0,1], f(u)) for a single HiPPO state reconstruction."""
    u = jnp.linspace(0.0, 1.0, nt)
    f = eval_func(u, state)
    return u, f


def eval_matrix_on_common_lags(eval_func, t_common: jnp.ndarray, H: float, n: int) -> jnp.ndarray:
    """
    Build L_common (mQ x n) on a common lag grid t_common in [-Hmax, 0].
    For lags t < -H (outside this system's window), rows are zero.
    For lags in [-H,0], use normalized u=(t+H)/H and evaluate with correct orientation.
    """
    mQ = t_common.shape[0]
    u = (t_common + H) / H                       # maps [-H,0] -> [0,1]
    valid = (u >= 0.0) & (u <= 1.0)

    # We'll build columns by evaluating basis vectors at the valid u's.
    I = jnp.eye(n)
    L = jnp.zeros((mQ, n), dtype=jnp.float32)

    u_valid = u[valid]
    # orientation fix: eval_func(1-u, *)
    cols = [eval_func(1.0 - u_valid, I[j]) for j in range(n)]  # each is (num_valid,)
    L_valid = jnp.stack(cols, axis=1)                          # (num_valid, n)

    L = L.at[valid, :].set(L_valid)
    return L


def plot_gp_oracle_zoh(
    ax,
    Hn: int,
    t_obs: jnp.ndarray,     # (N,) observed times, e.g. [-Hctx, ..., -1]
    y_obs: jnp.ndarray,     # (N,) observed values
    ells: jnp.ndarray,      # (M,)
    w: jnp.ndarray,         # (M,) sum w^2 = 1
    ridge: float = 1e-6,
    alpha_fill: float = 0.25,
    label_prefix: str = "oracle GP",
):
    """
    Plot GP posterior mean and +/-1 std for future in ZOH style on [0,Hn]:
      mean is constant on [k,k+1) equal to \mu(k), k=0..Hn-1
      band is constant on [k,k+1) equal to \mu(k)+/-\sigma(k)
    """
    # Build Kxx once
    D_xx = t_obs[:, None] - t_obs[None, :]
    Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0])
    alpha = jnp.linalg.solve(Kxx, y_obs)
    Lchol = jnp.linalg.cholesky(Kxx)

    # Query at integer times 0..Hn-1 (one per ZOH interval)
    t_int = jnp.arange(0, Hn, dtype=jnp.float32)  # (Hn,)

    D_qx = t_int[:, None] - t_obs[None, :]
    Kqx = mixture_kernel(D_qx, ells, w)
    mu = Kqx @ alpha

    # Posterior variance at integer times
    v = jax.scipy.linalg.solve_triangular(Lchol, Kqx.T, lower=True)  # (N,Hn)
    kqq = mixture_kernel(jnp.zeros_like(t_int), ells, w)             # (Hn,)
    var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0))
    std = jnp.sqrt(var + 1e-12)

    # Convert to step plotting:
    # edges: 0..Hn, value on [k,k+1) is mu[k]
    edges = jnp.arange(0, Hn + 1, dtype=jnp.float32)
    mu_extended = jnp.concatenate([jnp.array(mu), mu[-1]*jnp.ones(1)], 0)
    std_extended = jnp.concatenate([jnp.array(std), std[-1]*jnp.ones(1)], 0)
    ax.step(
        jnp.array(edges),
        mu_extended,
        where="post",
        linestyle="--",
        linewidth=2.0,
        label=f"{label_prefix} mean (ZOH)",
    )

    ax.fill_between(
        jnp.array(edges),
        jnp.array(mu_extended - std_extended),
        jnp.array(mu_extended + std_extended),
        step="post",
        alpha=alpha_fill,
        linewidth=0,
        label=f"{label_prefix} ±1 std (ZOH)",
    )


def plot_overlay_with_eigfns(
    cfg: Config,
    results: Dict[str, Dict[str, jnp.ndarray]],
    use_full_rank_Q: bool = True,
    k_eigs: int = 4,
    layout: str = "double",   # "single" or "double"
) -> None:
    """
    2x2 figure:
      Left col: overlay (short on top, long on bottom)
      Right col: top eigenfunctions of Q_hist (short on top, long on bottom)

    Styling:
      - Legend only in top-left
      - Left plots: grayscale (no Tableau colors)
      - Right plots: blue ramp for eigenfunctions (dark->light)
      - Left plots share y-limits
      - Two layout presets: single-column (skinny-ish) or two-column (wide & short)
    """
    # -------------------------
    # Figure size presets
    # -------------------------
    if layout == "single":
        # ~single column: skinny, roughly square
        figsize = (3.35, 3.35)  # inches (common single-column width ~3.3")
        width_ratios = [2.2, 1.0]
        hspace = 0.25
        wspace = 0.35
        fontsize = 8
    elif layout == "double":
        # ~two columns: short, squat
        figsize = (5.5, 2.9)    # inches (two-column width ~6.9")
        width_ratios = [1.7, 1.0]
        hspace = 0.18
        wspace = 0.30
        fontsize = 8
    else:
        raise ValueError("layout must be 'single' or 'double'")

    plt.rcParams.update({
        "font.size": fontsize,
        "axes.titlesize": fontsize,
        "axes.labelsize": fontsize,
        "legend.fontsize": fontsize - 1,
        "xtick.labelsize": fontsize - 1,
        "ytick.labelsize": fontsize - 1,
        "axes.linewidth": 0.7,
        "figure.dpi": 200,
        "savefig.dpi": 300,
    })

    fig, axes = plt.subplots(
        2, 2, figsize=figsize,
        gridspec_kw={"width_ratios": width_ratios},
        sharex=False, sharey=False
    )

    # -------------------------
    # Mixture kernel params (weights normalized so sum w^2 = 1)
    # -------------------------
    ells = jnp.array(cfg.gp_lengthscales, dtype=jnp.float32)
    if cfg.gp_weights is None:
        w = jnp.ones((len(cfg.gp_lengthscales),), dtype=jnp.float32)
    else:
        w = jnp.array(cfg.gp_weights, dtype=jnp.float32)
    w = w / jnp.sqrt(jnp.sum(w * w) + 1e-12)

    # Common axes setup
    Hmax = float(max(cfg.horizon_short, cfg.horizon_long))
    Hmax_int = int(round(Hmax))
    t0 = cfg.T_train  # ZOH: after ingesting x[t0-1], current time is t0

    # Condition ONLY on past samples occupying [-Hmax, -1]
    t_obs = jnp.arange(-Hmax_int, 0, dtype=jnp.float32)
    y_obs = jnp.array(results["long"]["x"][t0 - Hmax_int : t0], dtype=jnp.float32)

    cols = [("short", float(cfg.horizon_short)), ("long", float(cfg.horizon_long))]

    # For consistent y-lims across left panels, collect plotted y-extents
    left_ymins, left_ymaxs = [], []

    # Blue ramp for eigenfunctions: dark -> light
    # (use Matplotlib's "Blues" colormap but choose a range that avoids near-white)
    blues = plt.cm.Blues(np.linspace(0.85, 0.35, max(k_eigs, 1)))

    # Grayscale styles for left overlay curves (no Tableau colors)
    # Order: GP mean, GP band, true past, HiPPO hist, pred hist, forecast, forecast (pred)
    # We'll keep band as a light gray fill.
    style_true = dict(color="k", linewidth=0.5, alpha=1.0)
    style_hippo_hist = dict(color="maroon", linewidth=1.2, alpha=0.8)
    style_pred_hist = dict(color="firebrick", linewidth=1.2, alpha=0.8)
    style_forecast = dict(color="lightcoral", linewidth=1.2, alpha=0.8)
    style_forecast2 = dict(color="0.35", linewidth=1.2, linestyle=":")

    for row, (name, _) in enumerate(cols):
        r = results[name]
        x = r["x"]
        eval_func = r["eval_func"]
        S3 = r["S3"]
        T_d = r["T_d"]
        P_x = r["P_x"]

        H_hist = float(cfg.horizon_23)
        H_fut = float(r["horizon"])
        Hn = int(round(H_fut))

        # -------------------------
        # LEFT: overlay
        # -------------------------
        ax = axes[row, 0]

        # Oracle GP (ZOH mean + uncertainty) in grayscale
        ridge = 1e-6
        D_xx = t_obs[:, None] - t_obs[None, :]
        Kxx = mixture_kernel(D_xx, ells, w) + ridge * jnp.eye(t_obs.shape[0])
        alpha = jnp.linalg.solve(Kxx, y_obs)
        Lchol = jnp.linalg.cholesky(Kxx)

        t_int = jnp.arange(0, Hn, dtype=jnp.float32)
        D_qx = t_int[:, None] - t_obs[None, :]
        Kqx = mixture_kernel(D_qx, ells, w)
        mu = Kqx @ alpha

        v = jax.scipy.linalg.solve_triangular(Lchol, Kqx.T, lower=True)
        kqq = mixture_kernel(jnp.zeros_like(t_int), ells, w)
        var = jnp.maximum(0.0, kqq - jnp.sum(v * v, axis=0))
        std = jnp.sqrt(var + 1e-12)

        edges = jnp.arange(0, Hn + 1, dtype=jnp.float32)

        mu_extended = np.concatenate([np.array(mu), mu[-1]*np.ones(1)], 0)
        std_extended = np.concatenate([np.array(std), std[-1]*np.ones(1)], 0)
        ax.step(np.array(edges), mu_extended, where="post",
                color="0.15", linewidth=1.4, linestyle="-", label="Oracle GP mean")
        ax.fill_between(
            np.array(edges),
            np.array(mu_extended - std_extended),
            np.array(mu_extended + std_extended),
            step="post",
            color="0.85",
            alpha=0.8,
            linewidth=0,
            label="Oracle GP ±1σ"
        )

        # True past (ZOH) over [-Hmax, 0]
        past_vals_full = jnp.array(x[t0 - Hmax_int : t0])
        past_edges = jnp.arange(-Hmax_int, 1) * cfg.dt
        ax.step(np.array(past_edges[:-1]), np.array(past_vals_full), where="post",
                label="True past", **style_true)

        # HiPPO reconstructed history from S3 over [-H_hist, 0]
        nt = cfg.plot_nt
        t_hist = jnp.linspace(-H_hist, 0.0, nt)
        _, f_hist = reconstruct_window(eval_func, S3, nt)
        ax.plot(np.array(t_hist[::-1]), np.array(f_hist),
                label="HiPPO memory", **style_hippo_hist)

        # Predictive history: P_d S3
        S3_predhist = P_x @ S3
        _, f_predhist = reconstruct_window(eval_func, S3_predhist, nt)
        ax.plot(np.array(t_hist[::-1]), np.array(f_predhist),
                label="Predictive HiPPO memory", **style_pred_hist)

        # Forecast via predictive history
        t_fut = jnp.linspace(0.0, H_fut, nt)
        S_future2 = T_d @ S3_predhist
        _, f_fut2 = reconstruct_window(eval_func, S_future2, nt)
        ax.plot(np.array(t_fut[::-1]), np.array(f_fut2),
                label="HiPPO forecast", **style_forecast)

        # Formatting
        ax.axvline(0.0, linewidth=0.8, color="0.2")
        ax.set_xlim([-Hmax, Hmax])
        if row == 0:
            ax.set_title("Forecasts")

        # Only label y-axis on left column
        if row == 0:
            ax.set_ylabel(r"Short horizon ($\mathbf{H=4}$)")
        else:
            ax.set_ylabel(r"Long horizon ($\mathbf{H=32}$)")

        # Only bottom-left gets x-label
        ax.set_xlabel("Relative time" if row == 1 else "")

        # Remove legend from bottom-left
        if row == 0:
            ax.legend(loc="upper right", frameon=False, ncol=1, handlelength=2.5)
        else:
            ax.legend_.remove() if ax.get_legend() is not None else None

        # Collect y-lims for later syncing
        ylo, yhi = ax.get_ylim()
        left_ymins.append(ylo)
        left_ymaxs.append(yhi)

        # Make spines subtle
        for spine in ["top", "right"]:
            ax.spines[spine].set_visible(False)

        # -------------------------
        # RIGHT: eigenfunctions of Q_hist
        # -------------------------
        axr = axes[row, 1]

        # Choose T for Q
        if use_full_rank_Q:
            Sigma_xx = r["Sigma_xx"]
            Sigma_yx = r["Sigma_yx"]
            n = Sigma_xx.shape[0]
            T_full = Sigma_yx @ jnp.linalg.solve(Sigma_xx + cfg.ridge * jnp.eye(n), jnp.eye(n))
            T_for_Q = T_full
        else:
            T_for_Q = T_d

        Q = jnp.array(T_for_Q.T @ T_for_Q)

        # Lag grid and eval matrix
        mQ = 220
        t_common = jnp.linspace(-H_hist, 0.0, mQ)
        L_common = eval_matrix_on_common_lags(eval_func, t_common, H_hist, cfg.n)

        Q_hist = L_common @ Q @ L_common.T
        Q_hist = 0.5 * (Q_hist + Q_hist.T)

        # Eigs
        evals, evecs = jnp.linalg.eigh(Q_hist)  # ascending
        k = int(min(k_eigs, evecs.shape[1]))
        idx = jnp.argsort(evals)[::-1][:k]
        top_evals = evals[idx]
        top_evecs = evecs[:, idx]

        # sign fix: make value at lag 0 nonnegative
        signs = jnp.sign(top_evecs[-1, :] + 1e-12)
        top_evecs = top_evecs * signs

        # Plot eigenfunctions with blue ramp, no legend by default (cleaner)
        for i in range(k):
            axr.plot(np.array(t_common), np.array(top_evecs[:, i]),
                     color=blues[i], alpha=0.8, linewidth=1.4)

        axr.axvline(0.0, linewidth=0.8, color="k", alpha=0.6)
        axr.set_xlim([-H_hist, 0.0])
        axr.set_title(r"Top eigfns of $Q$")

        # Label only bottom-right x-axis
        axr.set_xlabel("lag" if row == 1 else "")
        axr.set_ylabel("" if row == 0 else "")

        # Subtle spines
        for spine in ["top", "right"]:
            axr.spines[spine].set_visible(False)

    # -------------------------
    # Sync y-limits for left panels
    # -------------------------
    ylo = float(min(left_ymins))
    yhi = float(max(left_ymaxs))
    # Add a tiny padding
    pad = 0.03 * (yhi - ylo + 1e-12)
    ylo -= pad
    yhi += pad
    axes[0, 0].set_ylim([ylo, yhi])
    axes[1, 0].set_ylim([ylo, yhi])

    # Tight layout control tuned for paper
    fig.subplots_adjust(left=0.10, right=0.98, bottom=0.12, top=0.92, wspace=wspace, hspace=hspace)

    plt.savefig(cfg.out_png, bbox_inches="tight")
    plt.close(fig)


# ----------------------------
# Main
# ----------------------------

def main():
    cfg = Config()
    cfg.T_total = cfg.T_train + int(max(cfg.horizon_short, cfg.horizon_long)) + 10

    results = run_two_forecasters(cfg)
    plot_overlay_with_eigfns(cfg, results)
    print(f"Saved: {cfg.out_png}")


if __name__ == "__main__":
    main()