"""
Multiscale HiPPO vs regular HiPPO
Explicit reconstruction of the past from 16 Legendre coefficients.

What we do:
- Sample OU-mixture signals.
- Roll out three multiscale systems + H regular HiPPO systems.
- For each lag L in `lags`:
    * Multiscale systems: interpret their (n*m) state at final time T-1 as an OP over
      the *normalized* time coordinate s \in [0,1] where s=0 is time T-1 (present) and
      s=1 is time T-1-L (past endpoint).
      We reconstruct x_hat(s) = \Sigma_{k=0}^{n-1} c_k L_k(s) using 128 s-samples, compare to the
      true signal along that window (linear interpolation at fractional indices), and
      accumulate mean MSE.
    * Regular HiPPO (each base timescale \tau_h): its coefficients represent the fixed window
      [T-1-\tau_h, T-1]. For each lag L we still evaluate the *last* L steps portion by
      sampling s \in [0,1] over the whole \tau_h window and comparing on the last L-step segment.
      (So if L > \tau_h, the earliest part is outside the window; we clamp to available region.)
- Plot average MSE vs lag (log x-axis), one curve per model; regular HiPPO gives H curves.

Important conventions:
- For multiscale: s=0 <-> time T (here T means last index T-1) and s=1 <-> time T-lag.
  So time(s) = (T-1) - s*lag.
- For regular HiPPO with window \tau_h: time(s) = (T-1) - s*\tau_h.
  We compare over the interval [T-1-lag, T-1] by restricting s to [0, lag/\tau_h] (if lag<=\tau_h),
  otherwise we compare over the full window and treat the extra as outside (0 target).
"""

from dataclasses import dataclass
from typing import Tuple

import jax
import jax.numpy as jnp
import jax.random as jr
from jax import jit
from jax.lax import scan
from jax.scipy.linalg import expm
from scipy.special import iv

jax.config.update("jax_enable_x64", True)

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

from tqdm import tqdm

from mpm import get_system_params


# ----------------------------
# 1) OU-mixture signal
# ----------------------------

def sample_ou_mixture_1d(
    key: jax.Array,
    T: int,
    ou_components: int,
    tau_min: float,
    tau_max: float,
    dt: float = 1.0,
    taus: jax.Array | None = None,
) -> Tuple[jax.Array, jax.Array]:
    if taus is None:
        key_tau, key_init, key_noise = jr.split(key, 3)
        log_tau = jr.uniform(
            key_tau, (ou_components,), minval=jnp.log(tau_min), maxval=jnp.log(tau_max)
        )
        taus = jnp.exp(log_tau)
    else:
        key_init, key_noise = jr.split(key, 2)

    a = jnp.exp(-dt / taus)
    b = jnp.sqrt(jnp.maximum(0.0, 1.0 - a * a))

    z0 = jr.normal(key_init, (ou_components,))
    eps = jr.normal(key_noise, (T - 1, ou_components))

    def step(z, e):
        z_next = a * z + b * e
        return z_next, z_next

    _, zs = scan(step, z0, eps)
    zs = jnp.concatenate([z0[None, :], zs], axis=0)
    x = jnp.sum(zs, axis=1)
    x = x / jnp.sqrt(float(ou_components))
    return x.astype(jnp.float32), taus


# ----------------------------
# 2) Multiscale matrices: C, B and dense ZOH discretization
# ----------------------------

def vecT(S):
    return jnp.reshape(S.T, (-1,))


def unvecT(v, n, m):
    return jnp.reshape(v, (m, n)).T


def legendre_shifted_mult_by_c_matrix(m: int) -> jax.Array:
    n = jnp.arange(m, dtype=jnp.float32)
    a_up = (n[:-1] + 1.0) / jnp.sqrt((2.0 * n[:-1] + 1.0) * (2.0 * n[:-1] + 3.0))
    a_dn = (n[1:]) / jnp.sqrt((2.0 * n[1:] + 1.0) * (2.0 * n[1:] - 1.0))
    Mx = jnp.zeros((m, m), dtype=jnp.float32)
    Mx = Mx.at[jnp.arange(m - 1), jnp.arange(1, m)].set(a_up)
    Mx = Mx.at[jnp.arange(1, m), jnp.arange(m - 1)].set(a_dn)
    C = 0.5 * jnp.eye(m, dtype=jnp.float32) + 0.5 * Mx
    return C


def legendre_orthonormal_Jx(m: int, dtype=jnp.float64) -> jax.Array:
    n = jnp.arange(m - 1, dtype=dtype)
    off = (n + 1.0) / jnp.sqrt((2.0 * n + 1.0) * (2.0 * n + 3.0))
    J = jnp.zeros((m, m), dtype=dtype)
    J = J.at[jnp.arange(m - 1), jnp.arange(1, m)].set(off)
    J = J.at[jnp.arange(1, m), jnp.arange(m - 1)].set(off)
    return J


def build_G_log_g(m: int, eps: float, dtype=jnp.float64) -> jax.Array:
    L = jnp.log(1.0 / jnp.asarray(eps, dtype=dtype))
    kappa = 0.5 * L
    Jx = legendre_orthonormal_Jx(m, dtype=dtype)
    G = jnp.sqrt(jnp.asarray(eps, dtype=dtype)) * expm(kappa * Jx)
    G = 0.5 * (G + G.T)
    return G


