source_code/src/mpm/synthetic.py

"""Synthetic datasets for quick experiments."""
from __future__ import annotations

import numpy as np


def make_induction_head_task(n: int = 100, T: int = 30, V: int = 20, seed: int = 42):
    np.random.seed(seed)
    perm = np.random.permutation(V)

    X = np.zeros((n, T), dtype=int)
    Y = np.zeros(n, dtype=int)

    X[:, 0] = np.random.randint(0, V, n)
    for i in range(n):
        for j in range(1, T - 2):
            if X[i, j - 1] == 0:
                X[i, j] = np.random.randint(1, V)
            else:
                X[i, j] = np.random.randint(0, V)
    X[:, -2] = np.random.randint(1, V, n)
    X[:, -1] = 0

    contains_zero = np.any(X[:, :-1] == 0, axis=1)
    missing_zero_indices = np.where(~contains_zero)[0]
    if len(missing_zero_indices) > 0:
        random_indices = np.random.randint(0, T - 2, size=len(missing_zero_indices))
        X[missing_zero_indices, random_indices] = 0

    for i in range(n):
        Y[i] = X[i][np.argwhere(X[i, :-1] == 0).flatten()[-1] + 1]

    vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V)
    vocab = vocab[perm]
    return vocab[X], vocab[Y]


def make_induction_head_seq_to_seq_task(
    n: int = 1, T: int = 1000, V: int = 20, seed: int = 42
):
    np.random.seed(seed)
    perm = np.random.permutation(V)

    X = np.zeros((n, T), dtype=int)
    Y = np.zeros((n, T), dtype=int)

    X[:, 0] = np.random.randint(0, V, n)
    for i in range(n):
        for j in range(1, T - 2):
            if X[i, j - 1] == 0:
                X[i, j] = np.random.randint(1, V)
            else:
                X[i, j] = np.random.randint(0, V)
    X[:, -2] = np.random.randint(1, V, n)
    X[:, -1] = 0

    contains_zero = np.any(X[:, :-1] == 0, axis=1)
    missing_zero_indices = np.where(~contains_zero)[0]
    if len(missing_zero_indices) > 0:
        random_indices = np.random.randint(0, T - 2, size=len(missing_zero_indices))
        X[missing_zero_indices, random_indices] = 0

    for i in range(n):
        current_val = 0
        flag = False
        for t in range(T):
            if X[i, t] == 0:
                flag = True
            elif flag:
                current_val = X[i, t]
                flag = False
            Y[i, t] = current_val

    vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V)
    vocab[0] *= 2
    vocab = vocab[perm]
    return vocab[X], vocab[Y]


