source_code/scripts/associative_memory_hippo.py

"""
Online (TBPTT) training: Fixed-ZOH HiPPO + Continuous-time Associative Memory
on the associative recall task (A/B tokens + WRITE).

This script trains a fixed-ZOH HiPPO model with a continuous-time associative
memory module for associative recall.

The model precomputes one ZOH discretization, keeps HiPPO state S and memory
bank C as separate states, applies exact memory updates each step, and trains
with TBPTT using stop_gradient at chunk boundaries.

"""
__date__ = "January 2026"

import numpy as np
import matplotlib.pyplot as plt

import jax
import jax.numpy as jnp
import jax.random as jr
from jax import jit, value_and_grad
from jax.lax import scan

import flax.linen as nn
import optax
from tqdm import tqdm

from mpm import get_system_params
from mpm.models import AssocMemHiPPO, legendre_orthonormal_basis01
from mpm.tasks.associative_memory import make_associative_recall_task_tokens


# ----------------------------
# TBPTT chunk padding
# ----------------------------
def pad_chunk(X, Y, t0: int, L: int):
    T = X.shape[0]
    t1 = min(t0 + L, T)
    ell = t1 - t0

    Xc = jnp.zeros((L, X.shape[1]), dtype=X.dtype)
    Yc = jnp.zeros((L, Y.shape[1]), dtype=Y.dtype)
    mask = jnp.zeros((L,), dtype=X.dtype)

    Xc = Xc.at[:ell].set(X[t0:t1])
    Yc = Yc.at[:ell].set(Y[t0:t1])
    mask = mask.at[:ell].set(1.0)
    return Xc, Yc, mask


