"""zenfro_v3 — Grammar-of-Time (GoT) mixture-of-priors generator.

Built on the   / fullctx-spectral-v12 infrastructure (prefetch chunk
producer, SciPy AR filters, FFT spectral stems, measurement artifacts), with a
new competitive prior: **Grammar of Time**.

Why a grammar, not another single-family mix
--------------------------------------------
Cascade scores a Toto2 backbone trained from scratch on your corpus. The
winning synthetic priors in the literature (Chronos-2, TempoPFN, ForecastPFN,
CauKer) are not single ARIMA draws — they are *compositions* of temporal
primitives (trend × season × noise × events × causal lags).   already
covers many primitives as mutually exclusive families. GoT goes one step
further: a substantial share of series is produced by sampling a short
**production** over stems and operators, so the model sees the same structural
atoms *combined* the way real series combine them.

Temporal grammar (informal CFG)::

    Series   → Production | Stem
    Production → Compose | Splice | Nested | Causal
    Compose  → Stem ⊕ Stem [⊕ Stem]          # additive / multiplicative phrase
    Splice   → Stem ‖ Stem [‖ Stem]          # clause boundary at breakpoints
    Nested   → Envelope ⋉ Carrier            # slow scale modulates fast scale
    Causal   → Driver ▷ Response             # lagged temporal DAG (1→1)
    Stem     → {trend_seasonal, ar2, integrated, spectral, ou, long_memory, …}
    Affix    → jump | hold | pulse | artifact   # applied after production

Determinism, code-only, bounded+finite contracts are unchanged: one
``np.random.default_rng(seed)``, allowlisted NumPy/SciPy, ``_sanitize`` gate.
"""

from __future__ import annotations

import json
from collections.abc import Iterator
from functools import lru_cache, partial
from pathlib import Path
from queue import Full, Queue
from threading import Event, Thread

import numpy as np
from scipy.signal import lfilter

from cascade.interface import DataGenerator

# Prefetch chunk size —   measured 2048 as the local sweet spot.
_CHUNK = 2048

_SEASONAL_PERIODS = np.array(
    [4, 7, 12, 24, 30, 48, 52, 90, 96, 144, 168, 183, 288, 336, 365, 672, 730],
    dtype=np.float64,
)
_SEASONAL_PROBS = np.array(
    [0.04, 0.12, 0.04, 0.16, 0.05, 0.06, 0.04, 0.03, 0.07, 0.03,
     0.13, 0.04, 0.04, 0.06, 0.07, 0.04, 0.05],
    dtype=np.float64,
)
_SEASONAL_PROBS /= _SEASONAL_PROBS.sum()

# ── family mixture ──────────────────────────────────────────────────────────
#   dynamics-heavy core retained; four GoT productions added. Default
# mass puts ~22% on compositional grammar so the model learns combination
# structure without drowning the proven single-stem families.
_FAMILIES: tuple[str, ...] = (
    "trend_seasonal_ar",
    "regime_shift",
    "multiplicative",
    "ar2",
    "integrated",
    "threshold_ar",
    "chaotic",
    "spectral_gp",
    "long_memory",
    "ou_stochastic_vol",
    "physical_sensors",
    "seasonal_counts",
    "intermittent",
    "pulse_outlier",
    "got_compose",   # Grammar: Stem ⊕ Stem [⊕ Stem]
    "got_splice",    # Grammar: Stem ‖ Stem at clause boundaries
    "got_nested",    # Grammar: Envelope ⋉ Carrier (multi-scale)
    "got_causal",    # Grammar: Driver ▷ lagged Response
)
_DEFAULT_WEIGHTS: dict[str, float] = {
    "trend_seasonal_ar": 0.10,
    "regime_shift": 0.10,
    "multiplicative": 0.06,
    "ar2": 0.12,
    "integrated": 0.10,
    "threshold_ar": 0.06,
    "chaotic": 0.03,
    "spectral_gp": 0.05,
    "long_memory": 0.05,
    "ou_stochastic_vol": 0.08,
    "physical_sensors": 0.015,
    "seasonal_counts": 0.015,
    "intermittent": 0.01,
    "pulse_outlier": 0.01,
    "got_compose": 0.08,
    "got_splice": 0.06,
    "got_nested": 0.05,
    "got_causal": 0.03,
}