def modified_spherical_in(n: jax.Array, z: jax.Array) -> jax.Array:
    n = jnp.asarray(n)
    z = jnp.asarray(z)
    v = n + 0.5
    return jnp.sqrt(jnp.pi / (2.0 * z)) * jnp.array(iv(np.array(v), np.array(z)))


def beta_for_exp_u_in_u_legendre(m: int, eps: float, dtype=jnp.float64) -> jax.Array:
    L = jnp.log(1.0 / jnp.asarray(eps, dtype=dtype))
    kappa = 0.5 * L
    n = jnp.arange(m, dtype=dtype)
    i_n = modified_spherical_in(n, kappa)
    beta = jnp.sqrt(jnp.asarray(eps, dtype=dtype)) * jnp.sqrt((2.0 * n + 1.0) * L) * i_n
    return beta


def build_B_from_b_log_g(b: jax.Array, m: int, eps: float, dtype=jnp.float64) -> jax.Array:
    beta = beta_for_exp_u_in_u_legendre(m, eps, dtype=dtype)
    return jnp.asarray(b, dtype=dtype)[:, None] * beta[None, :]


def build_B_from_b(b: jax.Array, m: int) -> jax.Array:
    n = b.shape[0]
    B = jnp.zeros((n, m), dtype=b.dtype)
    B = B.at[:, 0].set(0.5 * b)
    if m >= 2:
        B = B.at[:, 1].set((0.5 / jnp.sqrt(3.0)) * b)
    return B


def build_B_from_b_jeffreys(b: jax.Array, m: int, eps: float) -> jax.Array:
    b = jnp.asarray(b)
    mu0 = jnp.log(1.0 / jnp.asarray(eps, dtype=jnp.float64))
    mu1 = (1.0 - eps)
    mu2 = 0.5 * (1.0 - eps * eps)

    beta0 = (mu1 / jnp.sqrt(mu0)).astype(b.dtype)
    beta1 = jnp.sqrt(jnp.maximum(mu2 - (mu1 * mu1) / mu0, 0.0)).astype(b.dtype)

    B = jnp.zeros((b.shape[0], m), dtype=b.dtype)
    B = B.at[:, 0].set(beta0 * b)
    if m >= 2:
        B = B.at[:, 1].set(beta1 * b)
    return B


def jeffreys_eval_basis(g: jax.Array, a: jax.Array, b: jax.Array, eps: float, m: int) -> jax.Array:
    """
    Return Q(g) = [Q_0(g), ..., Q_{m-1}(g)] for Jeffreys-orthonormal polynomials.

    Recurrence:
      g Q_n = a_{n+1} Q_{n+1} + b_n Q_n + a_n Q_{n-1},  a_0=0.
    """
    g = jnp.asarray(g)
    mu0 = jnp.log(1.0 / jnp.asarray(eps, dtype=g.dtype))
    Q0 = 1.0 / jnp.sqrt(mu0)

    # Q_0
    Q = [jnp.broadcast_to(Q0, g.shape)]

    if m == 1:
        return jnp.stack(Q, axis=-1)

    # Q_1 = ((g - b0)/a1) Q0
    Q1 = ((g - b[0]) / a[1]) * Q[0]
    Q.append(Q1)

    # Q_{n+1} = ((g - b_n) Q_n - a_n Q_{n-1}) / a_{n+1}
    for n in range(1, m - 1):
        Qnp1 = ((g - b[n]) * Q[n] - a[n] * Q[n - 1]) / a[n + 1]
        Q.append(Qnp1)

    return jnp.stack(Q, axis=-1)  # (..., m)


def jeffreys_clenshaw(g: jax.Array, coeffs: jax.Array, a: jax.Array, b: jax.Array, eps: float) -> jax.Array:
    """
    Evaluate sum_{k=0}^{m-1} coeffs[k] Q_k(g) using Clenshaw.
    Here a has length >= m+2 (we return a_0..a_{m+2}), b has length >= m.
    """
    g = jnp.asarray(g)
    coeffs = jnp.asarray(coeffs)
    m = coeffs.shape[-1]

    mu0 = jnp.log(1.0 / jnp.asarray(eps, dtype=g.dtype))
    Q0 = 1.0 / jnp.sqrt(mu0)

    def one(x, c):
        d1 = jnp.asarray(0.0, dtype=x.dtype)
        d2 = jnp.asarray(0.0, dtype=x.dtype)
        for k in range(m - 1, -1, -1):
            d0 = c[k] + ((x - b[k]) / a[k + 1]) * d1 - (a[k + 1] / a[k + 2]) * d2
            d2, d1 = d1, d0
        return Q0 * d1

    xflat = g.reshape(-1)
    if coeffs.ndim == 1:
        cflat = jnp.broadcast_to(coeffs, (xflat.shape[0], m))
    else:
        cflat = coeffs.reshape((-1, m))
    out = jax.vmap(one)(xflat, cflat)
    return out.reshape(g.shape)