# ----------------------------
# Training (TBPTT)
# ----------------------------
def main():
    # ----------------------------
    # Config
    # ----------------------------
    seed = 42
    key = jr.PRNGKey(seed)

    # task
    episode_len = 12 # must be even
    num_episodes = 10000
    num_tokens = 12

    # online training steps (episodes sampled)
    num_steps = 3500

    # TBPTT
    tbptt_steps = episode_len

    # token dims / model dims
    d_in = 24
    d_model = 32
    write_hidden = 256
    out_hidden = 256

    # HiPPO
    n_hippo = 32
    measure = "legt"
    base_timescale = 2.0

    # Associative memory truncation
    n_assoc = 32

    # opt
    lr = 1e-3
    wd = 1e-4

    # ----------------------------
    # Dataset
    # ----------------------------
    X_np, Y_np, token_table, ids_np, meta_np = make_associative_recall_task_tokens(
        episode_len=episode_len,
        num_episodes=num_episodes,
        d_in=d_in,
        num_tokens=num_tokens,
        seed=seed,
    )
    X_all = jnp.array(X_np, dtype=jnp.float32)  # (N,T,d_in)
    Y_all = jnp.array(Y_np, dtype=jnp.float32)
    ids_all = np.array(ids_np)

    # ----------------------------
    # HiPPO system + fixed ZOH discretization (dt=1)
    # ----------------------------
    (A, b), _, _ = get_system_params(measure, n_hippo)

    # scale time
    A = A / base_timescale
    b = b / base_timescale

    A = jnp.array(A, dtype=jnp.float32)
    b = jnp.array(b, dtype=jnp.float32)

    I = jnp.eye(n_hippo, dtype=jnp.float32)
    A_d = jax.scipy.linalg.expm(A)  # dt = 1
    # b_d = \int_0^1 exp(A \tau) b d\tau = A^{-1}(A_d - I)b  (use solve for stability)
    rhs = (A_d - I) @ b
    b_d = jnp.linalg.solve(A, rhs)

    # ----------------------------
    # Init model + optimizer
    # ----------------------------
    model = AssocMemHiPPO(
        d_in=d_in,
        d_model=d_model,
        n_hippo=n_hippo,
        n_assoc=n_assoc,
        A_d=A_d,
        b_d=b_d,
        write_hidden=write_hidden,
        out_hidden=out_hidden,
    )

    key, k1 = jr.split(key)
    x0 = jnp.zeros((d_in,), dtype=jnp.float32)
    S0 = jnp.zeros((d_model, n_hippo), dtype=jnp.float32)
    C0 = jnp.zeros((d_model, n_assoc), dtype=jnp.float32)
    params = model.init(k1, x0, S0, C0)["params"]

    optimizer = optax.adamw(learning_rate=lr, weight_decay=wd)
    opt_state = optimizer.init(params)

    # ----------------------------
    # Chunk rollout + TBPTT update
    # ----------------------------
    def chunk_loss_and_state(params, S_init, C_init, Xc, Yc, mask):
        """
        One TBPTT chunk.
        S_init: (d_model,n_hippo)
        C_init: (d_model,n_assoc)
        Xc,Yc:  (L,d_in) padded
        mask:   (L,) in {0,1}
        """
        def step_fn(carry, inputs):
            S, C = carry
            x_t, y_t, m_t = inputs
            S_next, C_next, yhat, aux = model.apply({"params": params}, x_t, S, C)
            mse_t = jnp.mean((yhat - y_t) ** 2)
            return (S_next, C_next), (m_t * mse_t, m_t, yhat, aux["g_write"], aux["g_out"])

        (S_final, C_final), (mse_masked, m_sum, yhat_ts, g_write_ts, g_out_ts) = scan(
            step_fn,
            (S_init, C_init),
            (Xc, Yc, mask),
        )
        denom = jnp.maximum(jnp.sum(m_sum), 1.0)
        loss = jnp.sum(mse_masked) / denom
        aux = (yhat_ts, g_write_ts, g_out_ts)
        return loss, (S_final, C_final), aux

    @jit
    def tbptt_chunk_train_step(params, opt_state, S_init, C_init, Xc, Yc, mask):
        def loss_fn(p):
            loss, (S_final, C_final), aux = chunk_loss_and_state(p, S_init, C_init, Xc, Yc, mask)
            return loss, (S_final, C_final, aux)

        (loss, (S_final, C_final, aux)), grads = value_and_grad(loss_fn, has_aux=True)(params)
        updates, opt_state = optimizer.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, S_final, C_final, loss, aux

    # ----------------------------
    # Online training loop (episode-by-episode) with TBPTT
    # ----------------------------
    losses = []
    key = jr.PRNGKey(seed + 1)
    N = X_all.shape[0]

    pbar = tqdm(range(num_steps))
    smooth_loss = None

    for step in pbar:
        key, sub = jr.split(key)
        idx = jr.randint(sub, (), 0, N)
        X_ep = X_all[idx]  # (T,d_in)
        Y_ep = Y_all[idx]  # (T,d_in)

        # Reset states per episode
        S = jnp.zeros((d_model, n_hippo), dtype=jnp.float32)
        C = jnp.zeros((d_model, n_assoc), dtype=jnp.float32)

        # TBPTT over chunks
        T = X_ep.shape[0]
        n_chunks = (T + tbptt_steps - 1) // tbptt_steps
        loss_ep = 0.0

        for c in range(n_chunks):
            t0 = c * tbptt_steps
            Xc, Yc, mask = pad_chunk(X_ep, Y_ep, t0=t0, L=tbptt_steps)

            params, opt_state, S_end, C_end, loss_chunk, _ = tbptt_chunk_train_step(
                params, opt_state, S, C, Xc, Yc, mask
            )
            loss_ep = loss_ep + loss_chunk

            # stop gradients across TBPTT boundaries
            S = jax.lax.stop_gradient(S_end)
            C = jax.lax.stop_gradient(C_end)

        loss_ep = loss_ep / n_chunks
        losses.append(float(loss_ep))

        if smooth_loss is None:
            smooth_loss = losses[-1]
        else:
            smooth_loss = 0.99 * smooth_loss + 0.01 * losses[-1]
        pbar.set_description(f"loss: {smooth_loss:.6f}")

    # ----------------------------
    # Evaluate on one held-out episode
    # ----------------------------
    X_eval = jnp.concatenate([X_all[-2], X_all[-1]], 0)
    Y_eval = jnp.concatenate([Y_all[-2], Y_all[-1]], 0)

    def full_rollout_collect(params, X, Y):
        S = jnp.zeros((d_model, n_hippo), dtype=jnp.float32)
        C = jnp.zeros((d_model, n_assoc), dtype=jnp.float32)

        def step_fn(carry, inputs):
            S, C = carry
            x_t, y_t = inputs
            S_next, C_next, yhat, aux = model.apply({"params": params}, x_t, S, C)
            loss_t = jnp.mean((yhat - y_t) ** 2)
            return (S_next, C_next), (loss_t, yhat, aux["g_write"], aux["g_out"])

        _, (loss_ts, yhat_ts, g_write_ts, g_out_ts) = scan(step_fn, (S, C), (X, Y))
        return jnp.mean(loss_ts), yhat_ts, g_write_ts, g_out_ts

    loss_eval, preds_ts, g_write_ts, g_out_ts = full_rollout_collect(params, X_eval, Y_eval)
    preds = np.array(preds_ts)
    Yp = np.array(Y_eval)
    g_write_np = np.array(g_write_ts)
    g_out_np = np.array(g_out_ts)

    def cos_sim(a, b, eps=1e-8):
        na = np.linalg.norm(a, axis=-1)
        nb = np.linalg.norm(b, axis=-1)
        return np.sum(a * b, axis=-1) / (na * nb + eps)

    cos = cos_sim(preds, Yp)

    t = np.arange(len(Y_eval))
    fig, ax = plt.subplots(4, 1, figsize=(10, 9))
    ax[0].set_title(f"Associative recall (eval loss={float(loss_eval):.4g})")
    ax[0].plot(np.arange(len(losses)), losses)
    ax[0].set_ylabel("train loss")

    for k in [0, 1, 2]:
        ax[1].plot(t, Yp[:, k], label=f"Y dim{k}", linewidth=2)
        ax[1].plot(t, preds[:, k], linestyle="--", label=f"pred dim{k}", alpha=0.8)
    ax[1].legend(loc="upper left", ncol=3)
    ax[1].set_ylabel("dims")

    ax[2].plot(t, cos)
    ax[2].set_ylabel("cos(pred,target)")
    ax[2].set_ylim(-0.05, 1.05)

    ax[3].plot(t, g_write_np, label="g write")
    ax[3].plot(t, g_out_np, label="g out")
    ax[3].set_ylabel("g")
    ax[3].legend(loc="best")
    ax[3].set_xlabel("timestep")
    ax[3].set_ylim(0, None)

    plt.tight_layout()
    plt.savefig("temp.png")
    plt.close('all')

    make_temp2_plot(
        model=model,
        params=params,
        X_all=np.array(X_all),
        ids_all=ids_all,
        token_table=token_table,
        num_tokens=num_tokens,
        n_holdout=200,
        seed=0,
        x_grid_G=250,
    )
    print("Saved temp2.png")