class Generator(DataGenerator):
    """Grammar-of-Time mixture generator. Submit as ``generator.Generator``."""

    def __init__(self, config_dir: str, *, seed: int) -> None:
        cfg_path = Path(config_dir) / "config.json"
        cfg = 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}]")
        weights = dict(_DEFAULT_WEIGHTS)
        for k, v in dict(cfg.get("family_weights", {})).items():
            if k in weights:
                weights[k] = float(v)
        w = np.asarray([weights[f] for f in _FAMILIES], dtype=np.float64)
        if not np.all(np.isfinite(w)) or w.min() < 0 or w.sum() <= 0:
            raise ValueError("family_weights must be finite, non-negative, and not all zero")
        self._weights = w / w.sum()
        self._tr_hi_frac = float(cfg.get("tr_hi_frac", 0.25))
        self._tr_exc_lo = float(cfg.get("tr_exc_lo", 0.4))
        self._tr_exc_hi = float(cfg.get("tr_exc_hi", 2.5))
        self._gr_exc_lo = float(cfg.get("gr_exc_lo", 0.3))
        self._gr_exc_hi = float(cfg.get("gr_exc_hi", 1.5))
        self._sa_clean_frac = float(cfg.get("sa_clean_frac", 0.4))
        self._sa_clean_lo = float(cfg.get("sa_clean_lo", 0.02))
        self._sa_clean_hi = float(cfg.get("sa_clean_hi", 0.12))
        # GoT knobs — composition depth, splice density, nest scale span.
        self._got_depth = int(cfg.get("got_depth", 3))  # stems in a compose (2..4)
        self._got_mul_frac = float(cfg.get("got_mul_frac", 0.35))
        self._got_splice_cuts = int(cfg.get("got_splice_cuts", 2))
        self._got_nest_ratio = float(cfg.get("got_nest_ratio", 6.0))
        self._got_causal_lag_frac = float(cfg.get("got_causal_lag_frac", 0.08))
        self._artifact_scale = float(cfg.get("artifact_scale", 1.0))

    @property
    def name(self) -> str:
        return str(self._cfg.get("name", "zenfro-v3-grammar-of-time"))

    def generate(self, n_series: int) -> Iterator[np.ndarray]:
        if n_series <= 0:
            return
        rng = np.random.default_rng(self._seed)
        max_len = self._max_len
        builders = (
            partial(_trend_seasonal_ar, hi_frac=self._tr_hi_frac,
                    exc_lo=self._tr_exc_lo, exc_hi=self._tr_exc_hi,
                    clean_frac=self._sa_clean_frac,
                    clean_lo=self._sa_clean_lo, clean_hi=self._sa_clean_hi),
            _regime_shift,
            partial(_multiplicative, hi_frac=self._tr_hi_frac,
                    exc_lo=self._gr_exc_lo, exc_hi=self._gr_exc_hi),
            _ar2, _integrated, _threshold_ar, _chaotic, _spectral_gp,
            _long_memory, _ou_stochastic_vol, _physical_sensors,
            _seasonal_counts, _intermittent, _pulse_outlier,
            partial(_got_compose, depth=self._got_depth, mul_frac=self._got_mul_frac,
                    hi_frac=self._tr_hi_frac, exc_lo=self._tr_exc_lo, exc_hi=self._tr_exc_hi),
            partial(_got_splice, n_cuts=self._got_splice_cuts,
                    hi_frac=self._tr_hi_frac, exc_lo=self._tr_exc_lo, exc_hi=self._tr_exc_hi),
            partial(_got_nested, nest_ratio=self._got_nest_ratio),
            partial(_got_causal, lag_frac=self._got_causal_lag_frac),
        )
        queue: Queue[object] = Queue(maxsize=1)
        stop = Event()
        done = object()
        art_scale = self._artifact_scale

        def put(item: object) -> bool:
            while not stop.is_set():
                try:
                    queue.put(item, timeout=0.1)
                    return True
                except Full:
                    continue
            return False

        def produce() -> None:
            try:
                produced = 0
                while produced < n_series and not stop.is_set():
                    lengths = rng.integers(self._min_len, max_len + 1, size=_CHUNK)
                    fam_ids = rng.choice(len(_FAMILIES), size=_CHUNK, p=self._weights)
                    chunk: list[np.ndarray | None] = [None] * _CHUNK
                    for fam in range(len(_FAMILIES)):
                        idx = np.nonzero(fam_ids == fam)[0]
                        if idx.size == 0:
                            continue
                        block = builders[fam](rng, int(idx.size), max_len)
                        preserve_nonnegative = fam in (2, 10, 11, 12)
                        block = _sanitize(
                            _measurement_artifacts(
                                rng,
                                block,
                                preserve_nonnegative=preserve_nonnegative,
                                rate_scale=art_scale,
                            )
                        )
                        for row, series_i in enumerate(idx):
                            length = int(lengths[series_i])
                            chunk[series_i] = np.ascontiguousarray(
                                block[row, :length], dtype=np.float64
                            )
                    take = min(_CHUNK, n_series - produced)
                    if not put((chunk, take)):
                        return
                    produced += take
            except BaseException as exc:
                put(exc)
            finally:
                put(done)

        producer = Thread(target=produce, name="zenfro-v3-generator", daemon=True)
        producer.start()
        try:
            while True:
                item = queue.get()
                if item is done:
                    break
                if isinstance(item, BaseException):
                    raise item
                chunk, take = item
                for arr in chunk[:take]:
                    if arr is None:  # pragma: no cover
                        raise RuntimeError("internal: unfilled series slot")
                    yield arr
        finally:
            stop.set()
            producer.join(timeout=1.0)


# ── shared vectorised primitives ────────────────────────────────────────────


def _ar1_batch(innov: np.ndarray, phi: np.ndarray) -> np.ndarray:
    n, L = innov.shape
    x = np.empty((n, L), dtype=np.float64)
    p = phi.reshape(n)
    for i in range(n):
        x[i] = lfilter([1.0], [1.0, -float(p[i])], innov[i])
    return x


def _ar2_batch(innov: np.ndarray, a1: np.ndarray, a2: np.ndarray) -> np.ndarray:
    n, L = innov.shape
    x = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        x[i] = lfilter(
            [1.0], [1.0, -float(a1[i]), -float(a2[i])], innov[i]
        )
    return x


@lru_cache(maxsize=4)
def _seasonal_basis(L: int) -> tuple[np.ndarray, np.ndarray]:
    angle = (
        2.0
        * np.pi
        * np.arange(L, dtype=np.float64)[None, :]
        / _SEASONAL_PERIODS[:, None]
    )
    return np.sin(angle), np.cos(angle)


