"""aurora-blend-v2 — score-tuned long-context mixture with weather slice.

Cascade score = geomean(CRPS, MASE) on held-out windows (lower is better).
This prior is tuned so Toto2 learns *forecastable* multi-scale structure:

  * weather mass ~0.42 (king-92 skill signal: diurnal / radiation / cloud)
  * remaining mass on long-range AR / integrated / GP / seasonal (king-zenfro)
  * calendar-biased periods (24 / 168 / 7 / 12 / 720), higher SNR, quieter innov
  * weekly-structured intermittent demand for retail-like domains

Hot recurrences stay numba-jitted (cache=False for cascade's dynamic loader).
Chunked emission + fixed L=4096 keep generation wall-safe under the token wall.
"""

from __future__ import annotations

import json
from collections.abc import Callable, Iterator
from pathlib import Path
from typing import Any

import numpy as np
from numba import njit

from cascade.interface import DataGenerator

_BATCH = 512
# Block size for the segmented AR(1) scan (~2*sqrt(L) iters instead of L).
_AR1_BLOCK = 32

_PRIOR_NAMES: tuple[str, ...] = (
    "trend_seasonal_ar",
    "regime_shift",
    "multiplicative",
    "ar2",
    "integrated",
    "threshold_ar",
    "chaotic",
    "rff_gp",
    "intermittent",
    "pulse_outlier",
    "weather",
)

# Defaults mirror config.json; config overrides win.
_PRIOR_MASS: dict[str, float] = {
    "weather": 0.42,
    "trend_seasonal_ar": 0.15,
    "regime_shift": 0.08,
    "ar2": 0.09,
    "integrated": 0.09,
    "rff_gp": 0.07,
    "multiplicative": 0.05,
    "threshold_ar": 0.02,
    "chaotic": 0.01,
    "intermittent": 0.015,
    "pulse_outlier": 0.005,
}

# Period menu + sampling weights: bias toward cadences that dominate real eval
# (hourly diurnal, weekly, daily, monthly) while retaining long-context options.
_SEASON_PERIODS = np.asarray(
    [4.0, 7.0, 12.0, 24.0, 30.0, 52.0, 96.0, 144.0, 168.0, 336.0, 504.0, 672.0, 720.0],
    dtype=np.float64,
)
_SEASON_P = np.asarray(
    [0.02, 0.12, 0.10, 0.22, 0.04, 0.04, 0.03, 0.03, 0.18, 0.05, 0.04, 0.05, 0.08],
    dtype=np.float64,
)
_SEASON_P = _SEASON_P / _SEASON_P.sum()


# ── numba kernels (cache=False: cascade imports this file as a dynamic module) ─


@njit(cache=False)
def _jit_ar2(innov: np.ndarray, a1: np.ndarray, a2: np.ndarray) -> np.ndarray:
    n, L = innov.shape
    out = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        aa1, aa2 = a1[i], a2[i]
        out[i, 0] = innov[i, 0]
        if L > 1:
            out[i, 1] = aa1 * out[i, 0] + innov[i, 1]
        for t in range(2, L):
            out[i, t] = aa1 * out[i, t - 1] + aa2 * out[i, t - 2] + innov[i, t]
    return out


@njit(cache=False)
def _jit_setar(
    innov: np.ndarray,
    phi_pos: np.ndarray,
    phi_neg: np.ndarray,
    c_pos: np.ndarray,
    c_neg: np.ndarray,
) -> np.ndarray:
    n, L = innov.shape
    out = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        out[i, 0] = innov[i, 0]
        for t in range(1, L):
            prev = out[i, t - 1]
            if prev >= 0.0:
                v = c_pos[i] + phi_pos[i] * prev + innov[i, t]
            else:
                v = c_neg[i] + phi_neg[i] * prev + innov[i, t]
            if v > 1e6:
                v = 1e6
            elif v < -1e6:
                v = -1e6
            out[i, t] = v
    return out