def _cos_sim(a, b, eps=1e-8):
    # a: (..., d), b: (..., d) or (d,)
    a = np.asarray(a)
    b = np.asarray(b)
    na = np.linalg.norm(a, axis=-1)
    nb = np.linalg.norm(b, axis=-1)
    return np.sum(a * b, axis=-1) / (na * nb + eps)


def kernel_overlap_legendre01(x1, x2, n_assoc):
    """
    Normalized kernel overlap:
      k(x1,x2) / sqrt(k(x1,x1)k(x2,x2)),
    where k(x,z) = phi(x)^T phi(z) for orthonormal basis.
    """
    phi1 = np.array(legendre_orthonormal_basis01(jnp.asarray(x1), n_assoc))
    phi2 = np.array(legendre_orthonormal_basis01(jnp.asarray(x2), n_assoc))
    k12 = float(np.dot(phi1, phi2))
    k11 = float(np.dot(phi1, phi1))
    k22 = float(np.dot(phi2, phi2))
    return k12 / (np.sqrt(k11 * k22) + 1e-8)


def rollout_collect_episode(model, params, X_ep, ids_ep):
    """
    Runs one episode and collects aux + memory states.

    Returns dict with:
      x_key:   (T,)
      x_query: (T,)
      g_write: (T,)
      g_out:   (T,)
      C_ts:    (T, d_model, n_assoc)   post-step C
    """
    # infer shapes from model config
    d_model = model.d_model
    n_hippo = model.n_hippo
    n_assoc = model.n_assoc

    S0 = jnp.zeros((d_model, n_hippo), dtype=jnp.float32)
    C0 = jnp.zeros((d_model, n_assoc), dtype=jnp.float32)

    def step_fn(carry, x_t):
        S, C = carry
        S_next, C_next, y_hat, aux = model.apply({"params": params}, x_t, S, C)
        return (S_next, C_next), (aux["x_key"], aux["x_query"], aux["g_write"], aux["g_out"], C_next)

    (Sf, Cf), (x_key_ts, x_query_ts, g_write_ts, g_out_ts, C_ts) = jax.lax.scan(step_fn, (S0, C0), X_ep)
    return {
        "x_key": np.array(x_key_ts),
        "x_query": np.array(x_query_ts),
        "g_write": np.array(g_write_ts),
        "g_out": np.array(g_out_ts),
        "C_ts": np.array(C_ts),
        "ids": np.array(ids_ep),
    }