def _seasonal(rng: np.random.Generator, n: int, L: int, k_max: int = 3) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    sin_basis, cos_basis = _seasonal_basis(L)
    k = rng.integers(1, k_max + 1, size=n)
    out = np.zeros((n, L), dtype=np.float64)
    for j in range(k_max):
        active = np.nonzero(k > j)[0]
        per = rng.choice(_SEASONAL_PERIODS, size=n, p=_SEASONAL_PROBS)[:, None]
        amp = rng.uniform(0.2, 2.0, size=n)[:, None]
        phase = rng.uniform(0.0, 2.0 * np.pi, size=n)[:, None]
        basis_idx = np.searchsorted(_SEASONAL_PERIODS, per[active, 0])
        component = amp[active] * (
            sin_basis[basis_idx] * np.cos(phase[active])
            + cos_basis[basis_idx] * np.sin(phase[active])
        )
        modulated = np.nonzero((k > j) & (rng.random(n) < 0.35))[0]
        if modulated.size:
            modulated_local = np.searchsorted(active, modulated)
            modulated_arg = (
                2.0 * np.pi * t / per[modulated] + phase[modulated]
            )
            m_per = np.clip(
                per[modulated] * rng.uniform(
                    4.0, 12.0, size=(modulated.size, 1)
                ),
                32.0,
                2.0 * L,
            )
            m_phase = rng.uniform(
                0.0, 2.0 * np.pi, size=(modulated.size, 1)
            )
            slow = np.sin(2.0 * np.pi * t / m_per + m_phase)
            amp_mod = 1.0 + rng.uniform(
                0.05, 0.45, size=(modulated.size, 1)
            ) * slow
            phase_mod = rng.uniform(
                0.05, 0.75, size=(modulated.size, 1)
            ) * np.sin(2.0 * np.pi * t / (1.7 * m_per) - m_phase)
            component[modulated_local] = (
                amp[modulated]
                * amp_mod
                * np.sin(modulated_arg + phase_mod)
            )
        out[active] += component
    return out


def _sparse_jumps(rng: np.random.Generator, n: int, L: int, rate: float, scale) -> np.ndarray:
    mask = rng.random((n, L)) < rate
    mask[:, 0] = False
    rows, cols = np.nonzero(mask)
    jumps = np.zeros((n, L), dtype=np.float64)
    if rows.size == 0:
        return jumps
    s = np.asarray(scale, dtype=np.float64)
    event_scale = s if s.ndim == 0 else s.reshape(n)[rows]
    jumps[rows, cols] = rng.normal(0.0, 1.0, size=rows.size) * event_scale
    return jumps


def _row_standardize(x: np.ndarray) -> np.ndarray:
    x = x - x.mean(axis=1, keepdims=True)
    sd = x.std(axis=1, keepdims=True)
    return x / np.where(sd < 1e-12, 1.0, sd)


def _measurement_artifacts(
    rng: np.random.Generator,
    block: np.ndarray,
    *,
    preserve_nonnegative: bool,
    rate_scale: float = 1.0,
) -> np.ndarray:
    """Sparse real-measurement effects; ``rate_scale`` multiplies base rates."""
    original = np.asarray(block, dtype=np.float64)
    out = original.copy()
    n, L = out.shape
    rs = float(np.clip(rate_scale, 0.0, 3.0))

    reverse = rng.random(n) < (0.06 * rs)
    out[reverse] = out[reverse, ::-1]

    if not preserve_nonnegative:
        invert = rng.random(n) < (0.04 * rs)
        out[invert] *= -1.0

    for row in np.nonzero(rng.random(n) < (0.06 * rs))[0]:
        q = float(rng.uniform(0.03, 0.18))
        if rng.random() < 0.5:
            out[row] = np.minimum(out[row], np.quantile(out[row], 1.0 - q))
        else:
            out[row] = np.maximum(out[row], np.quantile(out[row], q))

    quantized = np.nonzero(rng.random(n) < (0.07 * rs))[0]
    if quantized.size:
        x = out[quantized]
        lo = x.min(axis=1, keepdims=True)
        hi = x.max(axis=1, keepdims=True)
        levels = rng.integers(16, 257, size=(quantized.size, 1))
        step = (hi - lo) / np.maximum(levels - 1, 1)
        safe_step = np.where(step < 1e-12, 1.0, step)
        out[quantized] = lo + np.rint((x - lo) / safe_step) * safe_step

    held = np.nonzero(rng.random(n) < (0.04 * rs))[0]
    if held.size:
        factors = rng.choice([2, 4, 8], size=held.size, p=[0.55, 0.30, 0.15])
        for factor in (2, 4, 8):
            rows = held[factors == factor]
            if rows.size:
                out[rows] = np.repeat(
                    out[rows, ::factor], factor, axis=1
                )[:, :L]
    degenerate = out.std(axis=1) < 1e-9
    out[degenerate] = original[degenerate]
    return out


# ── stem builders (  lineage) ────────────────────────────────────────


def _trend_seasonal_ar(rng: np.random.Generator, n: int, L: int, *,
                       hi_frac: float = 0.25, exc_lo: float = 0.4,
                       exc_hi: float = 3.0, clean_frac: float = 0.4,
                       clean_lo: float = 0.02, clean_hi: float = 0.12) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    level = rng.normal(0.0, 1.0, size=(n, 1))
    _hi = rng.random((n, 1)) < hi_frac
    exc = np.where(_hi, rng.normal(0.0, exc_hi, size=(n, 1)),
                   rng.normal(0.0, exc_lo, size=(n, 1)))
    tn = t / max(L - 1, 1)
    series = level + exc * tn + _seasonal(rng, n, L)
    phi = rng.uniform(0.0, 0.85, size=n)
    clean = rng.random((n, 1)) < clean_frac
    sigma = np.where(
        clean,
        rng.uniform(clean_lo, clean_hi, size=(n, 1)),
        rng.uniform(0.1, 0.6, size=(n, 1)),
    )
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    return series + _ar1_batch(innov, phi)