def zoh_discretize_dense(K: jax.Array, g: jax.Array, dt: float) -> Tuple[jax.Array, jax.Array]:
    D = K.shape[0]
    Z = jnp.zeros((D + 1, D + 1), dtype=K.dtype)
    Z = Z.at[:D, :D].set(K)
    Z = Z.at[:D, D].set(g)
    E = expm(dt * Z)
    A_d = E[:D, :D]
    B_d = E[:D, D]
    return A_d.astype(jnp.float32), B_d.astype(jnp.float32)


def jeffreys_nodes_weights(eps: float, Q: int, dtype=jnp.float64):
    eps = jnp.asarray(eps, dtype=dtype)
    u0 = jnp.log(eps)
    u1 = jnp.array(0.0, dtype=dtype)
    i = jnp.arange(Q, dtype=dtype)
    du = (u1 - u0) / Q
    u = u0 + (i + 0.5) * du
    g = jnp.exp(u)
    w = jnp.full((Q,), du, dtype=dtype)  # integral is \int f(exp(u)) du
    return g, w


def jeffreys_recurrence_discrete(m: int, eps: float, Q: int = 8192, dtype=jnp.float64):
    g, w = jeffreys_nodes_weights(eps, Q, dtype=dtype)

    def inner(x, y):
        return jnp.sum(w * x * y)

    mu0 = jnp.log(1.0 / jnp.asarray(eps, dtype=dtype))
    p_prev = jnp.zeros((Q,), dtype=dtype)
    p = jnp.full((Q,), 1.0 / jnp.sqrt(mu0), dtype=dtype)

    a = [jnp.array(0.0, dtype=dtype)]
    b = []
    a_n = jnp.array(0.0, dtype=dtype)

    for _n in range(m + 1):
        gp = g * p
        b_n = inner(gp, p)
        r = gp - b_n * p - a_n * p_prev
        a_np1 = jnp.sqrt(jnp.maximum(inner(r, r), 0.0))
        b.append(b_n)
        a.append(a_np1)
        p_next = jnp.where(a_np1 > 0, r / a_np1, jnp.zeros_like(r))
        p_prev, p = p, p_next
        a_n = a_np1

    # one more for Clenshaw safety
    gp = g * p
    b_mp1 = inner(gp, p)
    r = gp - b_mp1 * p - a_n * p_prev
    a_mp2 = jnp.sqrt(jnp.maximum(inner(r, r), 0.0))
    a.append(a_mp2)

    return jnp.stack(a), jnp.stack(b)


def jeffreys_mult_by_g_c_matrix_discrete(
    m: int, eps: float, Q: int = 8192, dtype=jnp.float64
):
    a, b = jeffreys_recurrence_discrete(m=m, eps=eps, Q=Q, dtype=dtype)
    diag = b[:m]
    off = a[1:m]
    C = jnp.zeros((m, m), dtype=dtype)
    C = C.at[jnp.arange(m), jnp.arange(m)].set(diag)
    if m >= 2:
        C = C.at[jnp.arange(m - 1), jnp.arange(1, m)].set(off)
        C = C.at[jnp.arange(1, m), jnp.arange(m - 1)].set(off)
    return 0.5 * (C + C.T)


# ----------------------------
# 3) Orthonormal shifted Legendre evaluation on [0,1]
# ----------------------------

def shifted_legendre_orthonormal_vals(s: jax.Array, n: int) -> jax.Array:
    """
    L_k(s) orthonormal on [0,1]:
      L_k(s) = sqrt(2k+1) P_k(2s-1)
    Returns shape (..., n)
    """
    s = jnp.asarray(s)
    x = 2.0 * s - 1.0

    # P_0, P_1
    P0 = jnp.ones_like(x)
    if n == 1:
        return P0[..., None] * jnp.sqrt(1.0)

    P1 = x
    Ps = [P0, P1]
    for k in range(1, n - 1):
        Pkp1 = ((2 * k + 1) * x * Ps[k] - k * Ps[k - 1]) / (k + 1)
        Ps.append(Pkp1)

    P = jnp.stack(Ps[:n], axis=-1)  # (..., n)
    scale = jnp.sqrt(2.0 * jnp.arange(n, dtype=s.dtype) + 1.0)
    return P * scale


# ----------------------------
# 4) Config
# ----------------------------

@dataclass
class Config:
    T: int = 30_000
    dt: float = 1.0

    n: int = 16
    m: int = 128
    measure: str = "legt"
    base_timescale: float = 5.0
    hippo_base_timescales: Tuple[float, ...] = (10.0, 100.0, 1000.0, 10000.0)
    epsilon: float = 1e-4

    num_iters: int = 64
    num_lags: int = 20
    num_samps: int = 128  # samples in [0,1] for reconstruction
    lag_min: int = 10
    lag_max: int = 30000

    # Signal
    ou_components: int = 20
    tau_min: float = 2.0
    tau_max: float = 2000.0


# ----------------------------
# 5) Run: roll out systems and evaluate MSE-vs-lag
# ----------------------------