def find_target_write_time(ids_ep, num_tokens):
    """
    For this task structure:
      - WRITE token id is 2*num_tokens
      - final A token is at time T-2 (even index)
      - find last earlier occurrence of that A token among even indices < T-2
      - write time is the B immediately after that earlier A: t_write = tA + 1
    Returns: (t_write, a_id, b_vocab_id)
      where b_vocab_id is in [num_tokens, 2*num_tokens-1]
    """
    WRITE_ID = 2 * num_tokens
    T = len(ids_ep)
    assert ids_ep[-1] == WRITE_ID, "expected WRITE at final position"

    a_id = int(ids_ep[T - 2])  # final A token id in [0,num_tokens-1]
    A_times = np.arange(0, T - 1, 2)  # even positions excluding final WRITE
    A_ids = ids_ep[A_times]

    # find last earlier occurrence among A_times excluding last A
    matches = np.where(A_ids[:-1] == a_id)[0]
    if len(matches) == 0:
        return None  # should not happen if dataset generator enforced it
    j_last = int(matches[-1])
    tA = int(A_times[j_last])
    t_write = tA + 1
    b_vocab = int(ids_ep[t_write])
    return t_write, a_id, b_vocab


def decode_memory_curve(C_bank, out_kernel, out_bias, x_grid):
    """
    C_bank: (d_model, n_assoc) associative memory banks at some time
    out_kernel: (d_model, d_in)
    out_bias: (d_in,)
    x_grid: (G,) in [0,1]

    Returns:
      Y_mem: (G, d_in) decoded token-space vector as a function of x
    """
    d_model, n_assoc = C_bank.shape
    G = len(x_grid)

    # build Phi(x): (G, n_assoc)
    Phi = np.array(legendre_orthonormal_basis01(jnp.asarray(x_grid), n_assoc))  # (G, n_assoc)

    # r(x) = C_bank @ phi(x): (d_model,) for each x => (G, d_model)
    R = Phi @ C_bank.T  # (G, d_model)

    # out_proj: y = R @ W + b  => (G, d_in)
    Y = R @ out_kernel + out_bias[None, :]
    return Y