@njit(cache=False)
def _jit_chaos(
    pick_sine: np.ndarray,
    r_log: np.ndarray,
    r_sin: np.ndarray,
    x0: np.ndarray,
    L: int,
) -> np.ndarray:
    n = x0.shape[0]
    out = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        cur = x0[i]
        out[i, 0] = cur
        sine = pick_sine[i]
        rl, rs = r_log[i], r_sin[i]
        for t in range(1, L):
            if sine:
                cur = rs * np.sin(np.pi * cur)
            else:
                cur = rl * cur * (1.0 - cur)
            if cur < 0.0:
                cur = 0.0
            elif cur > 1.0:
                cur = 1.0
            out[i, t] = cur
    return out


@njit(cache=False)
def _jit_holds(series: np.ndarray, hold: np.ndarray) -> None:
    n, L = series.shape
    for i in range(n):
        for t in range(1, L):
            if hold[i, t]:
                series[i, t] = series[i, t - 1]


# ── thin numpy wrappers ─────────────────────────────────────────────────────


def _ar1(innov: np.ndarray, phi: np.ndarray, S: int = _AR1_BLOCK) -> np.ndarray:
    """AR(1) via segmented (block) scan — ~2*sqrt(L) Python iters, same law as loop.

    Weather/cloud call this many times per batch; a pure-numpy scan beats a
    per-row numba loop here and feeds the trainer more tokens before the wall.
    """
    n, L = innov.shape
    p = np.asarray(phi, dtype=np.float64).reshape(n)
    if L < 2 * S:
        x = np.empty((n, L), dtype=np.float64)
        x[:, 0] = innov[:, 0]
        for t in range(1, L):
            x[:, t] = p * x[:, t - 1] + innov[:, t]
        return x
    B = L // S
    body = B * S
    main = innov[:, :body].reshape(n, B, S)
    y = np.empty((n, B, S), dtype=np.float64)
    y[:, :, 0] = main[:, :, 0]
    pcol = p[:, None]
    for s in range(1, S):
        y[:, :, s] = pcol * y[:, :, s - 1] + main[:, :, s]
    r = p ** S
    ylast = y[:, :, S - 1]
    X = np.empty((n, B), dtype=np.float64)
    X[:, 0] = ylast[:, 0]
    for b in range(1, B):
        X[:, b] = r * X[:, b - 1] + ylast[:, b]
    carry_in = np.empty((n, B), dtype=np.float64)
    carry_in[:, 0] = 0.0
    carry_in[:, 1:] = X[:, :-1]
    ppow = p[:, None] ** np.arange(1, S + 1, dtype=np.float64)[None, :]
    x = y + carry_in[:, :, None] * ppow[:, None, :]
    x = x.reshape(n, body)
    if body == L:
        return x
    out = np.empty((n, L), dtype=np.float64)
    out[:, :body] = x
    prev = out[:, body - 1]
    for t in range(body, L):
        prev = p * prev + innov[:, t]
        out[:, t] = prev
    return out


def _ar2(innov: np.ndarray, a1: np.ndarray, a2: np.ndarray) -> np.ndarray:
    return _jit_ar2(
        np.ascontiguousarray(innov),
        np.ascontiguousarray(a1.reshape(-1)),
        np.ascontiguousarray(a2.reshape(-1)),
    )


def _pick_periods(rng: np.random.Generator, n: int) -> np.ndarray:
    return rng.choice(_SEASON_PERIODS, size=n, p=_SEASON_P)[:, None]


def _harmonics(rng: np.random.Generator, n: int, L: int, *, depth: int = 3) -> np.ndarray:
    """Sum of up to ``depth`` sinusoids; periods biased to calendar cadences."""
    t = np.arange(L, dtype=np.float64)[None, :]
    n_comp = rng.integers(1, depth + 1, size=n)
    acc = np.zeros((n, L), dtype=np.float64)
    for j in range(depth):
        on = (n_comp > j).astype(np.float64)[:, None]
        period = _pick_periods(rng, n)
        # Slightly stronger amps on the first component (dominant seasonal).
        lo, hi = (0.4, 2.2) if j == 0 else (0.15, 1.2)
        amp = rng.uniform(lo, hi, size=n)[:, None]
        phase = rng.uniform(0.0, 2.0 * np.pi, size=n)[:, None]
        acc += on * amp * np.sin(2.0 * np.pi * t / period + phase)
    return acc