def run(seed: int = 0):
    cfg = Config()
    key = jr.PRNGKey(seed)

    # HiPPO continuous params
    (A, b), _, _ = get_system_params(cfg.measure, cfg.n)
    print("A", A.shape, "b", b.shape)
    A = (A / cfg.base_timescale).astype(jnp.float32)
    b = (b / cfg.base_timescale).astype(jnp.float32)

    (A_nm, b_nm), _, _ = get_system_params(cfg.measure, cfg.n)
    A_nm = A_nm.astype(jnp.float32)
    b_nm = b_nm.astype(jnp.float32)

    # Multiscale (c)
    C = legendre_shifted_mult_by_c_matrix(cfg.m).astype(jnp.float32)
    print("C", C.shape)
    B = build_B_from_b(b, cfg.m).astype(jnp.float32)
    K_ms = jnp.kron(C.T, A)
    g_ms = vecT(B)
    Ams_d, Bms_d = zoh_discretize_dense(K_ms, g_ms, dt=cfg.dt)

    # Multiscale (Jeffreys)
    C_jeff = jeffreys_mult_by_g_c_matrix_discrete(cfg.m, cfg.epsilon, Q=8192, dtype=jnp.float64).astype(jnp.float32)
    B_jeff = build_B_from_b_jeffreys(b, cfg.m, cfg.epsilon).astype(jnp.float32)
    K_ms_jeff = jnp.kron(C_jeff.T, A)
    g_ms_jeff = vecT(B_jeff)
    Amsj_d, Bmsj_d = zoh_discretize_dense(K_ms_jeff, g_ms_jeff, dt=cfg.dt)

    # Multiscale (log-g)
    G = build_G_log_g(cfg.m, cfg.epsilon, dtype=jnp.float64).astype(jnp.float32)
    B_log = build_B_from_b_log_g(b, cfg.m, cfg.epsilon, dtype=jnp.float64).astype(jnp.float32)
    K_ms_log = jnp.kron(G.T, A)
    g_ms_log = vecT(B_log)
    Amsg_d, Bmsg_d = zoh_discretize_dense(K_ms_log, g_ms_log, dt=cfg.dt)

    # Regular HiPPO at multiple base timescales
    hippo_base_timescales = jnp.array(cfg.hippo_base_timescales, dtype=jnp.float32)

    def hippo_one(scale):
        Kk = (A_nm / scale).astype(jnp.float32)
        gk = (b_nm / scale).astype(jnp.float32)
        return zoh_discretize_dense(Kk, gk, dt=cfg.dt)

    Ahp_d, Bhp_d = jax.vmap(hippo_one)(hippo_base_timescales)

    # Rollout step (stores states after ingesting u_t)
    @jit
    def step(carry, u_t):
        s_ms, s_msj, s_msg, s_hp = carry
        s_ms_next = Ams_d @ s_ms + Bms_d * u_t
        s_msj_next = Amsj_d @ s_msj + Bmsj_d * u_t
        s_msg_next = Amsg_d @ s_msg + Bmsg_d * u_t
        s_hp_next = jnp.einsum("hnm,hm->hn", Ahp_d, s_hp) + Bhp_d * u_t
        return (s_ms_next, s_msj_next, s_msg_next, s_hp_next), (s_ms_next, s_msj_next, s_msg_next, s_hp_next)

    # Lags (log-spaced, unique, >= lag_min and <= lag_max)
    lag_min = max(1, int(cfg.lag_min))
    lag_max = min(int(cfg.lag_max), cfg.T - 1)
    lags = jnp.exp(jnp.linspace(jnp.log(float(lag_min)), jnp.log(float(lag_max)), cfg.num_lags))
    lags = jnp.unique(jnp.clip(jnp.round(lags).astype(jnp.int32), lag_min, lag_max))

    ms_sum = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    msj_sum = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    msg_sum = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    hp_sum = jnp.zeros((len(cfg.hippo_base_timescales), lags.shape[0]), dtype=jnp.float64)
    ms_sum_sq = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    msj_sum_sq = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    msg_sum_sq = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    hp_sum_sq = jnp.zeros((len(cfg.hippo_base_timescales), lags.shape[0]), dtype=jnp.float64)
    oracle_sum    = jnp.zeros((lags.shape[0],), dtype=jnp.float64)
    oracle_sum_sq = jnp.zeros((lags.shape[0],), dtype=jnp.float64)

    # Precompute s-grid and Legendre values on [0,1]
    s_grid = jnp.linspace(0.0, 1.0, cfg.num_samps, dtype=jnp.float32) # (S,)
    L_grid_flip = shifted_legendre_orthonormal_vals(1.0 - s_grid, cfg.n).astype(jnp.float32) # (S,n)

    # Oracle projection matrix P = Phi @ (Phi.T @ Phi)^{-1} @ Phi.T  shape (S, S)
    Phi_oracle = L_grid_flip.astype(jnp.float64)  # (S, n)
    P_oracle = Phi_oracle @ jnp.linalg.solve(Phi_oracle.T @ Phi_oracle, Phi_oracle.T)  # (S, S)

    
    # Helper: linear interpolation of signal at fractional indices
    def interp_signal(sig: jax.Array, t_float: jax.Array) -> jax.Array:
        """
        sig: (T,)
        t_float: (...,) float time index
        Returns sig(t_float) via linear interpolation, with out-of-range -> 0.
        """
        T = sig.shape[0]
        t0 = jnp.floor(t_float).astype(jnp.int32)
        t1 = t0 + 1
        w = t_float - t0.astype(t_float.dtype)

        in0 = (t0 >= 0) & (t0 < T)
        in1 = (t1 >= 0) & (t1 < T)
        t0c = jnp.clip(t0, 0, T - 1)
        t1c = jnp.clip(t1, 0, T - 1)
        v0 = sig[t0c] * in0.astype(sig.dtype)
        v1 = sig[t1c] * in1.astype(sig.dtype)
        return (1.0 - w) * v0 + w * v1

    # Helper: compute per-lag MSE for one coefficient vector c (n,) representing [T..T-lag]
    def mse_for_coeffs_over_lag(sig: jax.Array, coeff: jax.Array, lag: int) -> jax.Array:
        """
        coeff: (n,) in shifted orthonormal Legendre basis over s\in[0,1]
        s=0 at time T (index T-1), s=1 at time T-lag.
        Compare on 128 s-grid points.
        """
        # predicted curve on s_grid
        yhat = L_grid_flip @ coeff  # (S,)

        # true curve: time(s) = (T-1) - s*lag
        Tlast = sig.shape[0] - 1
        t_float = Tlast - s_grid.astype(jnp.float32) * float(lag)
        ytrue = interp_signal(sig, t_float)  # (S,)

        return jnp.mean((yhat.astype(jnp.float64) - ytrue.astype(jnp.float64)) ** 2)


    def mse_for_coeffs_over_lag_regular(sig, coeff, tau, lag):
        tau = jnp.asarray(tau, jnp.float32)
        lagf = jnp.asarray(lag, jnp.float32)

        Tlast = sig.shape[0] - 1

        # Evaluate over the *lag* window
        t_float = Tlast - s_grid * lagf
        ytrue = interp_signal(sig, t_float)

        # Map those times into the HiPPO window coordinate
        s_tau = (Tlast - t_float) / tau          # = s_grid * lag/tau
        valid = (s_tau >= 0.0) & (s_tau <= 1.0)

        s_eval = 1.0 - s_tau
        L_hp = shifted_legendre_orthonormal_vals(jnp.clip(s_eval, 0.0, 1.0), cfg.n).astype(jnp.float32)

        yhat_in = L_hp @ coeff                    # defined everywhere but only meaningful on valid
        yhat = jnp.where(valid, yhat_in, 0.0)     # zero prediction outside window

        return jnp.mean((yhat.astype(jnp.float64) - ytrue.astype(jnp.float64))**2)


    def oracle_mse_for_lag(sig: jax.Array, lag: int) -> jax.Array:
        """MSE of the best possible projection of the true signal onto the n-Legendre basis."""
        Tlast = sig.shape[0] - 1
        t_float = Tlast - s_grid.astype(jnp.float32) * float(lag)
        ytrue = interp_signal(sig, t_float).astype(jnp.float64)  # (S,)
        yhat = P_oracle @ ytrue                                    # (S,)
        return jnp.mean((ytrue - yhat) ** 2)


    # Main Monte Carlo loop
    example = None
    for _itr in tqdm(range(cfg.num_iters)):
        key, subkey = jr.split(key)
        signal, _taus = sample_ou_mixture_1d(
            subkey,
            cfg.T,
            cfg.ou_components,
            cfg.tau_min,
            cfg.tau_max,
            dt=cfg.dt,
        )
        # Roll out states across the whole signal
        s_ms0 = jnp.zeros((cfg.n * cfg.m,), dtype=jnp.float32)
        s_msj0 = jnp.zeros((cfg.n * cfg.m,), dtype=jnp.float32)
        s_msg0 = jnp.zeros((cfg.n * cfg.m,), dtype=jnp.float32)
        s_hp0 = jnp.zeros((len(cfg.hippo_base_timescales), cfg.n), dtype=jnp.float32)

        (s_ms_T, s_msj_T, s_msg_T, s_hp_T), states = scan(
            step, (s_ms0, s_msj0, s_msg0, s_hp0), signal
        )
        # Use final states at T-1
        S_ms_T  = unvecT(s_ms_T,  cfg.n, cfg.m)   # (n,m)
        S_msj_T = unvecT(s_msj_T, cfg.n, cfg.m)
        S_msg_T = unvecT(s_msg_T, cfg.n, cfg.m)

        # s_hp_T: (H,n)

        if _itr == 0:
            example = {
                "signal": np.array(signal),
                "s_ms_T": np.array(s_ms_T),
                "s_msj_T": np.array(s_msj_T),
                "s_msg_T": np.array(s_msg_T),
                "s_hp_T": np.array(s_hp_T),
            }

        # For each lag, extract coefficients and compute MSE
        for li in range(lags.shape[0]):
            lag = int(lags[li])

            # Multiscale (c): interpret the *lag window* using timescale-axis coordinate c = base_timescale / lag,
            # clipped into [0,1] since the c-basis is on [0,1].
            c = float(cfg.base_timescale) / float(lag)
            c = float(np.clip(c, 0.0, 1.0))
            q_c = shifted_legendre_orthonormal_vals(jnp.array([c], dtype=jnp.float32), cfg.m)[0]  # (m,)
            coeff_ms = (unvecT(s_ms_T, cfg.n, cfg.m) @ q_c).astype(jnp.float32)  # (n,)

            # Multiscale (Jeffreys g): g = base_timescale / lag, clipped into [eps,1]
            g = float(cfg.base_timescale) / float(lag)
            g = float(np.clip(g, float(cfg.epsilon), 1.0))
            # Evaluate Jeffreys basis via Clenshaw by feeding one-hot coeffs (cheap at m=16)
            eye_m = jnp.eye(cfg.m, dtype=jnp.float32)  # (m,m)
            a_rec, b_rec = jeffreys_recurrence_discrete(cfg.m, cfg.epsilon, Q=8192, dtype=jnp.float64)
            a_rec = a_rec.astype(jnp.float32)
            b_rec = b_rec.astype(jnp.float32)
            q_g = jeffreys_eval_basis(g=jnp.array(g, dtype=jnp.float32), a=a_rec, b=b_rec, eps=cfg.epsilon, m=cfg.m)  # (m,)
            coeff_msj = unvecT(s_msj_T, cfg.n, cfg.m) @ q_g  # (n,)

            # Multiscale (log-g): u = log(g) in [log eps, 0], use mapped Legendre on u
            # We'll implement basis eval in u by mapping to x in [-1,1] and using orthonormal Legendre.
            u0 = float(np.log(cfg.epsilon))
            L = -u0
            u_val = float(np.log(g))
            x = (2.0 * u_val - u0) / L  # in [-1,1]
            # compute orthonormal Legendre phi_n(x)=sqrt((2n+1)/2)P_n(x), then scale by sqrt(2/L) for uniform-u
            # reuse shifted_legendre_orthonormal_vals by mapping: shifted on [0,1] isn't convenient; do direct P_n.
            # We'll do direct recurrence for P_n at scalar x.
            xj = jnp.array(x, dtype=jnp.float32)
            P = [jnp.array(1.0, dtype=jnp.float32)]
            if cfg.m >= 2:
                P.append(xj)
            for k in range(1, cfg.m - 1):
                Pkp1 = ((2 * k + 1) * xj * P[k] - k * P[k - 1]) / (k + 1)
                P.append(Pkp1)
            P = jnp.stack(P[:cfg.m], axis=0)  # (m,)
            phi_scale = jnp.sqrt((2.0 * jnp.arange(cfg.m, dtype=jnp.float32) + 1.0) / 2.0)
            phi = P * phi_scale
            q_u = jnp.sqrt(2.0 / L) * phi  # (m,)
            coeff_msg = (unvecT(s_msg_T, cfg.n, cfg.m) @ q_u).astype(jnp.float32)

            # MSEs
            mse_ms = mse_for_coeffs_over_lag(signal, coeff_ms, lag)
            mse_msj = mse_for_coeffs_over_lag(signal, coeff_msj, lag)
            mse_msg = mse_for_coeffs_over_lag(signal, coeff_msg, lag)
            ms_sum = ms_sum.at[li].add(mse_ms)
            msj_sum = msj_sum.at[li].add(mse_msj)
            msg_sum = msg_sum.at[li].add(mse_msg)
            ms_sum_sq = ms_sum_sq.at[li].add(mse_ms**2)
            msj_sum_sq = msj_sum_sq.at[li].add(mse_msj**2)
            msg_sum_sq = msg_sum_sq.at[li].add(mse_msg**2)
            mse_oracle = oracle_mse_for_lag(signal, lag)
            oracle_sum    = oracle_sum.at[li].add(mse_oracle)
            oracle_sum_sq = oracle_sum_sq.at[li].add(mse_oracle ** 2)

            # Regular HiPPO systems
            for hi in range(len(cfg.hippo_base_timescales)):
                tau = float(cfg.hippo_base_timescales[hi])
                coeff_hp = s_hp_T[hi].astype(jnp.float32)  # (n,)
                mse_hp = mse_for_coeffs_over_lag_regular(signal, coeff_hp, tau=tau, lag=lag)
                hp_sum = hp_sum.at[hi, li].add(mse_hp)
                hp_sum_sq = hp_sum_sq.at[hi, li].add(mse_hp**2)

    num_iters = float(cfg.num_iters)
    ms_avg_mses = ms_sum / num_iters
    msj_avg_mses = msj_sum / num_iters
    msg_avg_mses = msg_sum / num_iters
    hp_avg_mses = hp_sum / num_iters
    ms_sem = jnp.sqrt(jnp.maximum(ms_sum_sq / num_iters - ms_avg_mses**2, 0.0) / num_iters)
    msj_sem = jnp.sqrt(jnp.maximum(msj_sum_sq / num_iters - msj_avg_mses**2, 0.0) / num_iters)
    msg_sem = jnp.sqrt(jnp.maximum(msg_sum_sq / num_iters - msg_avg_mses**2, 0.0) / num_iters)
    hp_sem = jnp.sqrt(jnp.maximum(hp_sum_sq / num_iters - hp_avg_mses**2, 0.0) / num_iters)
    oracle_avg_mses = oracle_sum / num_iters
    oracle_sem = jnp.sqrt(jnp.maximum(oracle_sum_sq / num_iters - oracle_avg_mses**2, 0.0) / num_iters)

    return np.array(lags), (
        np.array(ms_avg_mses),
        np.array(msj_avg_mses),
        np.array(msg_avg_mses),
        np.array(hp_avg_mses),
        np.array(ms_sem),
        np.array(msj_sem),
        np.array(msg_sem),
        np.array(hp_sem),
        np.array(oracle_avg_mses),
        np.array(oracle_sem),
    ), example