# ----------------------------
# Diagnostics on holdout episodes
# ----------------------------
def make_temp2_plot(
    model,
    params,
    X_all,
    ids_all,
    token_table,
    num_tokens: int,
    n_holdout: int = 200,
    seed: int | None = None,
    x_grid_G: int = 200,
):
    rng = np.random.default_rng(seed)
    N, T, d_in = X_all.shape

    # Choose holdout indices: here we just sample from the *end* chunk of the dataset.
    # If you have a true split, replace this block.
    holdout_pool = np.arange(int(0.8 * N), N)
    sel = rng.choice(holdout_pool, size=min(n_holdout, len(holdout_pool)), replace=False)

    # Collect per-episode (read_addr, write_addr, a_id, b_vocab)
    read_addrs = []
    write_addrs = []
    a_ids = []
    b_vocabs = []
    per_ep_rollouts = []  # keep a few for plot 3 selection

    for idx in sel:
        X_ep = X_all[idx]
        ids_ep = ids_all[idx]
        # rollout
        roll = rollout_collect_episode(model, params, X_ep, ids_ep)
        per_ep_rollouts.append(roll)

        # identify write time corresponding to the retrieved association
        out = find_target_write_time(roll["ids"], num_tokens=num_tokens)
        if out is None:
            continue
        t_write, a_id, b_vocab = out

        # read address at WRITE time
        t_read = T - 1
        x_read = float(roll["x_query"][t_read])

        # write address at the *write time* (where the relevant B appears)
        x_write = float(roll["x_key"][t_write])

        read_addrs.append(x_read)
        write_addrs.append(x_write)
        a_ids.append(a_id)
        b_vocabs.append(b_vocab)

    read_addrs = np.array(read_addrs)
    write_addrs = np.array(write_addrs)
    a_ids = np.array(a_ids, dtype=np.int32)
    b_vocabs = np.array(b_vocabs, dtype=np.int32)

    # ----------------------------
    # Plot 1: scatter read vs write addresses, colored by A token
    # ----------------------------
    fig, axarr = plt.subplots(3, 1, figsize=(10, 14))
    ax1, ax2, ax3 = axarr

    cmap = plt.get_cmap("tab10")
    colors = [cmap(i % 10) for i in range(num_tokens)]

    for a in range(num_tokens):
        m = (a_ids == a)
        if np.any(m):
            ax1.scatter(write_addrs[m], read_addrs[m], s=20, alpha=0.8, color=colors[a], label=f"A{a}")

    ax1.set_title("1) Read vs write address for the retrieved association (colored by A token)")
    ax1.set_xlabel("write address x_key at relevant [a;b] time")
    ax1.set_ylabel("read address x_query at [a;WRITE] time")
    ax1.set_xlim(-0.02, 1.02)
    ax1.set_ylim(-0.02, 1.02)
    ax1.plot([0,1], [0,1], c='k', alpha=0.5, ls='--')
    ax1.grid(True, alpha=0.2)
    ax1.legend(loc="upper right", ncol=2, fontsize=9)

    # # ----------------------------
    # Plot 2: kernel overlap between mean write addresses per A token + null band
    # ----------------------------
    mean_write_addr = np.full((num_tokens,), np.nan, dtype=np.float32)
    for a in range(num_tokens):
        m = (a_ids == a)
        if np.any(m):
            mean_write_addr[a] = float(np.mean(write_addrs[m]))

    # Pair overlaps for A tokens
    pair_overlaps = []
    pair_labels = []
    for i in range(num_tokens):
        for j in range(i + 1, num_tokens):
            if np.isfinite(mean_write_addr[i]) and np.isfinite(mean_write_addr[j]):
                ov = kernel_overlap_legendre01(mean_write_addr[i], mean_write_addr[j], n_assoc=model.n_assoc)
                pair_overlaps.append(ov)
                pair_labels.append((i, j))
    pair_overlaps = np.array(pair_overlaps, dtype=np.float32)

    # ----------------------------
    # Null KDE of max |overlap| across num_tokens random addresses
    # ----------------------------
    def max_abs_pair_overlap_from_addresses(xs: jnp.ndarray, n_assoc: int) -> jnp.ndarray:
        """
        xs: (M,) addresses in [0,1]
        returns: scalar max_{i<j} | <phi(xi),phi(xj)> / (||phi(xi)|| ||phi(xj)||) |
        """
        Phi = legendre_orthonormal_basis01(xs, n_assoc)              # (M, n_assoc)
        Phi = Phi / (jnp.linalg.norm(Phi, axis=1, keepdims=True) + 1e-8)
        G = Phi @ Phi.T                                              # (M, M), diag ~ 1
        M = xs.shape[0]
        # upper triangle mask without diag
        mask = jnp.triu(jnp.ones((M, M), dtype=G.dtype), k=1)
        return jnp.max(jnp.abs(G) * mask)

    max_abs_pair_overlap_from_addresses_jit = jax.jit(
        max_abs_pair_overlap_from_addresses, static_argnames=("n_assoc",)
    )

    def null_max_abs_overlap_distribution(
        key: jax.Array,
        num_tokens: int,
        n_assoc: int,
        n_trials: int,
    ) -> jnp.ndarray:
        """
        Draw n_trials sets of num_tokens addresses ~ Uniform[0,1],
        return distribution of max abs pair overlap. Shape (n_trials,).
        """
        # sample all addresses at once: (n_trials, num_tokens)
        U = jr.uniform(key, shape=(n_trials, num_tokens), minval=0.0, maxval=1.0)

        # vectorize across trials
        f = lambda xs: max_abs_pair_overlap_from_addresses_jit(xs, n_assoc=n_assoc)
        return jax.vmap(f)(U)  # (n_trials,)


    # ---- observed statistic from mean write addresses (computed earlier)
    # mean_write_addr: (num_tokens,) with NaNs possible if token missing
    valid = np.isfinite(mean_write_addr)
    xs_obs = mean_write_addr[valid].astype(np.float32)
    m_eff = xs_obs.shape[0]

    if m_eff < 2:
        ax2.set_title("2) Not enough observed tokens with write addresses to compute overlaps")
    else:
        # observed max abs overlap
        T_obs = float(
            max_abs_pair_overlap_from_addresses_jit(jnp.asarray(xs_obs), n_assoc=int(model.n_assoc))
        )

        # null distribution
        n_trials = 1000  # increase if you want smoother KDE
        key_null = jr.PRNGKey(123)  # or fold in your main key/seed
        T_null = np.array(
            null_max_abs_overlap_distribution(
                key_null, num_tokens=m_eff, n_assoc=int(model.n_assoc), n_trials=n_trials
            ),
            dtype=np.float32
        )

        # ---- KDE plot (SciPy if available; else histogram)
        ax2.cla()

        from scipy.stats import gaussian_kde

        kde = gaussian_kde(T_null)
        xs = np.linspace(0.0, max(1e-3, float(np.max(T_null)) * 1.05), 400)
        ys = kde(xs)

        ax2.plot(xs, ys, linewidth=2.5, label=f"Null KDE of max |overlap| (n={n_trials})")
        ax2.fill_between(xs, 0.0, ys, alpha=0.25)
        ax2.axvline(T_obs, linewidth=3.0, linestyle="--", label=f"Observed max |overlap| = {T_obs:.3f}")

        # useful summary stats
        p = float(np.mean(T_null >= T_obs))
        q95 = float(np.quantile(T_null, 0.95))
        q99 = float(np.quantile(T_null, 0.99))

        ax2.set_title(
            "2) Max pairwise |kernel overlap| among A-token mean write addresses\n"
            f"null >= observed fraction p~={p:.3f}   (null 95%={q95:.3f}, 99%={q99:.3f})"
        )
        ax2.set_xlabel("max_{i<j} |normalized kernel overlap|")
        ax2.set_ylabel("density")
        ax2.set_xlim(left=0.0)
        ax2.grid(True, alpha=0.2)
        ax2.legend(loc="upper right", fontsize=9)

    # ----------------------------
    # Plot 3: before/after write: cosine similarity curves for each B token
    # ----------------------------
    # Choose one episode with a confident write (largest g_write at target write time)
    best = None
    best_score = -1.0
    best_info = None

    for roll in per_ep_rollouts:
        out = find_target_write_time(roll["ids"], num_tokens=num_tokens)
        if out is None:
            continue
        t_write, a_id, b_vocab = out
        score = float(roll["g_write"][t_write])
        if score > best_score:
            best_score = score
            best = roll
            best_info = (t_write, a_id, b_vocab)

    if best is None:
        ax3.set_title("3) (Could not find valid episode for before/after write plot)")
    else:
        t_write, a_id, b_vocab = best_info

        # We have C_ts as post-step. Approximate:
        #  - "before" write: C just before processing t_write -> use C_ts[t_write-1]
        #  - "after"  write: C after processing t_write -> use C_ts[t_write]
        # If t_write==0 (shouldn't happen), fallback to t_write.
        C_before = best["C_ts"][max(t_write - 1, 0)]
        C_after  = best["C_ts"][t_write]

        # out_proj params
        W_out = np.array(params["out_proj"]["kernel"])  # (d_model, d_in)
        b_out = np.array(params["out_proj"]["bias"])    # (d_in,)

        x_grid = np.linspace(0.0, 1.0, x_grid_G, dtype=np.float32)

        Y_before = decode_memory_curve(C_before, W_out, b_out, x_grid)  # (G, d_in)
        Y_after  = decode_memory_curve(C_after,  W_out, b_out, x_grid)  # (G, d_in)

        # B token vectors in token_table: vocab ids num_tokens..2*num_tokens-1
        # We'll plot cos( memory(x), token_b ) as a function of x.
        for b_id in range(num_tokens):
            vocab_id = num_tokens + b_id
            v = token_table[vocab_id]  # (d_in,)
            c_before = _cos_sim(Y_before, v[None, :])
            c_after  = _cos_sim(Y_after,  v[None, :])

            # Plot after as solid, before as dashed (light)
            ax3.plot(x_grid, c_after, linewidth=2.0, alpha=0.85, label=f"B{b_id}" if num_tokens <= 10 else None)
            ax3.plot(x_grid, c_before, linewidth=1.2, alpha=0.35, linestyle="--")

        # Mark the write address at t_write (where b_vocab was observed) for reference
        xw = float(best["x_key"][t_write])
        ax3.axvline(xw, linestyle=":", linewidth=2.0, alpha=0.8)

        written_b = int(b_vocab - num_tokens)  # in 0..num_tokens-1
        ax3.set_title(
            "3) Before (dashed) vs After (solid) write: cos(memory(x), B-token)\n"
            f"selected episode: g_write@t_write={best_score:.3f}, wrote B{written_b} at x_key={xw:.3f}"
        )
        ax3.set_xlabel("x in [0,1] (OP address)")
        ax3.set_ylabel("cosine similarity")
        # ax3.set_ylim(-0.05, 1.05)
        ax3.grid(True, alpha=0.2)
        if num_tokens <= 10:
            ax3.legend(loc="upper right", ncol=2, fontsize=9)

    plt.tight_layout()
    plt.savefig("temp2.png", dpi=200)
    plt.close(fig)
    return mean_write_addr