def _jumps(rng: np.random.Generator, n: int, L: int, rate: float, scale) -> np.ndarray:
    hit = rng.random((n, L)) < rate
    mag = rng.normal(0.0, 1.0, size=(n, L))
    s = np.asarray(scale, dtype=np.float64)
    if s.ndim == 1:
        s = s[:, None]
    out = hit * mag * s
    out[:, 0] = 0.0
    return out


def _finite(block: np.ndarray) -> np.ndarray:
    x = np.asarray(block, dtype=np.float64)
    x = np.nan_to_num(x, nan=0.0, posinf=1e6, neginf=-1e6)
    return np.clip(x, -1e6, 1e6)


# ── family emitters ─────────────────────────────────────────────────────────


def emit_trend_seasonal(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    hi_frac: float,
    exc_lo: float,
    exc_hi: float,
) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    level = rng.normal(0.0, 1.0, size=(n, 1))
    heavy = rng.random((n, 1)) < hi_frac
    slope = np.where(
        heavy,
        rng.normal(0.0, exc_hi, size=(n, 1)),
        rng.normal(0.0, exc_lo, size=(n, 1)),
    )
    tn = t / max(L - 1, 1)
    signal = level + slope * tn + _harmonics(rng, n, L, depth=3)
    # More persistent AR noise, lower sigma → higher SNR / easier forecasts
    phi = rng.uniform(0.35, 0.92, size=n)
    sigma = rng.uniform(0.06, 0.40, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    return signal + _ar1(innov, phi)


def emit_regime(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    # Slightly fewer breaks than classic (2/L): long stable regimes are learnable.
    level = np.cumsum(_jumps(rng, n, L, rate=2.0 / L, scale=2.0), axis=1)
    log_vol = np.cumsum(_jumps(rng, n, L, rate=2.0 / L, scale=0.4), axis=1)
    vol = np.exp(np.clip(log_vol, -3.0, 3.0)) * rng.uniform(0.08, 0.40, size=(n, 1))
    noise = rng.normal(0.0, 1.0, size=(n, L)) * vol
    seas = _harmonics(rng, n, L, depth=2) * rng.uniform(0.2, 1.0, size=(n, 1))
    return level + seas + noise


def emit_multiplicative(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    hi_frac: float,
    exc_lo: float,
    exc_hi: float,
) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    heavy = rng.random((n, 1)) < hi_frac
    g = np.where(
        heavy,
        rng.normal(0.0, exc_hi, size=(n, 1)),
        rng.normal(0.0, exc_lo, size=(n, 1)),
    )
    tn = t / max(L - 1, 1)
    base = np.exp(g * tn + rng.normal(0.0, 0.25, size=(n, 1)))
    amp = rng.uniform(0.15, 0.65, size=(n, 1))
    per = rng.choice(np.array([7.0, 12.0, 24.0, 52.0, 168.0, 720.0]), size=n)[:, None]
    seas = 1.0 + amp * np.sin(2.0 * np.pi * t / per + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1)))
    # Quieter multiplicative noise
    noise = 1.0 + rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.015, 0.10, size=(n, 1))
    scale = rng.uniform(1.0, 50.0, size=(n, 1))
    return scale * base * np.clip(seas, 0.05, None) * np.clip(noise, 0.05, None)