def get_rbf_sample(
    T: int,
    *,
    dt: float = 1.0,
    lengthscale: float = 10.0,
    sigma: float = 1.0,
    mu: float = 0.0,
    seed: int | None = None,
    embed_factor: int = 2,
    jitter: float = 1e-12,
    clip_neg_eigs: bool = True,
) -> np.ndarray:
    """
    Efficiently sample from a 1D zero-mean GP on an evenly spaced grid using an RBF kernel,
    via circulant embedding + FFT (O(M log M)).

    GP prior on indices {0,...,T-1} with covariance:
        k(i,j) = sigma^2 * exp(-( (i-j)*dt )^2 / (2*lengthscale^2))

    Args:
        T: number of samples
        dt: grid spacing
        lengthscale: RBF lengthscale ℓ (>0)
        sigma: marginal std (>=0)
        mu: mean
        seed: RNG seed
        embed_factor: embedding size multiplier; M is next power of 2 >= embed_factor*T.
                      2 is standard; increase (e.g. 4) if negative eigenvalues occur.
        jitter: small diagonal jitter added to k(0) to help numerical stability
        clip_neg_eigs: if True, clip tiny negative FFT eigenvalues to 0.

    Returns:
        z: (T,) sample
    """
    if T <= 0:
        return np.zeros((0,), dtype=np.float64)
    if dt <= 0:
        raise ValueError("dt must be > 0")
    if lengthscale <= 0:
        raise ValueError("lengthscale must be > 0")
    if sigma < 0:
        raise ValueError("sigma must be >= 0")
    if embed_factor < 2:
        raise ValueError("embed_factor should be >= 2 for circulant embedding")

    rng = np.random.default_rng(seed)

    # Choose embedding size M (power of 2 for fast FFT).
    M_min = embed_factor * T
    M = 1 << (M_min - 1).bit_length()
    if M % 2 != 0:
        M += 1  # keep even for clean Nyquist handling

    # Build Toeplitz first column for size T: k[0..T-1]
    d = np.arange(T, dtype=np.float64) * dt
    k_col = (sigma ** 2) * np.exp(-0.5 * (d / lengthscale) ** 2)
    k_col[0] += jitter

    # Circulant embedding first column c of length M:
    # c = [k0, k1, ..., k_{T-1}, 0, 0, ..., 0, k_{T-1}, ..., k1]
    c = np.zeros(M, dtype=np.float64)
    c[:T] = k_col
    c[M - (T - 1) :] = k_col[1:][::-1]  # tail mirror (exclude k0)

    # Eigenvalues of the circulant covariance matrix (should be >= 0)
    lam = np.real(np.fft.fft(c))
    if clip_neg_eigs:
        lam = np.maximum(lam, 0.0)
    else:
        if np.min(lam) < -1e-10:
            raise ValueError(
                f"Circulant embedding not PSD (min eigenvalue {np.min(lam)}). "
                "Increase embed_factor or enable clip_neg_eigs."
            )
        lam = np.maximum(lam, 0.0)

    # Sample in Fourier domain with conjugate symmetry so the time series is real.
    # numpy FFT conventions: ifft includes 1/M normalization.
    # To match covariance C = (1/M) F^H diag(lam) F, use:
    #   x = sqrt(M) * ifft( sqrt(lam) * z ),  with z ~ CN(0, I) and conjugate symmetry.
    Z = np.zeros(M, dtype=np.complex128)

    # k=0 (DC) and k=M/2 (Nyquist) are real-valued in the symmetric FFT.
    Z[0] = rng.normal()
    Z[M // 2] = rng.normal()

    # Positive frequencies 1..M/2-1: complex normals with Var=1 (CN(0,1))
    re = rng.normal(size=(M // 2 - 1))
    im = rng.normal(size=(M // 2 - 1))
    Z[1 : M // 2] = (re + 1j * im) / np.sqrt(2.0)
    # Enforce conjugate symmetry
    Z[M // 2 + 1 :] = np.conj(Z[1 : M // 2][::-1])

    # Scale by sqrt eigenvalues
    Y = np.sqrt(lam) * Z

    # Back to time domain, take real part
    x_full = (np.sqrt(M) * np.fft.ifft(Y)).real

    return mu + x_full[:T]


def rbf_kernel_1d(i: int, j: int, *, sigma: float, lengthscale: float, dt: float) -> float:
    d = (i - j) * dt
    return (sigma ** 2) * np.exp(-0.5 * (d / lengthscale) ** 2)


def oracle_mse_rbf_periodic_censor(
    P: int,
    L: int,
    *,
    sigma: float = 1.0,
    lengthscale: float = 10.0,
    dt: float = 1.0,
    num_periods: int = 50,
    burn_periods: int = 10,
    jitter: float = 1e-10,
) -> dict:
    """
    Oracle MSE for the censored-copy task when Z_t ~ GP(mean, RBF kernel).

    Censoring pattern (matches ``make_censored_task``):
      observed at index i iff (i % P) < (P//2).
    But due to the task's lagging, we only ever observe indices >= L (because X is input_signal[L:]).

    At model time t, the oracle conditions on all observed indices <= (t+L), i.e. fixed-lag smoothing.

    Returns:
      {
        "mse": (T,) oracle per-timestep MSE trace,
        "avg_mse_last_period": scalar average over the last period (after burn-in),
        "avg_mse_post_burn": scalar average over all timesteps after burn-in,
      }
    """
    if P <= 0 or P % 2 != 0:
        raise ValueError("P must be a positive even integer.")
    if L < 0:
        raise ValueError("L must be >= 0")
    if sigma < 0 or lengthscale <= 0 or dt <= 0:
        raise ValueError("Require sigma>=0, lengthscale>0, dt>0")
    if num_periods <= 0:
        raise ValueError("num_periods must be > 0")
    if burn_periods < 0 or burn_periods >= num_periods:
        raise ValueError("burn_periods must be in [0, num_periods-1]")

    m = P // 2
    T = num_periods * P               # number of Y_t evaluated
    N = T + L                         # original time indices potentially observable up to t+L
    burn_T = burn_periods * P

    # Observation availability on original time axis
    # Only indices >= L are observable via X (since X corresponds to original indices L..L+T-1).
    observed = np.zeros(N, dtype=bool)
    for i in range(L, N):
        observed[i] = (i % P) < m

    # Incremental Cholesky factor of K_obs (lower-triangular)
    obs_idx: list[int] = []
    chol_L = np.zeros((0, 0), dtype=np.float64)

    def add_observation(i_new: int):
        """Add new observed index to the Cholesky factor (noise-free GP with jitter)."""
        nonlocal chol_L, obs_idx
        if len(obs_idx) == 0:
            k_nn = rbf_kernel_1d(i_new, i_new, sigma=sigma, lengthscale=lengthscale, dt=dt) + jitter
            chol_L = np.array([[np.sqrt(k_nn)]], dtype=np.float64)
            obs_idx.append(i_new)
            return

        # k between new point and existing obs
        k_vec = np.array(
            [rbf_kernel_1d(i_new, j, sigma=sigma, lengthscale=lengthscale, dt=dt) for j in obs_idx],
            dtype=np.float64,
        )  # (M,)

        # Solve L w = k_vec
        w = np.linalg.solve(chol_L, k_vec)  # (M,)
        k_nn = rbf_kernel_1d(i_new, i_new, sigma=sigma, lengthscale=lengthscale, dt=dt) + jitter
        diag_sq = k_nn - float(w @ w)
        diag = np.sqrt(max(diag_sq, jitter))

        # Build expanded Cholesky
        M = len(obs_idx)
        L_new = np.zeros((M + 1, M + 1), dtype=np.float64)
        L_new[:M, :M] = chol_L
        L_new[M, :M] = w
        L_new[M, M] = diag

        chol_L = L_new
        obs_idx.append(i_new)

    def posterior_var(test_t: int) -> float:
        """Var(Z_test_t | observed indices obs_idx), using current Cholesky."""
        k_tt = rbf_kernel_1d(test_t, test_t, sigma=sigma, lengthscale=lengthscale, dt=dt)

        if len(obs_idx) == 0:
            return k_tt

        if obs_idx[-1] == test_t or test_t in obs_idx:
            return 0.0

        k_tO = np.array(
            [rbf_kernel_1d(test_t, j, sigma=sigma, lengthscale=lengthscale, dt=dt) for j in obs_idx],
            dtype=np.float64,
        )
        # alpha = L^{-1} k_tO
        alpha = np.linalg.solve(chol_L, k_tO)
        var = k_tt - float(alpha @ alpha)
        return max(var, 0.0)

    # Main loop: advance horizon h, and output MSE for t = h-L once h>=L
    mse = np.zeros(T, dtype=np.float64)

    for h in range(N):
        if observed[h]:
            add_observation(h)

        if h >= L:
            t = h - L
            if t < T:
                mse[t] = posterior_var(t)

    # Averages (useful “oracle loss” scalars)
    post_burn = mse[burn_T:] if burn_T < T else mse
    avg_post_burn = float(np.mean(post_burn)) if post_burn.size else float(np.mean(mse))

    return avg_post_burn


def get_ou_sample(
    T: int,
    dt: float = 1.0,
    tau: float = 10.0,
    sigma: float = 1.0,
    mu: float = 0.0,
    seed: int | None = None,
) -> np.ndarray:
    """
    Sample a stationary Ornstein–Uhlenbeck process at discrete times.

    Continuous-time OU (one common parametrization):
        dX_t = -(1/tau) (X_t - mu) dt + sigma dW_t

    Discretization (exact transition):
        X_{t+dt} = mu + rho (X_t - mu) + eps
        rho = exp(-dt/tau)
        eps ~ N(0, q),  q = (sigma^2 * tau / 2) * (1 - rho^2)

    Stationary distribution:
        X_t ~ N(mu, sigma^2 * tau / 2)

    Args:
        T: number of samples to return
        dt: sampling interval
        tau: relaxation time constant (> 0)
        sigma: diffusion scale (>= 0)
        mu: mean
        seed: RNG seed

    Returns:
        x: (T,) numpy array
    """
    if T <= 0:
        return np.zeros((0,), dtype=np.float64)
    if tau <= 0:
        raise ValueError("tau must be > 0")
    if sigma < 0:
        raise ValueError("sigma must be >= 0")

    rng = np.random.default_rng(seed)

    rho = np.exp(-dt / tau)
    var_stationary = (sigma**2) * tau / 2.0

    # Exact conditional variance for step dt
    q = var_stationary * (1.0 - rho**2)

    x = np.empty((T,), dtype=np.float64)
    x[0] = mu + np.sqrt(var_stationary) * rng.standard_normal()

    if T > 1:
        noise = np.sqrt(q) * rng.standard_normal(size=T - 1)
        for t in range(T - 1):
            x[t + 1] = mu + rho * (x[t] - mu) + noise[t]

    return x


def oracle_mse_censored_task(
    P: int,
    L: int,
    *,
    dt: float = 1.0,
    tau: float = 10.0,
    sigma: float = 1.0,
) -> float:
    """
    Oracle (minimum expected) per-timestep MSE for the censored-copy task
    under a stationary OU process, with periodic censoring pattern:

      - Period length: P
      - First half of each period: uncensored (perfect observation => MSE=0)
      - Second half: censored (missing observations)
      - Lookahead (in original time index units): L  (the task lag)

    Assumes the oracle uses all uncensored samples up to time t+L to predict Z_t.
    Because OU is Gaussian Markov (AR(1)), the per-step conditional variance is:

      If future endpoint u is NOT yet observed:
        Var(Z_t | Z_s) = v * (1 - rho^(2 d1))

      If future endpoint u IS observed:
        Var(Z_t | Z_s, Z_u) = v * (1 - a^2 - b^2 + 2ab rho^D) / (1 - rho^(2D)),
        which simplifies to:
        v * (1 - rho^(2 d1) - rho^(2 d2) + rho^(2D)) / (1 - rho^(2D)).

    Here:
      rho = exp(-dt/tau)
      v   = stationary variance = sigma^2 * tau / 2
      m   = P/2 (must be integer; require even P)
      s   = last uncensored time in the period, u = first uncensored time next period
      D   = u - s = m + 1
      For censored positions: k = 1..m with d1 = k and d2 = D - k.

    Returns:
      Long-run average MSE per timestep (averaged over one period).
    """
    if P <= 0:
        raise ValueError("P must be positive")
    if P % 2 != 0:
        raise ValueError("This oracle formula assumes even P (because censoring uses P//2).")
    if L < 0:
        raise ValueError("L must be >= 0")
    if tau <= 0:
        raise ValueError("tau must be > 0")
    if sigma < 0:
        raise ValueError("sigma must be >= 0")

    m = P // 2  # censored block length
    rho = np.exp(-dt / tau)
    v = (sigma**2) * tau / 2.0

    D = m + 1  # distance between last uncensored and next uncensored in this pattern
    rho2D = rho ** (2 * D)
    denom = 1.0 - rho2D

    mse_sum = 0.0

    # Uncensored half contributes 0, so only sum over censored half (k=1..m).
    for k in range(1, m + 1):
        d1 = k
        d2 = D - k

        if d2 > L:
            # No observed sample after the censor block yet (given lookahead L)
            mse_k = v * (1.0 - rho ** (2 * d1))
        else:
            # Smoothing with both endpoints available
            # v * (1 - rho^(2 d1) - rho^(2 d2) + rho^(2D)) / (1 - rho^(2D))
            mse_k = v * (1.0 - rho ** (2 * d1) - rho ** (2 * d2) + rho2D) / denom

        mse_sum += mse_k

    # Average over all P timesteps in a period
    return mse_sum / P


def oracle_info_gain(
    X: np.ndarray,
    *,
    lag: int,
    censor_val: float,
    dt: float = 1.0,
    tau: float = 10.0,
    sigma: float = 1.0,
    mu: float = 0.0,
    atol: float = 0.0,
    eps: float = 1e-12,
) -> np.ndarray:
    """
    Oracle information-gain (salience) trace for the censored OU task.

    info_gain[t] = 0.5 * log(
        Var(Z_t | obs <= t+lag-1) / Var(Z_t | obs <= t+lag)
    )

    This measures how much the *newest* input sample at time (t+lag)
    reduces uncertainty about the current target Z_t.

    Args:
        X: (T,) input signal (censored OU samples)
        lag: fixed lag used in the task
        censor_val: value used to mark censored samples
        dt, tau, sigma, mu: OU parameters
        atol: optional tolerance for detecting censor_val
        eps: numerical stability

    Returns:
        info_gain: (T,) oracle salience / information-gain trace
    """
    X = np.asarray(X, dtype=np.float64)
    T = X.shape[0]

    rho = np.exp(-dt / tau)
    v = (sigma**2) * tau / 2.0  # stationary variance

    # Original-time indexing
    N = T + lag

    observed = np.zeros(N, dtype=bool)
    obs_value = np.zeros(N, dtype=np.float64)

    if atol > 0:
        obs_mask = np.abs(X - censor_val) > atol
    else:
        obs_mask = X != censor_val

    ks = lag + np.arange(T)
    observed[ks] = obs_mask
    obs_value[ks] = X

    # last observed <= k
    last_obs_leq = np.full(N, -1, dtype=int)
    last = -1
    for k in range(N):
        if observed[k]:
            last = k
        last_obs_leq[k] = last

    # next observed >= k
    next_obs_geq = np.full(N, N, dtype=int)
    nxt = N
    for k in range(N - 1, -1, -1):
        if observed[k]:
            nxt = k
        next_obs_geq[k] = nxt

    def posterior_var(t: int, horizon: int) -> float:
        """Var(Z_t | uncensored obs with indices <= horizon)."""
        if horizon < 0:
            return v

        horizon = min(horizon, N - 1)

        s = last_obs_leq[min(t, horizon)] if t >= 0 else -1
        u = next_obs_geq[t] if t <= horizon and next_obs_geq[t] <= horizon else N

        if s == t and s != -1:
            return 0.0

        if s == -1 and u == N:
            return v
        if s == -1:
            d = u - t
            return v * (1.0 - rho ** (2 * d))
        if u == N:
            d = t - s
            return v * (1.0 - rho ** (2 * d))

        d1 = t - s
        d2 = u - t
        D = d1 + d2
        rho2D = rho ** (2 * D)
        return v * (1.0 - rho ** (2 * d1) - rho ** (2 * d2) + rho2D) / max(
            1.0 - rho2D, eps
        )

    info_gain = np.zeros(T, dtype=np.float64)

    for t in range(T):
        var_prev = posterior_var(t, t + lag - 1)
        var_post = posterior_var(t, t + lag)
        info_gain[t] = 0.5 * np.log((var_prev + eps) / (var_post + eps))

    return info_gain




def oracle_info_gain_rbf(
    X: np.ndarray,
    *,
    lag: int,
    censor_val: float,
    dt: float = 1.0,
    lengthscale: float = 10.0,
    sigma: float = 1.0,
    atol: float = 0.0,
    jitter: float = 1e-10,
    eps: float = 1e-12,
    max_obs: int | None = None,
) -> np.ndarray:
    """
    Oracle information-gain trace for the censored task under an RBF-kernel GP prior.

    Prior:
      Z ~ GP(mu, k),  k(i,j)=sigma^2 * exp(-((i-j)*dt)^2/(2*ell^2))
    Observations:
      Noise-free: observe Z_k exactly at uncensored k; censored => missing.

    At model time t:
      horizon h = t + lag (original-time index of current input X[t])
      IG_t = 0.5 * log( Var(Z_t | obs<=h-1) / Var(Z_t | obs<=h) )

    Args:
        X: (T,) input stream; X[t] is either Z_{t+lag} (uncensored) or censor_val (censored)
        lag: task lag L
        censor_val: sentinel for missing observations
        dt, lengthscale, sigma: RBF GP kernel params
        atol: optional tolerance for detecting censor_val
        jitter: diagonal jitter for numerical stability
        eps: stability for log ratio
        max_obs: cap number of retained observed points (approximation). If None, keeps all.

    Returns:
        info_gain: (T,) oracle IG trace
    """
    if max_obs is None:
        max_obs = int(lag + 3 * lengthscale / dt)
    X = np.asarray(X, dtype=np.float64)
    T = X.shape[0]
    if lag < 0:
        raise ValueError("lag must be >= 0")
    if dt <= 0 or lengthscale <= 0:
        raise ValueError("dt and lengthscale must be > 0")
    if sigma < 0:
        raise ValueError("sigma must be >= 0")

    # Original-time axis indices potentially involved: 0..N-1 where N=T+lag
    N = T + lag

    # Which original-time indices are observed (only k=lag..lag+T-1 are ever presented via X)
    observed = np.zeros(N, dtype=bool)
    if atol > 0:
        obs_mask = np.abs(X - censor_val) > atol
    else:
        obs_mask = X != censor_val
    ks = lag + np.arange(T)
    observed[ks] = obs_mask

    # RBF kernel helpers
    inv_ell2 = 1.0 / (lengthscale * lengthscale)
    sigma2 = sigma * sigma

    def k_vec(t_idx: int, obs_idx: np.ndarray) -> np.ndarray:
        d = (t_idx - obs_idx).astype(np.float64) * dt
        return sigma2 * np.exp(-0.5 * (d * d) * inv_ell2)

    def k_tt(_: int) -> float:
        return sigma2  # RBF has k(t,t)=sigma^2

    # Maintain Cholesky of K_obs (lower triangular), for current retained obs_idx list
    obs_idx: list[int] = []
    chol_L = np.zeros((0, 0), dtype=np.float64)

    def rebuild_cholesky():
        nonlocal chol_L
        if len(obs_idx) == 0:
            chol_L = np.zeros((0, 0), dtype=np.float64)
            return
        idx = np.array(obs_idx, dtype=np.int64)
        d = (idx[:, None] - idx[None, :]).astype(np.float64) * dt
        K = sigma2 * np.exp(-0.5 * (d * d) * inv_ell2)
        K[np.diag_indices_from(K)] += jitter
        chol_L = np.linalg.cholesky(K)

    def add_observation(i_new: int):
        """Rank-1 append update; if we truncate (drop oldest), rebuild."""
        nonlocal chol_L

        # Append then (optional) truncate to max_obs
        obs_idx.append(i_new)
        if max_obs is not None and len(obs_idx) > max_obs:
            # Drop oldest; rebuilding is simplest/stable (max_obs should be modest).
            obs_idx.pop(0)
            rebuild_cholesky()
            return

        # Incremental update when no drop
        if chol_L.shape[0] == 0:
            chol_L = np.array([[np.sqrt(k_tt(i_new) + jitter)]], dtype=np.float64)
            return

        idx = np.array(obs_idx[:-1], dtype=np.int64)  # previous obs
        kv = k_vec(i_new, idx)  # (M,)
        w = np.linalg.solve(chol_L, kv)  # (M,)
        diag_sq = (k_tt(i_new) + jitter) - float(w @ w)
        diag = np.sqrt(max(diag_sq, jitter))

        M = chol_L.shape[0]
        L_new = np.zeros((M + 1, M + 1), dtype=np.float64)
        L_new[:M, :M] = chol_L
        L_new[M, :M] = w
        L_new[M, M] = diag
        chol_L = L_new

    def posterior_var(t_idx: int) -> float:
        """Var(Z_t | current obs_idx)."""
        if len(obs_idx) == 0:
            return k_tt(t_idx)
        idx = np.array(obs_idx, dtype=np.int64)
        ktO = k_vec(t_idx, idx)
        alpha = np.linalg.solve(chol_L, ktO)
        var = k_tt(t_idx) - float(alpha @ alpha)
        return max(var, 0.0)

    # Main horizon sweep: at horizon h, before adding obs at h we have O_{h-1}
    info_gain = np.zeros(T, dtype=np.float64)

    for h in range(N):
        # model time corresponding to this horizon
        if h >= lag:
            t = h - lag

            # Var before incorporating potential obs at h
            var_prev = posterior_var(t)

            # Incorporate obs at h if present
            if observed[h]:
                add_observation(h)

            # Var after (if censored, obs set unchanged so var_post=var_prev)
            var_post = posterior_var(t)

            info_gain[t] = 0.5 * np.log((var_prev + eps) / (var_post + eps))

        else:
            # horizons before we can even define t>=0: still need to update obs set if any,
            # but in this dataset observed indices start at lag anyway.
            if observed[h]:
                add_observation(h)

    return info_gain



def make_censored_task(
    T: int = 1000,
    lag: int = 10,
    source: str = "get_ou_sample",
    censor_period: int = 40,
    seed: int = 0,
    censor_val: float = -3.0,
    **signal_kwargs,
):
    np.random.seed(seed)

    if source == "whitesignal":
        dt = signal_kwargs.pop("dt", 0.01)
        freq = signal_kwargs.pop("freq", 1.0)
        rms = signal_kwargs.pop("rms", 0.5)
        output_signal = whitesignal(period=(T + lag) * dt, dt=dt, freq=freq, rms=rms, **signal_kwargs)
    elif source == "ou":
        output_signal = get_ou_sample(T=(T+lag), **signal_kwargs)
    elif source  == "rbf":
        output_signal = get_rbf_sample(T=(T+lag), **signal_kwargs)
    else:
        raise ValueError(f"Unknown source '{source}'. Use 'whitesignal' or 'ou'.")

    input_signal = np.copy(output_signal)
    mask = (np.arange(T+lag) % censor_period) >= (censor_period // 2)

    input_signal[mask] = censor_val
 
    return input_signal[lag:], output_signal[:-lag]


def make_copying_task(
    T: int = 1000,
    lag: int = 10,
    source: str = "get_ou_sample",
    seed: int = 0,
    **signal_kwargs,
):
    """
    Build a simple sequence-copying task from a continuous signal.

    The input is generated either by :func:`whitesignal` or
    :func:`get_ou_sample`. The output is the same signal shifted forward by
    ``lag`` time steps, forcing a sequence model to retain information over that
    window to predict correctly.

    Parameters
    ----------
    T : int, optional
        Length of the sequence.
    lag : int, optional
        Number of time steps to shift the target output relative to the input.
    source : {"whitesignal", "ou"}, optional
        Which generator to use. "ou" selects :func:`get_ou_sample`.
    seed : int, optional
        Random seed used for reproducibility.
    **signal_kwargs :
        Additional keyword arguments forwarded to the signal generator.

    Returns
    -------
    input_signal : ndarray, shape (T,)
        The driving input sequence.
    target_signal : ndarray, shape (T,)
        The delayed copy of ``input_signal``.
    """

    if lag <= 0:
        raise ValueError("lag must be positive to form a copying task")
    if lag >= T:
        raise ValueError("lag must be smaller than T to produce a valid shift")

    np.random.seed(seed)

    if source == "whitesignal":
        dt = signal_kwargs.pop("dt", 0.01)
        freq = signal_kwargs.pop("freq", 1.0)
        rms = signal_kwargs.pop("rms", 0.5)
        input_signal = whitesignal(period=T * dt, dt=dt, freq=freq, rms=rms, **signal_kwargs)
    elif source in {"ou", "get_ou_sample"}:
        input_signal = get_ou_sample(T=T, **signal_kwargs)
    else:
        raise ValueError(f"Unknown source '{source}'. Use 'whitesignal' or 'ou'.")

    target_signal = np.zeros_like(input_signal)
    target_signal[lag:] = input_signal[: T - lag]

    return input_signal, target_signal


def make_multiplexing_task(
    T: int = 1000, K: int = 3, lag: int = 3, repeat: int = 1, seed: int = 42
):
    np.random.seed(seed)
    assert T % repeat == 0
    multiplex_pattern = np.random.randint(0, K, size=T // repeat)
    multiplex_pattern = np.repeat(multiplex_pattern, repeat)

    orig_seqs = np.stack([get_ou_sample(T=T) for _ in range(K)], axis=1)
    orig_seqs = 1.0 / (1.0 + np.exp(-orig_seqs))

    all_seqs = 0.8 * orig_seqs + 0.1
    all_seqs = all_seqs + np.arange(K)[None]
    all_seqs -= np.mean(all_seqs)
    all_seqs /= np.std(all_seqs)

    X = np.zeros(T)
    Y = np.zeros((T, K))
    histories = np.zeros((K, lag))
    task_idx = np.zeros(K, dtype=int)
    for i, k in enumerate(multiplex_pattern):
        X[i] = all_seqs[task_idx[k], k]
        histories[k] = np.roll(histories[k], -1)
        histories[k, -1] = orig_seqs[task_idx[k], k]
        task_idx[k] += 1
        Y[i] = np.mean(histories, axis=1)
    return X, Y


def make_induction_head_multioutput_s2s_task(
    T: int = 1000, V: int = 30, K: int = 5, seed: int = 42
):
    np.random.seed(seed)

    X = np.zeros(T, dtype=int)
    Y = np.zeros((T, K), dtype=int)

    X[0] = np.random.randint(0, V)
    for j in range(1, T):
        if X[j - 1] < K:
            X[j] = np.random.randint(K, V)
        else:
            X[j] = np.random.randint(0, V)

    current_vals = np.zeros(K, dtype=int)
    flag = False
    for t in range(T):
        if t > 0:
            Y[t] = Y[t - 1]
        if X[t] < K:
            flag = True
        elif flag:
            flag = False
            Y[t, X[t - 1]] = X[t]

    vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V - K)
    special_vocab = vocab[0] + np.linspace(-3, -1, K)
    vocab = np.concatenate([special_vocab, vocab], axis=0)
    return vocab[X], vocab[Y]


def make_implicit_measure_task(T: int = 100, filter_size: int = 15, seed: int = 42):
    np.random.seed(seed)
    t = np.linspace(0, 1, T)
    clean_signal = np.sin(8 * np.pi * t)
    mask = np.zeros(T)
    mask[:-filter_size] = 1.0
    observed = clean_signal.copy()
    noise = 0.05 * np.random.randn(T)
    observed += noise
    observed[mask == 0] = -4.0
    return observed, clean_signal

def make_simple_repetition_task(T: int = 100, shift: int = 10, seed: int = 42):

    np.random.seed(seed)

    t = np.linspace(0, 1, T)
    observed = np.sin(8 * np.pi * t)
    noise = 0.2 * np.random.randn(T)
    observed += noise

    output = observed.copy()
    output = output[:-shift]
    output = np.concatenate([np.zeros(shift), output])

    return observed, output


def whitesignal(period, dt, freq, rms=0.5, batch_shape=(), seed=None):
    """
    Copied from github.com/state-spaces/s4

    Produces output signal of length period / dt, band-limited to frequency freq
    Output shape (*batch_shape, period/dt)
    Adapted from the nengo library
    """
    assert not (freq is not None and freq < 1.0 / period)
    assert freq <= 0.5 / dt

    if seed is not None:
        np.random.seed(seed)

    n_coefficients = int(np.ceil(period / dt / 2.0))
    shape = batch_shape + (n_coefficients + 1,)
    sigma = rms * np.sqrt(0.5)
    coefficients = 1j * np.random.normal(0.0, sigma, size=shape)
    coefficients[..., -1] = 0.0
    coefficients += np.random.normal(0.0, sigma, size=shape)
    coefficients[..., 0] = 0.0
    set_to_zero = np.fft.rfftfreq(2 * n_coefficients, d=dt) > freq
    coefficients *= 1 - set_to_zero
    power_correction = np.sqrt(1.0 - np.sum(set_to_zero, dtype=float) / n_coefficients)
    if power_correction > 0:
        coefficients /= power_correction
    coefficients *= np.sqrt(2 * n_coefficients)
    signal = np.fft.irfft(coefficients, axis=-1)
    return signal


def wray_and_green_output(input_signal, a=2.0, m=0.3, k=0.08, tau_max=50, scaling=4e-3):
    """Implement the system described in Wray and Green (1994)."""
    T = len(input_signal)

    # Make the filter.
    mu = lambda t: a / m * np.exp(-k * t) * np.sin(m * t)
    tau_vals = np.arange(tau_max)
    filter = mu(tau_vals)[::-1]
    filter = scaling * np.outer(filter, filter)

    # Pad the input signal.
    output_signal = np.zeros_like(input_signal)
    input_signal = np.concatenate([np.zeros(tau_max - 1), input_signal], axis=0)

    # Do a convolution.
    for i in range(T):
        input_slice = input_signal[i : i + tau_max]
        output_signal[i] = np.sum(filter * np.outer(input_slice, input_slice))

    return output_signal


__all__ = [
    "get_ou_sample",
    "make_copying_task",
    "make_induction_head_multioutput_s2s_task",
    "make_induction_head_seq_to_seq_task",
    "make_induction_head_task",
    "make_implicit_measure_task",
    "make_multiplexing_task",
    "whitesignal",
    "wray_and_green_output",
]