def _regime_shift(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    level = np.cumsum(_sparse_jumps(rng, n, L, rate=3.0 / L, scale=2.0), axis=1)
    log_vol = np.cumsum(_sparse_jumps(rng, n, L, rate=3.0 / L, scale=0.5), axis=1)
    vol = np.exp(np.clip(log_vol, -3.0, 3.0)) * rng.uniform(0.1, 0.5, size=(n, 1))
    noise = rng.normal(0.0, 1.0, size=(n, L)) * vol
    seas = _seasonal(rng, n, L, k_max=2) * rng.uniform(0.0, 1.0, size=(n, 1))
    slope = rng.normal(0.0, 1.0 / L, size=(n, 1)) + np.cumsum(
        _sparse_jumps(rng, n, L, rate=2.0 / L, scale=4.0 / L), axis=1
    )
    piecewise_trend = np.cumsum(slope, axis=1)
    return level + piecewise_trend + seas + noise


def _multiplicative(rng: np.random.Generator, n: int, L: int, *,
                    hi_frac: float = 0.25, exc_lo: float = 0.3, exc_hi: float = 2.0) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    _hg = rng.random((n, 1)) < hi_frac
    gexc = np.where(_hg, 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_level = np.exp(gexc * tn + rng.normal(0.0, 0.3, size=(n, 1)))
    amp = rng.uniform(0.1, 0.6, size=(n, 1))
    seasonal_shape = _seasonal(rng, n, L, k_max=1)
    seasonal_sd = seasonal_shape.std(axis=1, keepdims=True)
    seasonal_shape /= np.where(seasonal_sd < 1e-12, 1.0, seasonal_sd)
    seas = 1.0 + amp * seasonal_shape
    noise = 1.0 + rng.normal(0.0, 1.0, size=(n, L)) * rng.uniform(0.02, 0.15, size=(n, 1))
    scale = rng.uniform(1.0, 50.0, size=(n, 1))
    return scale * base_level * np.clip(seas, 0.05, None) * np.clip(noise, 0.05, None)


def _ar2(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    p1 = rng.uniform(0.3, 0.98, size=n)
    p2 = rng.uniform(-0.6, 0.6, size=n)
    a2 = p2
    a1 = p1 * (1.0 - p2)
    sigma = rng.uniform(0.2, 0.8, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    x = _ar2_batch(innov, a1, a2)
    drift = rng.normal(0.0, 0.005, size=(n, 1)) * np.arange(L, dtype=np.float64)[None, :]
    return x + drift


def _integrated(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    order2 = rng.random(n) < 0.35
    drift = rng.normal(0.0, 0.02, size=(n, 1))
    sigma = rng.uniform(0.2, 1.0, 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)
    o2 = order2[:, None]
    return np.where(o2, walk2 / max(L, 1) ** 0.5, walk)


def _threshold_ar(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    phi_hi = rng.uniform(0.3, 0.9, size=n)
    phi_lo = rng.uniform(-0.9, 0.3, size=n)
    const_hi = rng.normal(0.0, 0.3, size=n)
    const_lo = rng.normal(0.0, 0.3, size=n)
    sigma = rng.uniform(0.2, 0.7, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    x = np.empty((n, L), dtype=np.float64)
    x[:, 0] = innov[:, 0]
    for t in range(1, L):
        prev = x[:, t - 1]
        hi = prev >= 0.0
        phi = np.where(hi, phi_hi, phi_lo)
        const = np.where(hi, const_hi, const_lo)
        x[:, t] = np.clip(const + phi * prev + innov[:, t], -1e6, 1e6)
    return x


def _chaotic(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    use_sine = rng.random(n) < 0.5
    r_log = rng.uniform(3.6, 4.0, size=n)
    r_sin = rng.uniform(0.85, 1.0, size=n)
    x0 = rng.uniform(0.05, 0.95, size=n)
    x = np.empty((n, L), dtype=np.float64)
    cur = x0.copy()
    x[:, 0] = cur
    for t in range(1, L):
        nxt_log = r_log * cur * (1.0 - cur)
        nxt_sin = r_sin * np.sin(np.pi * cur)
        cur = np.where(use_sine, nxt_sin, nxt_log)
        cur = np.clip(cur, 0.0, 1.0)
        x[:, t] = cur
    return x


def _spectral_gp(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    f = np.fft.rfftfreq(L)[None, :]
    lengthscale = np.exp(rng.uniform(np.log(8.0), np.log(256.0), size=(n, 1)))
    envelope = np.exp(-0.5 * (2.0 * np.pi * lengthscale * f) ** 2)
    z = rng.standard_normal((n, f.shape[1])) + 1j * rng.standard_normal((n, f.shape[1]))
    z[:, 0] = 0.0
    x = np.fft.irfft(z * np.sqrt(envelope), n=L, axis=1)
    sd = x.std(axis=1, keepdims=True)
    return x / np.where(sd < 1e-12, 1.0, sd)


def _long_memory(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    f = np.fft.rfftfreq(L)
    safe_f = np.maximum(f, 1.0 / L)[None, :]
    beta = rng.uniform(-0.6, 2.4, size=(n, 1))
    amp = safe_f ** (-0.5 * beta)
    multiscale = rng.random((n, 1)) < 0.4
    split_idx = rng.integers(8, max(9, f.size // 3), size=(n, 1))
    split_f = np.maximum(split_idx / L, 1.0 / L)
    beta_hi = rng.uniform(-0.6, 2.8, size=(n, 1))
    above = np.arange(f.size)[None, :] > split_idx
    amp_hi = split_f ** (-0.5 * beta) \
        * (safe_f / split_f) ** (-0.5 * beta_hi)
    amp = np.where(multiscale & above, amp_hi, amp)
    amp[:, 0] = 0.0
    z = rng.standard_normal((n, f.size)) + 1j * rng.standard_normal((n, f.size))
    x = np.fft.irfft(z * amp, n=L, axis=1)
    integrate = rng.random(n) < 0.25
    if integrate.any():
        x[integrate] = np.cumsum(x[integrate], axis=1)
    x -= x.mean(axis=1, keepdims=True)
    sd = x.std(axis=1, keepdims=True)
    return x / np.where(sd < 1e-12, 1.0, sd)


def _ou_stochastic_vol(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    switch_rate = np.exp(rng.uniform(np.log(0.001), np.log(0.15), size=(n, 1)))
    switches = rng.random((n, L)) < switch_rate
    switches[:, 0] = rng.random(n) < 0.5
    regime = np.bitwise_and(np.cumsum(switches, axis=1), 1).astype(np.int8)

    slow = rng.random((n, 1)) < 0.5
    phi = np.where(
        slow,
        rng.uniform(0.995, 0.9995, size=(n, 1)),
        rng.uniform(0.90, 0.99, size=(n, 1)),
    )
    mu0 = rng.normal(-2.0, 1.0, size=(n, 1))
    mu1 = rng.normal(2.0, 1.0, size=(n, 1))
    mean = np.where(regime == 0, mu0, mu1)
    seasonal_on = rng.random((n, 1)) < 0.6
    mean += seasonal_on * _seasonal(rng, n, L, k_max=3) \
        * rng.uniform(0.5, 3.0, size=(n, 1))

    sigma0 = rng.lognormal(np.log(0.3), 0.3, size=(n, 1))
    sigma1 = rng.lognormal(np.log(1.5), 0.5, size=(n, 1))
    base_sigma = np.where(regime == 0, sigma0, sigma1)
    log_vol = np.cumsum(
        _sparse_jumps(rng, n, L, rate=8.0 / L, scale=0.35), axis=1
    )
    log_vol -= log_vol.mean(axis=1, keepdims=True)
    vol = base_sigma * np.exp(np.clip(log_vol, -1.5, 1.5))

    eps = rng.standard_normal((n, L))
    heavy = np.nonzero(rng.random(n) < 0.35)[0]
    if heavy.size:
        eps[heavy] = (
            rng.standard_t(4.0, size=(heavy.size, L)) / np.sqrt(2.0)
        )
    shocks = rng.random((n, L)) < (3.0 / L)
    shock_rows, shock_cols = np.nonzero(shocks)
    eps[shock_rows, shock_cols] += rng.normal(
        0.0, 5.0, size=shock_rows.size
    )

    innovation_scale = np.sqrt(np.maximum(1.0 - phi * phi, 1e-6))
    drive = (1.0 - phi) * mean + innovation_scale * vol * eps
    out = np.empty((n, L), dtype=np.float64)
    out[:, 0] = mean[:, 0] + vol[:, 0] * eps[:, 0]
    for i in range(n):
        p = float(phi[i, 0])
        out[i, 1:] = lfilter(
            [1.0], [1.0, -p], drive[i, 1:], zi=[p * out[i, 0]]
        )[0]

    scale = np.exp(rng.uniform(np.log(0.1), np.log(50.0), size=(n, 1)))
    shift = rng.uniform(-100.0, 100.0, size=(n, 1))
    return out * scale + shift


def _physical_sensors(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    seasonal = _seasonal(rng, n, L, k_max=2)
    smooth = _spectral_gp(rng, n, L)
    fronts = np.cumsum(
        _sparse_jumps(rng, n, L, rate=5.0 / L, scale=1.0), axis=1
    )
    base = (
        seasonal * rng.uniform(0.3, 2.0, size=(n, 1))
        + smooth * rng.uniform(0.2, 1.2, size=(n, 1))
        + fronts * rng.uniform(0.2, 1.0, size=(n, 1))
    )

    kind = rng.integers(0, 4, size=n)
    out = base.copy()

    bounded = kind == 1
    if bounded.any():
        gain = rng.uniform(0.8, 3.5, size=(int(bounded.sum()), 1))
        midpoint = rng.uniform(-0.8, 0.8, size=(int(bounded.sum()), 1))
        out[bounded] = 100.0 / (1.0 + np.exp(-gain * (base[bounded] - midpoint)))

    pressure = kind == 2
    if pressure.any():
        count = int(pressure.sum())
        walk = np.cumsum(rng.standard_normal((count, L)), axis=1) / np.sqrt(L)
        level = rng.uniform(900.0, 1100.0, size=(count, 1))
        out[pressure] = level + rng.uniform(2.0, 15.0, size=(count, 1)) * walk \
            + 2.0 * fronts[pressure] + 0.5 * seasonal[pressure]

    magnitude = kind == 3
    if magnitude.any():
        count = int(magnitude.sum())
        gusts = (rng.random((count, L)) < (8.0 / L)) \
            * rng.lognormal(0.0, 0.8, size=(count, L))
        power = rng.uniform(1.0, 1.6, size=(count, 1))
        out[magnitude] = np.abs(base[magnitude]) ** power + gusts

    return out


def _seasonal_counts(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    period = rng.choice(
        _SEASONAL_PERIODS, size=(n, 1), p=_SEASONAL_PROBS
    )
    phase = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
    amp = rng.uniform(0.15, 0.8, size=(n, 1))
    log_rate = amp * np.sin(2.0 * np.pi * t / period + phase)
    second = rng.random((n, 1)) < 0.55
    log_rate += second * (0.5 * amp) * np.sin(
        4.0 * np.pi * t / period + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
    )
    calendar = rng.random((n, 1)) < 0.35
    day_period = rng.choice([24, 48, 96, 144], size=(n, 1))
    day_idx = (np.floor_divide(np.arange(L)[None, :], day_period) % 7).astype(np.int64)
    day_factors = rng.normal(0.0, 0.12, size=(n, 7))
    day_factors[:, 5:] += rng.uniform(-0.8, 0.3, size=(n, 1))
    calendar_effect = np.take_along_axis(day_factors, day_idx, axis=1)
    log_rate += calendar * calendar_effect
    excursion = rng.uniform(-0.5, 0.5, size=(n, 1))
    log_rate += excursion * t / max(L - 1, 1)

    impulses = (
        (rng.random((n, L)) < (2.0 / L))
        * rng.uniform(1.0, 10.0, size=(n, L))
    )
    burst = _ar1_batch(impulses, rng.uniform(0.85, 0.995, size=(n, 1)))
    base = np.exp(rng.uniform(np.log(3.0), np.log(3000.0), size=(n, 1)))
    lam = base * np.exp(np.clip(log_rate, -5.0, 5.0)) * (1.0 + burst)
    np.clip(lam, 0.0, 1.0e7, out=lam)

    overdispersed = rng.random((n, 1)) < 0.5
    shape = rng.uniform(0.5, 4.0, size=(n, 1))
    mixed = lam * rng.gamma(shape, 1.0 / shape, size=(n, L))
    return rng.poisson(np.where(overdispersed, mixed, lam)).astype(np.float64)


def _intermittent(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    base_p = rng.uniform(0.03, 0.35, size=(n, 1))
    period = rng.choice([7.0, 12.0, 24.0, 48.0, 168.0], size=(n, 1))
    season = rng.uniform(0.2, 1.2, size=(n, 1)) * np.sin(
        2.0 * np.pi * t / period + rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
    )
    logit = np.log(base_p / (1.0 - base_p)) + season
    p = 1.0 / (1.0 + np.exp(-logit))
    occur = (rng.random((n, L)) < p).astype(np.float64)
    magnitude = (
        rng.gamma(shape=2.0, scale=1.0, size=(n, L))
        * rng.uniform(1.0, 10.0, size=(n, 1))
        * np.exp(0.25 * season)
    )
    baseline = rng.uniform(0.0, 0.5, size=(n, 1))
    return baseline + occur * magnitude


def _pulse_outlier(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    base = _spectral_gp(rng, n, L) * rng.uniform(0.5, 2.0, size=(n, 1))
    base += _seasonal(rng, n, L, k_max=1) * rng.uniform(0.0, 1.0, size=(n, 1))
    sharp = _sparse_jumps(
        rng, n, L, rate=3.0 / L, scale=rng.uniform(3.0, 8.0, size=n)
    )
    impulses = _sparse_jumps(
        rng, n, L, rate=2.0 / L, scale=rng.uniform(2.0, 7.0, size=n)
    )
    recovery = _ar1_batch(impulses, rng.uniform(0.75, 0.995, size=n))
    series = base + sharp + recovery

    starts = rng.random((n, L)) < (2.0 / L)
    starts[:, 0] = False
    for row in range(n):
        for start in np.nonzero(starts[row])[0]:
            run = int(rng.integers(3, 65))
            end = min(int(start) + run, L)
            series[row, start:end] = series[row, start - 1]
    return series


# ── Grammar-of-Time productions ─────────────────────────────────────────────
# Stem pool for compositions. Order is fixed so RNG draw sequences stay stable
# across config-only weight changes to non-GoT families.


# Lightweight stem pool for GoT productions. Full   builders stay as
# top-level families; compositions need many stems per row, so these stay
# FFT/AR/seasonal only — no Python-over-t recurrences.
def _stem_ar_seasonal(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    seas = _seasonal(rng, n, L, k_max=2)
    phi = rng.uniform(0.1, 0.9, size=n)
    sigma = rng.uniform(0.15, 0.7, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    return _row_standardize(seas * rng.uniform(0.3, 1.2, size=(n, 1)) + _ar1_batch(innov, phi))


def _stem_integrated(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    sigma = rng.uniform(0.2, 1.0, size=(n, 1))
    walk = np.cumsum(rng.normal(0.0, 1.0, size=(n, L)) * sigma, axis=1)
    return _row_standardize(walk)


def _sample_stem(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """Draw one cheap stem family per row and fill a (n, L) block."""
    builders = (
        _stem_ar_seasonal,
        _spectral_gp,
        _long_memory,
        _stem_integrated,
        _ar2,
    )
    fam = rng.integers(0, len(builders), size=n)
    out = np.empty((n, L), dtype=np.float64)
    for k, builder in enumerate(builders):
        idx = np.nonzero(fam == k)[0]
        if idx.size == 0:
            continue
        out[idx] = _row_standardize(builder(rng, int(idx.size), L))
    return out


def _got_compose(rng: np.random.Generator, n: int, L: int, *,
                 depth: int = 3, mul_frac: float = 0.35,
                 hi_frac: float = 0.25, exc_lo: float = 0.4,
                 exc_hi: float = 2.5) -> np.ndarray:
    """Production: Stem ⊕ Stem [⊕ Stem] — additive or multiplicative phrase.

    Each row stacks ``depth`` standardised stems. A minority use multiplicative
    agreement (level × seasonal-like factor), matching TempoPFN-style compound
    structure rather than a single process family.
    """
    depth = int(np.clip(depth, 2, 4))
    # Always draw ``depth`` stems so the RNG stream is depth-stable.
    stems = [_sample_stem(rng, n, L) for _ in range(depth)]
    n_active = rng.integers(2, depth + 1, size=n)

    out = np.zeros((n, L), dtype=np.float64)
    use_mul = rng.random(n) < mul_frac
    for j, stem in enumerate(stems):
        active = (n_active > j)[:, None]
        w = rng.uniform(0.4, 1.6, size=(n, 1))
        # Additive branch.
        add_mask = active & (~use_mul[:, None])
        out = np.where(add_mask, out + w * stem, out)
        # Multiplicative branch: first stem is the carrier; later stems modulate.
        if j == 0:
            out = np.where(use_mul[:, None], stem, out)
        else:
            factor = 1.0 + 0.35 * w * stem
            out = np.where(active & use_mul[:, None], out * factor, out)

    # Mild length-normalised trend affix so compose rows still carry forecastable
    # drift without exploding scale.
    t = np.arange(L, dtype=np.float64)[None, :] / max(L - 1, 1)
    _hi = rng.random((n, 1)) < hi_frac
    exc = np.where(_hi, rng.normal(0.0, exc_hi * 0.5, size=(n, 1)),
                   rng.normal(0.0, exc_lo * 0.5, size=(n, 1)))
    out = out + exc * t

    # Sparse punctuation affix (jumps) on a minority of rows.
    punct = rng.random(n) < 0.4
    if punct.any():
        jumps = np.cumsum(
            _sparse_jumps(rng, n, L, rate=2.5 / L, scale=1.5), axis=1
        )
        out[punct] = out[punct] + jumps[punct]

    scale = np.exp(rng.uniform(np.log(0.2), np.log(40.0), size=(n, 1)))
    shift = rng.uniform(-50.0, 50.0, size=(n, 1))
    return out * scale + shift


def _got_splice(rng: np.random.Generator, n: int, L: int, *,
                n_cuts: int = 2, hi_frac: float = 0.25,
                exc_lo: float = 0.4, exc_hi: float = 2.5) -> np.ndarray:
    """Production: Stem ‖ Stem — clause boundaries splice different dynamics.

    Unlike  's piecewise level jumps inside one process, each clause is
    an independent stem; breakpoints teach structural change of *generating law*.
    """
    n_cuts = int(np.clip(n_cuts, 1, 4))
    n_clauses = n_cuts + 1
    clauses = [_sample_stem(rng, n, L) for _ in range(n_clauses)]

    # Cut positions in (0.15L, 0.85L), sorted per row.
    cuts = np.sort(
        rng.integers(max(1, L // 8), max(2, (7 * L) // 8), size=(n, n_cuts)),
        axis=1,
    )
    # Enforce strictly increasing cuts with a small gap.
    for c in range(1, n_cuts):
        cuts[:, c] = np.maximum(cuts[:, c], cuts[:, c - 1] + max(8, L // 32))
    cuts = np.clip(cuts, 1, L - 2)

    out = clauses[0].copy()
    t_idx = np.arange(L)[None, :]
    for c in range(n_cuts):
        after = t_idx >= cuts[:, c:c + 1]
        out = np.where(after, clauses[c + 1], out)

    # Soft blend near each cut so the splice is a transition, not a hard glitch
    # (real regime changes often ramp over a few steps).
    blend_w = max(4, L // 128)
    for c in range(n_cuts):
        cut = cuts[:, c:c + 1]
        dist = (t_idx - cut).astype(np.float64)
        gate = np.clip(0.5 + dist / (2.0 * blend_w), 0.0, 1.0)
        left = clauses[c]
        right = clauses[c + 1]
        near = np.abs(dist) <= blend_w
        blended = (1.0 - gate) * left + gate * right
        out = np.where(near, blended, out)

    # Optional level offset between clauses (structural break magnitude).
    level_jump = rng.normal(0.0, 1.5, size=(n, n_cuts))
    for c in range(n_cuts):
        after = t_idx >= cuts[:, c:c + 1]
        out = np.where(after, out + level_jump[:, c:c + 1], out)

    t = np.arange(L, dtype=np.float64)[None, :] / max(L - 1, 1)
    _hi = rng.random((n, 1)) < hi_frac
    exc = np.where(_hi, rng.normal(0.0, exc_hi * 0.4, size=(n, 1)),
                   rng.normal(0.0, exc_lo * 0.4, size=(n, 1)))
    out = out + exc * t

    scale = np.exp(rng.uniform(np.log(0.2), np.log(40.0), size=(n, 1)))
    shift = rng.uniform(-50.0, 50.0, size=(n, 1))
    return out * scale + shift


def _got_nested(rng: np.random.Generator, n: int, L: int, *,
                nest_ratio: float = 6.0) -> np.ndarray:
    """Production: Envelope ⋉ Carrier — slow scale nests a fast carrier.

    Hierarchical seasonality / synoptic weather / business-cycle nesting: a
    smooth long-scale envelope modulates amplitude (and sometimes phase) of a
    faster seasonal or AR carrier. This is the multi-scale grammar Chronos-style
    priors emphasise but   only touches via modulated seasonality.
    """
    nest_ratio = float(np.clip(nest_ratio, 2.0, 24.0))
    t = np.arange(L, dtype=np.float64)[None, :]

    # Slow envelope: spectral GP with long lengthscale, or long sinusoid.
    use_gp_env = rng.random(n) < 0.55
    env = np.empty((n, L), dtype=np.float64)
    gp_rows = np.nonzero(use_gp_env)[0]
    sin_rows = np.nonzero(~use_gp_env)[0]
    if gp_rows.size:
        # Force long lengthscales for the envelope.
        f = np.fft.rfftfreq(L)[None, :]
        lengthscale = np.exp(
            rng.uniform(np.log(64.0), np.log(min(512.0, L / 2.0)), size=(gp_rows.size, 1))
        )
        envelope = np.exp(-0.5 * (2.0 * np.pi * lengthscale * f) ** 2)
        z = rng.standard_normal((gp_rows.size, f.shape[1])) + 1j * rng.standard_normal(
            (gp_rows.size, f.shape[1])
        )
        z[:, 0] = 0.0
        g = np.fft.irfft(z * np.sqrt(envelope), n=L, axis=1)
        env[gp_rows] = _row_standardize(g)
    if sin_rows.size:
        per = rng.uniform(L / nest_ratio, L / 1.5, size=(sin_rows.size, 1))
        phase = rng.uniform(0.0, 2.0 * np.pi, size=(sin_rows.size, 1))
        env[sin_rows] = np.sin(2.0 * np.pi * t / per + phase)

    # Fast carrier: seasonal bank and/or AR(2).
    carrier = _seasonal(rng, n, L, k_max=3)
    carrier = _row_standardize(carrier)
    mix_ar = rng.random(n) < 0.45
    if mix_ar.any():
        ar = _row_standardize(_ar2(rng, n, L))
        w = rng.uniform(0.3, 0.7, size=(n, 1))
        carrier = np.where(mix_ar[:, None], w * carrier + (1.0 - w) * ar, carrier)

    amp = 1.0 + rng.uniform(0.3, 1.4, size=(n, 1)) * env
    # Phase wobble as a small quadrature mix with a lagged carrier — vectorised,
    # no per-row roll. Equivalent spirit: slow envelope nudges fast phase.
    carrier_lag = np.empty_like(carrier)
    carrier_lag[:, 0] = carrier[:, 0]
    carrier_lag[:, 1:] = carrier[:, :-1]
    wobble = rng.uniform(0.0, 0.35, size=(n, 1)) * env
    out = amp * (carrier + wobble * carrier_lag)

    # Residual noise scaled by envelope intensity (prosody).
    sigma = rng.uniform(0.05, 0.35, size=(n, 1)) * (0.5 + 0.5 * np.abs(env))
    phi = rng.uniform(0.0, 0.8, size=n)
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    out = out + _ar1_batch(innov, phi)

    scale = np.exp(rng.uniform(np.log(0.2), np.log(40.0), size=(n, 1)))
    shift = rng.uniform(-50.0, 50.0, size=(n, 1))
    return out * scale + shift


def _delay_batch(x: np.ndarray, lags: np.ndarray) -> np.ndarray:
    """Causal delay with edge hold, vectorised over a small lag vocabulary.

    Rows sharing a lag are shifted in one slice copy — O(#unique_lags) passes
    instead of a Python loop over n.
    """
    n, L = x.shape
    out = np.empty_like(x)
    # Hold initial value for the lag prefix.
    for lag in np.unique(lags):
        rows = np.nonzero(lags == lag)[0]
        if rows.size == 0:
            continue
        lag_i = int(lag)
        block = x[rows]
        delayed = np.empty_like(block)
        delayed[:, :lag_i] = block[:, :1]
        delayed[:, lag_i:] = block[:, :-lag_i]
        out[rows] = delayed
    return out


def _got_causal(rng: np.random.Generator, n: int, L: int, *,
                lag_frac: float = 0.08) -> np.ndarray:
    """Production: Driver ▷ Response — lagged temporal causal chain.

    Inspired by Chronos-2 / CauKer temporal causal graphs, specialised to a
    univariate observable: the emitted series is a response driven by a latent
    driver with a drawn lag and FIR-like coupling, plus its own AR residual.
    Teaches lead-lag structure that pure mixture families never emit.
    """
    lag_frac = float(np.clip(lag_frac, 0.01, 0.25))
    # Discrete lag menu keeps _delay_batch on a handful of unique values.
    lag_menu = np.unique(
        np.clip(
            (np.array([0.01, 0.02, 0.04, 0.06, 0.08, 0.12, 0.16, 0.20]) * L).astype(np.int64),
            1,
            max(1, int(L * lag_frac)),
        )
    )
    lags = rng.choice(lag_menu, size=n)

    driver = _sample_stem(rng, n, L)
    use_parent2 = rng.random(n) < 0.4
    parent2 = _sample_stem(rng, n, L)

    a0 = rng.uniform(0.2, 1.2, size=(n, 1))
    a1 = rng.uniform(0.3, 1.5, size=(n, 1))
    b = rng.uniform(0.2, 1.0, size=(n, 1))
    lags2 = rng.choice(lag_menu, size=n)

    delayed = _delay_batch(driver, lags)
    resp = a0 * driver + a1 * delayed
    if use_parent2.any():
        delayed2 = _delay_batch(parent2, lags2)
        resp = np.where(use_parent2[:, None], resp + b * delayed2, resp)

    phi = rng.uniform(0.2, 0.9, size=n)
    sigma = rng.uniform(0.1, 0.5, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    resp = resp + _ar1_batch(innov, phi)
    seas_on = rng.random((n, 1)) < 0.5
    resp = resp + seas_on * _seasonal(rng, n, L, k_max=2) * rng.uniform(
        0.1, 0.8, size=(n, 1)
    )

    scale = np.exp(rng.uniform(np.log(0.2), np.log(40.0), size=(n, 1)))
    shift = rng.uniform(-50.0, 50.0, size=(n, 1))
    return resp * scale + shift


# ── final safety gate ───────────────────────────────────────────────────────


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