# ----------------------------
# 6) Plot
# ----------------------------

def save_mse_npz(out, filepath: str = "mse_results.npz") -> None:
    lags, (ms_mse, msj_mse, msg_mse, hp_mse, ms_sem, msj_sem, msg_sem, hp_sem, oracle_mse, oracle_sem), example = out
    save_kwargs = dict(
        lags=np.array(lags),
        ms_mse=np.array(ms_mse),
        msj_mse=np.array(msj_mse),
        msg_mse=np.array(msg_mse),
        hp_mse=np.array(hp_mse),
        ms_sem=np.array(ms_sem),
        msj_sem=np.array(msj_sem),
        msg_sem=np.array(msg_sem),
        hp_sem=np.array(hp_sem),
        oracle_mse=np.array(oracle_mse),
        oracle_sem=np.array(oracle_sem),
    )
    if example is not None:
        for k, v in example.items():
            save_kwargs[f"example_{k}"] = np.array(v)
    np.savez(filepath, **save_kwargs)
    print(f"Saved MSE data to {filepath}")


def load_mse_npz(filepath: str = "mse_results.npz"):
    data = np.load(filepath)

    lags = data["lags"]

    ms_mse = data["ms_mse"]
    msj_mse = data["msj_mse"]
    msg_mse = data["msg_mse"]
    hp_mse = data["hp_mse"]

    ms_sem = data["ms_sem"]
    msj_sem = data["msj_sem"]
    msg_sem = data["msg_sem"]
    hp_sem = data["hp_sem"]
    oracle_mse = data["oracle_mse"]
    oracle_sem = data["oracle_sem"]

    example_keys = ["signal", "s_ms_T", "s_msj_T", "s_msg_T", "s_hp_T"]
    if all(f"example_{k}" in data for k in example_keys):
        example = {k: data[f"example_{k}"] for k in example_keys}
    else:
        example = None

    out = (
        lags,
        (ms_mse, msj_mse, msg_mse, hp_mse,
         ms_sem, msj_sem, msg_sem, hp_sem,
         oracle_mse, oracle_sem),
        example,
    )

    return out