@jax.jit(static_argnames=("n_assoc",))
def null_kernel_overlaps_legendre01(u1: jnp.ndarray, u2: jnp.ndarray, n_assoc: int) -> jnp.ndarray:
    """
    Vectorized (Option A) null kernel overlaps for orthonormal Legendre basis on [0,1].

    Computes normalized overlaps:
      ov[i] = <phi(u1[i]), phi(u2[i])> / (||phi(u1[i])|| * ||phi(u2[i])||)

    Args:
      u1, u2: shape (N,) in [0,1]
      n_assoc: truncation/order (python int is fine)

    Returns:
      overlaps: shape (N,) float32
    """
    u1 = jnp.asarray(u1)
    u2 = jnp.asarray(u2)

    # Phi1, Phi2: (N, n_assoc)
    Phi1 = legendre_orthonormal_basis01(u1, n_assoc)
    Phi2 = legendre_orthonormal_basis01(u2, n_assoc)

    # Dot products and norms: (N,)
    dot12 = jnp.sum(Phi1 * Phi2, axis=-1)
    n1 = jnp.sqrt(jnp.sum(Phi1 * Phi1, axis=-1) + 1e-8)
    n2 = jnp.sqrt(jnp.sum(Phi2 * Phi2, axis=-1) + 1e-8)

    return (dot12 / (n1 * n2)).astype(jnp.float32)