def emit_ar2(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    # Bias toward near-unit-root persistence (finance / sensor domains).
    p1 = rng.uniform(0.45, 0.985, size=n)
    p2 = rng.uniform(-0.45, 0.45, size=n)
    a1 = p1 * (1.0 - p2)
    a2 = p2
    sigma = rng.uniform(0.15, 0.55, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    path = _ar2(innov, a1, a2)
    drift = rng.normal(0.0, 0.004, size=(n, 1)) * np.arange(L, dtype=np.float64)[None, :]
    return path + drift


def emit_integrated(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    twice = rng.random(n) < 0.30
    drift = rng.normal(0.0, 0.015, size=(n, 1))
    sigma = rng.uniform(0.15, 0.75, size=(n, 1))
    steps = rng.normal(0.0, 1.0, size=(n, L)) * sigma + drift
    walk = np.cumsum(steps, axis=1)
    walk2 = np.cumsum(walk, axis=1)
    return np.where(twice[:, None], walk2 / max(L, 1) ** 0.5, walk)


def emit_threshold(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    phi_pos = rng.uniform(0.4, 0.92, size=n)
    phi_neg = rng.uniform(-0.7, 0.35, size=n)
    c_pos = rng.normal(0.0, 0.25, size=n)
    c_neg = rng.normal(0.0, 0.25, size=n)
    sigma = rng.uniform(0.15, 0.55, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    return _jit_setar(
        np.ascontiguousarray(innov),
        np.ascontiguousarray(phi_pos),
        np.ascontiguousarray(phi_neg),
        np.ascontiguousarray(c_pos),
        np.ascontiguousarray(c_neg),
    )


def emit_chaotic(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    # Milder maps (less fully chaotic) so a tiny residual mass is not pure noise.
    pick_sine = rng.random(n) < 0.55
    r_log = rng.uniform(3.4, 3.9, size=n)
    r_sin = rng.uniform(0.80, 0.97, size=n)
    x0 = rng.uniform(0.05, 0.95, size=n)
    return _jit_chaos(
        np.ascontiguousarray(pick_sine),
        np.ascontiguousarray(r_log),
        np.ascontiguousarray(r_sin),
        np.ascontiguousarray(x0),
        L,
    )


def emit_rff(rng: np.random.Generator, n: int, L: int, *, features: int = 24) -> np.ndarray:
    """Stationary GP via RFF; longer lengthscales → smoother, more forecastable.

    K=24 matches the v16fast quality-neutral speed cut — same smoothness class,
    ~1.5x fewer feature loops so GP mass does not starve the wall.
    """
    t = np.arange(L, dtype=np.float64)[None, :]
    lengthscale = rng.uniform(40.0, 320.0, size=(n, 1))
    acc = np.zeros((n, L), dtype=np.float64)
    scale = np.sqrt(2.0 / features)
    for _ in range(features):
        omega = rng.normal(0.0, 1.0, size=(n, 1)) / lengthscale
        phase = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
        np.add(acc, np.cos(omega * t + phase), out=acc)
    return scale * acc


def emit_intermittent(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """Zero-inflated demand with a mild weekly occurrence modulation (retail-like)."""
    t = np.arange(L, dtype=np.float64)[None, :]
    p0 = rng.uniform(0.06, 0.35, size=(n, 1))
    week = 0.55 + 0.45 * (
        0.5 + 0.5 * np.sin(2.0 * np.pi * t / 7.0 + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1)))
    )
    p = np.clip(p0 * week, 0.02, 0.55)
    occur = (rng.random((n, L)) < p).astype(np.float64)
    mag = rng.gamma(shape=2.0, scale=1.0, size=(n, L)) * rng.uniform(1.0, 8.0, size=(n, 1))
    floor = rng.uniform(0.0, 0.4, size=(n, 1))
    return floor + occur * mag


def emit_pulse(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    base = emit_rff(rng, n, L, features=24) * rng.uniform(0.5, 2.0, size=(n, 1))
    base += _harmonics(rng, n, L, depth=1) * rng.uniform(0.2, 1.0, size=(n, 1))
    spikes = _jumps(rng, n, L, rate=4.0 / L, scale=rng.uniform(2.5, 6.5, size=n))
    series = base + spikes
    hold = rng.random((n, L)) < (2.0 / L)
    hold[:, 0] = False
    _jit_holds(series, hold)
    return series


# ── weather: sliced archetype mixture (base / radiation / cloud) ────────────


def _cloud_cover(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """U-shaped [0, 100] cloud fraction via multi-timescale red noise + clip."""
    t = np.arange(L, dtype=np.float64)[None, :]

    def _z_ar(phi_lo: float, phi_hi: float) -> np.ndarray:
        phi = rng.uniform(phi_lo, phi_hi, size=n)
        z = _ar1(rng.normal(0.0, 1.0, size=(n, L)), phi)
        return (z - z.mean(axis=1, keepdims=True)) / (z.std(axis=1, keepdims=True) + 1e-9)

    syn = _z_ar(0.990, 0.9990)
    mid = _z_ar(0.895, 0.95)
    fast = _z_ar(0.75, 0.87)

    g = rng.normal(0.0, 1.0, size=(n, L))
    spike = (rng.random((n, L)) < 0.05) * rng.uniform(5.0, 11.0, size=(n, L))
    heavy = g * (1.0 + spike)
    rough = heavy - 0.30 * np.concatenate([np.zeros((n, 1)), heavy[:, :-1]], axis=1)
    rough = rough / (rough.std(axis=1, keepdims=True) + 1e-9)

    diur = np.sin(2.0 * np.pi * t / 24.0 + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1)))
    u = (
        0.95 * rng.uniform(0.8, 1.2, size=(n, 1)) * syn
        + 0.82 * rng.uniform(0.7, 1.3, size=(n, 1)) * mid
        + 0.55 * rng.uniform(0.7, 1.3, size=(n, 1)) * fast
        + 0.30 * rng.uniform(0.7, 1.3, size=(n, 1)) * rough
        + 0.44 * rng.uniform(0.3, 1.4, size=(n, 1)) * diur
    )
    centre = np.clip(rng.normal(61.0, 29.0, size=(n, 1)), 2.0, 98.0)
    span = 57.0 * rng.uniform(0.82, 1.18, size=(n, 1))
    return np.round(np.clip(centre + span * u, 0.0, 100.0))


def emit_weather(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """Mixture of diurnal base, night-floored radiation, and cloud cover.

    Archetypes assigned first; expensive multi-scale red-noise only on cloud rows.
    v2: quieter AR innov + optional monthly harmonic for longer contexts.
    """
    t = np.arange(L, dtype=np.float64)[None, :]
    is_rad = rng.random(n) < 0.25
    is_cloud = (rng.random(n) < 0.24) & ~is_rad
    is_wbp = ~is_cloud
    out = np.empty((n, L), dtype=np.float64)

    wbp = np.nonzero(is_wbp)[0]
    m = wbp.size
    if m:
        rad_rows = is_rad[wbp][:, None]
        level = rng.normal(0.0, 1.0, size=(m, 1))
        a1 = rng.uniform(0.55, 2.1, size=(m, 1))
        p1 = rng.uniform(0.0, 2.0 * np.pi, size=(m, 1))
        a2 = rng.uniform(0.12, 0.65, size=(m, 1))
        p2 = rng.uniform(0.0, 2.0 * np.pi, size=(m, 1))
        seas = a1 * np.sin(2.0 * np.pi * t / 24.0 + p1) + a2 * np.sin(2.0 * np.pi * t / 12.0 + p2)
        weekly_on = (rng.random((m, 1)) < 0.45).astype(np.float64)
        seas += weekly_on * rng.uniform(0.12, 0.55, size=(m, 1)) * np.sin(
            2.0 * np.pi * t / 168.0 + rng.uniform(0.0, 2.0 * np.pi, size=(m, 1))
        )
        # Monthly harmonic (hourly ~720) — multi-week structure inside L=4096
        month_on = (rng.random((m, 1)) < 0.35).astype(np.float64)
        seas += month_on * rng.uniform(0.08, 0.40, size=(m, 1)) * np.sin(
            2.0 * np.pi * t / 720.0 + rng.uniform(0.0, 2.0 * np.pi, size=(m, 1))
        )
        y_amp = rng.uniform(0.25, 1.6, size=(m, 1))
        y_per = rng.uniform(2000.0, 9000.0, size=(m, 1))
        yearly = y_amp * np.sin(2.0 * np.pi * t / y_per + rng.uniform(0.0, 2.0 * np.pi, size=(m, 1)))
        phi = rng.uniform(0.65, 0.96, size=m)
        sigma = rng.uniform(0.04, 0.18, size=(m, 1))
        noise = _ar1(rng.normal(0.0, 1.0, size=(m, L)) * sigma, phi)
        base = level + seas + yearly + noise

        thr = rng.uniform(0.2, 0.6, size=(m, 1)) * a1
        ramp = rng.uniform(0.85, 1.65, size=(m, 1))
        diurnal = ramp * np.maximum(a1 * np.sin(2.0 * np.pi * t / 24.0 + p1) - thr, 0.0)
        night = diurnal <= 0.0
        rad = level + diurnal + yearly + np.where(night, noise * 0.12, noise)
        out[wbp] = np.where(rad_rows, rad, base)

    cloud = np.nonzero(is_cloud)[0]
    if cloud.size:
        out[cloud] = _cloud_cover(rng, int(cloud.size), L)
    return out


# ── generator entry point ───────────────────────────────────────────────────


class Generator(DataGenerator):
    """Mixture-of-priors synthesizer. Submit as ``generator.Generator``."""

    def __init__(self, config_dir: str, *, seed: int) -> None:
        cfg_path = Path(config_dir) / "config.json"
        cfg: dict[str, Any] = (
            json.loads(cfg_path.read_text(encoding="utf-8")) if cfg_path.is_file() else {}
        )
        self._cfg = cfg
        self._seed = int(seed)
        self._min_len = int(cfg.get("min_length", 64))
        self._max_len = int(cfg.get("max_length", 4096))
        if self._min_len < 1 or self._max_len < self._min_len:
            raise ValueError(f"invalid length band [{self._min_len}, {self._max_len}]")

        mass = dict(_PRIOR_MASS)
        for key, val in dict(cfg.get("family_weights", {})).items():
            if key in mass:
                mass[key] = float(val)
        weights = np.asarray([mass[name] for name in _PRIOR_NAMES], dtype=np.float64)
        if not np.all(np.isfinite(weights)) or weights.min() < 0 or weights.sum() <= 0:
            raise ValueError("family_weights must be finite, non-negative, and not all zero")
        self._weights = weights / weights.sum()

        self._tr_hi = float(cfg.get("tr_hi_frac", 0.20))
        self._tr_lo = float(cfg.get("tr_exc_lo", 0.35))
        self._tr_hi_s = float(cfg.get("tr_exc_hi", 2.0))
        self._gr_lo = float(cfg.get("gr_exc_lo", 0.25))
        self._gr_hi = float(cfg.get("gr_exc_hi", 1.4))
        self._fixed = self._min_len == self._max_len

    @property
    def name(self) -> str:
        return str(self._cfg.get("name", "aurora-blend-v2"))

    def _builders(self) -> tuple[Callable[..., np.ndarray], ...]:
        return (
            lambda rng, n, L: emit_trend_seasonal(
                rng, n, L, hi_frac=self._tr_hi, exc_lo=self._tr_lo, exc_hi=self._tr_hi_s
            ),
            emit_regime,
            lambda rng, n, L: emit_multiplicative(
                rng, n, L, hi_frac=self._tr_hi, exc_lo=self._gr_lo, exc_hi=self._gr_hi
            ),
            emit_ar2,
            emit_integrated,
            emit_threshold,
            emit_chaotic,
            emit_rff,
            emit_intermittent,
            emit_pulse,
            emit_weather,
        )

    def generate(self, n_series: int) -> Iterator[np.ndarray]:
        if n_series <= 0:
            return
        rng = np.random.default_rng(self._seed)
        builders = self._builders()
        L_max = self._max_len
        done = 0
        while done < n_series:
            # Always draw a full batch so series i is a pure function of (seed, i).
            lengths = rng.integers(self._min_len, L_max + 1, size=_BATCH)
            picks = rng.choice(len(_PRIOR_NAMES), size=_BATCH, p=self._weights)
            slots: list[np.ndarray | None] = [None] * _BATCH
            for fam_id, build in enumerate(builders):
                idx = np.nonzero(picks == fam_id)[0]
                if idx.size == 0:
                    continue
                block = _finite(build(rng, int(idx.size), L_max))
                if self._fixed:
                    for row, slot in enumerate(idx):
                        slots[int(slot)] = np.ascontiguousarray(block[row], dtype=np.float64)
                else:
                    for row, slot in enumerate(idx):
                        L = int(lengths[slot])
                        slots[int(slot)] = np.ascontiguousarray(block[row, :L], dtype=np.float64)
            take = min(_BATCH, n_series - done)
            for arr in slots[:take]:
                if arr is None:  # pragma: no cover
                    raise RuntimeError("unfilled series slot")
                yield arr
            done += take