def plot_combined(
        out,
        cfg: Config,
        plot_reconstructions: bool = True,
        plot_other_multiscale: bool = False,
        filename: str = "multiscale_figure.png",
    ):
    lags, (ms_mse, msj_mse, msg_mse, hp_mse, ms_sem, msj_sem, msg_sem, hp_sem, oracle_mse, oracle_sem), example = out
    
    if plot_reconstructions and example is None:
        raise ValueError("Expected example data for reconstruction plot.")

    multiscale_colors = {
        "c": "tab:orange",
        "jeffreys": "tab:green",
        "log_g": "tab:red",
    }
    hippo_colors = plt.cm.Blues(np.linspace(0.35, 0.85, len(cfg.hippo_base_timescales)))

    if plot_reconstructions:
        fig, (ax_example, ax_mse) = plt.subplots(2, 1, figsize=(5.5, 5))
    else:
        fig, ax_mse = plt.subplots(figsize=(8,3))

    if plot_reconstructions:
        signal = jnp.asarray(example["signal"])
        s_msg_T = jnp.asarray(example["s_msg_T"])
        s_hp_T = jnp.asarray(example["s_hp_T"])

        Tlast = int(signal.shape[0] - 1)
        horizons = [30, 100, 200, 300, 1000]
        max_horizon = 300
        t_plot = jnp.arange(Tlast - max_horizon, Tlast + 1, dtype=jnp.float32)

        def s_from_time(t, horizon):
            return (Tlast - t) / horizon

        def log_g_coeffs_for_horizon(horizon: int) -> jax.Array:
            g = float(cfg.base_timescale) / float(horizon)
            g = float(np.clip(g, float(cfg.epsilon), 1.0))
            u0 = float(np.log(cfg.epsilon))
            L = -u0
            u_val = float(np.log(g))
            x = (2.0 * u_val - u0) / L
            xj = jnp.array(x, dtype=jnp.float32)
            P0 = jnp.array(1.0, dtype=jnp.float32)
            Ps = [P0]
            if cfg.m >= 2:
                Ps.append(xj)
            for k in range(1, cfg.m - 1):
                Ps.append(((2 * k + 1) * xj * Ps[k] - k * Ps[k - 1]) / (k + 1))
            P = jnp.stack(Ps[:cfg.m], axis=0)
            phi_scale = jnp.sqrt((2.0 * jnp.arange(cfg.m, dtype=jnp.float32) + 1.0) / 2.0)
            phi = P * phi_scale
            q_u = jnp.sqrt(2.0 / L) * phi
            return (unvecT(s_msg_T, cfg.n, cfg.m) @ q_u).astype(jnp.float32)

        t_np = np.array(t_plot).astype(np.int32)
        t_fig = t_np - np.max(t_np)
        true_signal = np.array(signal)[t_np]

        ax_example.plot(t_fig, true_signal, color="k", alpha=0.7, linewidth=1.5, label="True signal")
        red_shades = plt.cm.Reds(np.linspace(0.35, 0.9, len(horizons)))
        for horizon, color in zip(horizons, red_shades):
            s_win = s_from_time(t_plot, float(horizon))
            valid = (s_win >= 0.0) & (s_win <= 1.0)

            # Make a safe s for basis evaluation (anything inside [0,1] works for invalid points)
            s_safe = jnp.where(valid, s_win, 0.0)

            L_plot = shifted_legendre_orthonormal_vals(1.0 - s_safe, cfg.n).astype(jnp.float32)
            coeff_msg = log_g_coeffs_for_horizon(horizon)
            y_msg = L_plot @ coeff_msg

            # Now mask for plotting (NaNs break the line outside valid)
            y_msg = jnp.where(valid, y_msg, jnp.nan)

            label = "Multiscale HiPPO" if horizon == 200 else None
            ax_example.plot(
                t_fig,
                y_msg,
                color=color,
                alpha=0.9,
                lw=2.0,
                label=label,
            )

        for hi, tau in enumerate(cfg.hippo_base_timescales):
            if tau < 100 or tau > 1000:
                continue
            s_raw = (Tlast - t_plot) / tau
            valid = (s_raw >= 0.0) & (s_raw <= 1.0)
            s_eval = (1.0 - s_raw)
            L_hp = shifted_legendre_orthonormal_vals(jnp.clip(s_eval, 0.0, 1.0), cfg.n)
            y_hp = L_hp @ s_hp_T[hi]
            y_hp = np.where(np.array(valid), np.array(y_hp), np.nan)
            label = "Vanilla HiPPOs" if tau == 100 else None
            ax_example.plot(
                t_fig,
                y_hp,
                linestyle="--",
                alpha=0.9,
                lw=2.0,
                color=hippo_colors[hi],
                label=label,
            )

        ax_example.set_xlabel("Relative Time")
        ax_example.set_ylabel("Signal")
        ax_example.set_ylim(-1.3, 3)
        ax_example.set_title("Signal Reconstruction")
        ax_example.grid(True, alpha=0.25)
        ax_example.legend(frameon=False, ncol=1)

    if plot_other_multiscale:
        ax_mse.plot(
            lags,
            ms_mse,
            color=multiscale_colors["c"],
            label="Multiscale (Basic)",
        )
        ax_mse.fill_between(
            lags,
            ms_mse - ms_sem,
            ms_mse + ms_sem,
            color=multiscale_colors["c"],
            alpha=0.2,
            linewidth=0,
        )
        ax_mse.plot(
            lags,
            msj_mse,
            color=multiscale_colors["jeffreys"],
            label="Multiscale (Jeffreys-style)",
        )
        ax_mse.fill_between(
            lags,
            msj_mse - msj_sem,
            msj_mse + msj_sem,
            color=multiscale_colors["jeffreys"],
            alpha=0.2,
            linewidth=0,
        )
    ax_mse.plot(
        lags,
        msg_mse,
        color=multiscale_colors["log_g"],
        # label="Multiscale (Log-Timescale)",
    )
    ax_mse.fill_between(
        lags,
        msg_mse - msg_sem,
        msg_mse + msg_sem,
        color=multiscale_colors["log_g"],
        alpha=0.2,
        linewidth=0,
    )
    if plot_other_multiscale:
        ax_mse.legend(frameon=False, ncol=1)

    for hi, tau in enumerate(cfg.hippo_base_timescales):
        ax_mse.plot(
            lags,
            hp_mse[hi],
            linestyle="--",
            color=hippo_colors[hi],
            # label=f"HiPPO (base={tau:g})",
        )
        ax_mse.fill_between(
            lags,
            hp_mse[hi] - hp_sem[hi],
            hp_mse[hi] + hp_sem[hi],
            color=hippo_colors[hi],
            alpha=0.2,
            linewidth=0,
        )

    ax_mse.plot(
        lags,
        oracle_mse,
        color="gray",
        linestyle="-",
        linewidth=1.5,
    )
    ax_mse.fill_between(
        lags,
        oracle_mse - oracle_sem,
        oracle_mse + oracle_sem,
        color="gray",
        alpha=0.2,
        linewidth=0,
    )

    custom_lines = [
        Line2D([0], [0], color="tab:blue", lw=2.0, linestyle="--"),
        Line2D([0], [0], color="tab:red", lw=2.0),
        Line2D([0], [0], color="gray", lw=1.5),
    ]

    ax_mse.legend(
        custom_lines,
        ["Vanilla HiPPOs", "Multiscale HiPPO", "Oracle Error"],
        loc="lower right",
        frameon=False,
    )

    ax_mse.set_xscale("log")
    ax_mse.set_ylim(0, 1.06)
    ax_mse.set_xlabel("Horizon")
    ax_mse.set_ylabel("Average MSE")
    ax_mse.set_title("Reconstruction error from 16 Legendre coeffs")
    ax_mse.grid(True, which="both", alpha=0.25)

    fig.tight_layout()
    fig.savefig(filename, dpi=200)
    fig.savefig("multiscale_figure.pdf")
    print(f"Saved combined plot to {filename}")


if __name__ == "__main__":
    # cfg = Config()
    # out = run(seed=0)
    # save_mse_npz(out)
    # plot_combined(out, cfg, plot_reconstructions=True, plot_other_multiscale=False)

    cfg = Config()
    out = load_mse_npz()
    plot_combined(out, cfg, plot_reconstructions=True, plot_other_multiscale=False)