@jax.jit(static_argnames=("n_assoc", "grid_n"))
def kernel_overlap_grid_legendre01(
    n_assoc: int,
    grid_n: int,
) -> jnp.ndarray:
    """
    Compute pairwise normalized kernel overlap on a uniform grid in [0,1].

    Args:
      n_assoc: truncation/order of orthonormal Legendre basis
      grid_n:  number of grid points (includes endpoints 0 and 1)

    Returns:
      K: (grid_n, grid_n) array with
         K[i,j] = <phi(x_i), phi(x_j)> /
                  (||phi(x_i)|| ||phi(x_j)||)
    """
    # Uniform grid including endpoints
    x = jnp.linspace(0.0, 1.0, grid_n)

    # Basis evaluations: Phi[i,k] = p_k(x_i)
    Phi = legendre_orthonormal_basis01(x, n_assoc)  # (grid_n, n_assoc)

    # Normalize basis vectors
    Phi_norm = Phi / (jnp.linalg.norm(Phi, axis=1, keepdims=True) + 1e-8)

    # Pairwise normalized overlaps
    K = Phi_norm @ Phi_norm.T  # (grid_n, grid_n)

    return K.astype(jnp.float32)


ARR = [
    0.45298209190368655,
    0.21947033051401377,
    0.5032657034256879,
    0.2627926245331764,
    0.5524154046307439,
    0.4039508490001454,
    0.6963000237941742,
    0.3547567844390869,
    0.7415585688182286,
    0.3075451672077179,
    0.6019835743037137,
    0.6492930816279517,
] # Empirical mean write & read addresses

def make_kernel_plot(mean_write_addr=ARR):
    grid = kernel_overlap_grid_legendre01(32, 256)
    grid = np.clip(np.array(grid), -1, 1)
    print(np.min(grid), np.max(grid))
    xx, yy = np.meshgrid(mean_write_addr, mean_write_addr, indexing='xy')
    plt.scatter(xx, yy, marker='+', c='k', s=9.0, alpha=0.7)
    plt.imshow(grid, vmin=-1, vmax=1, extent=(0,1,0,1), cmap='bwr', origin='lower')
    plt.xlabel(r'Address $x_{query}$')
    plt.ylabel(r'Address $x_{key}$')
    plt.title(r'OP Memory Kernel $K(x_{key}, x_{query})$ ($n_{assoc}=32$)')
    plt.colorbar()
    plt.savefig('op_memory_kernel.png')
    plt.savefig('op_memory_kernel.pdf')
    plt.close('all')



if __name__ == "__main__":
    mean_write_addr = main()
    # make_kernel_plot()