"""zenfro_v2 — weather-tuned mixture + Grammar of Time.

Cascade trains Toto2 from scratch on this corpus alone and scores held-out
windows of length 64 after a 4096-step context. The prior must therefore teach
*forecastable* multi-scale structure under a fixed compute wall.

Evolution from zenfro_v1
-----------------------
v1's weather-heavy blend (diurnal / radiation / cloud) remains the skill
backbone. v2 adds:

  * Grammar of Time (GoT) productions — compose / splice / nested / causal —
    so Toto2 sees temporal *phrases*, not only mutually exclusive families
    (Chronos-2 / TempoPFN / CauKer style compositionality).
  * Coupled calendar seasonality (24↔168, 7↔52, …) and slow amp/phase drift.
  * Clean low-noise seasonal mode for sharp periodic reconstruction.
  * Slightly quieter innovations + retained numba kernels for wall throughput.

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
_AR1_BLOCK = 32

_PRIOR_NAMES: tuple[str, ...] = (
    "trend_seasonal_ar",
    "regime_shift",
    "multiplicative",
    "ar2",
    "integrated",
    "threshold_ar",
    "chaotic",
    "rff_gp",
    "intermittent",
    "pulse_outlier",
    "weather",
    # Grammar-of-Time productions
    "got_compose",
    "got_splice",
    "got_nested",
    "got_causal",
)

# Weather remains the largest single stem; ~22% mass on GoT so combination
# structure is learned without drowning the proven weather signal.
_PRIOR_MASS: dict[str, float] = {
    "weather": 0.34,
    "trend_seasonal_ar": 0.11,
    "regime_shift": 0.06,
    "ar2": 0.08,
    "integrated": 0.08,
    "rff_gp": 0.05,
    "multiplicative": 0.04,
    "threshold_ar": 0.015,
    "chaotic": 0.01,
    "intermittent": 0.012,
    "pulse_outlier": 0.003,
    "got_compose": 0.08,
    "got_splice": 0.06,
    "got_nested": 0.05,
    "got_causal": 0.03,
}

# Broader calendar bank than v1; bias toward diurnal / weekly / monthly.
_SEASON_PERIODS = np.asarray(
    [
        4.0,
        7.0,
        12.0,
        15.0,
        24.0,
        30.0,
        48.0,
        52.0,
        60.0,
        96.0,
        144.0,
        168.0,
        240.0,
        336.0,
        365.0,
        504.0,
        672.0,
        720.0,
    ],
    dtype=np.float64,
)
_SEASON_P = np.asarray(
    [
        0.02,
        0.11,
        0.06,
        0.03,
        0.18,
        0.04,
        0.05,
        0.04,
        0.04,
        0.05,
        0.04,
        0.14,
        0.03,
        0.04,
        0.03,
        0.03,
        0.04,
        0.07,
    ],
    dtype=np.float64,
)
_SEASON_P = _SEASON_P / _SEASON_P.sum()

# Coupled calendar pairs teach daily↔weekly and short↔long interactions.
_SEASON_PAIRS = np.asarray(
    [
        [24.0, 168.0],
        [7.0, 365.0],
        [12.0, 52.0],
        [48.0, 336.0],
        [96.0, 672.0],
        [15.0, 60.0],
        [24.0, 720.0],
    ],
    dtype=np.float64,
)


# ── 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."""
    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 _row_standardize(x: np.ndarray) -> np.ndarray:
    mu = x.mean(axis=1, keepdims=True)
    sd = x.std(axis=1, keepdims=True)
    return (x - mu) / np.where(sd < 1e-9, 1.0, sd)


def _harmonics(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    depth: int = 3,
    couple_frac: float = 0.35,
    modulate_frac: float = 0.30,
) -> np.ndarray:
    """Sum of up to ``depth`` sinusoids with optional calendar coupling + drift."""
    t = np.arange(L, dtype=np.float64)[None, :]
    n_comp = rng.integers(1, depth + 1, size=n)
    pair = _SEASON_PAIRS[rng.integers(0, len(_SEASON_PAIRS), size=n)]
    use_pair = rng.random(n) < couple_frac
    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)
        if j < 2:
            period = np.where(use_pair[:, None], pair[:, j : j + 1], period)
        lo, hi = (0.45, 2.3) if j == 0 else (0.15, 1.15)
        amp = rng.uniform(lo, hi, size=(n, 1))
        phase = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
        arg = 2.0 * np.pi * t / period + phase
        # Minority: slow amplitude / phase modulation (TempoPFN-style complex seas).
        mod_on = (rng.random(n) < modulate_frac) & (n_comp > j)
        if mod_on.any():
            m_per = np.clip(
                period * rng.uniform(4.0, 12.0, size=(n, 1)),
                32.0,
                2.0 * L,
            )
            m_phase = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
            slow = np.sin(2.0 * np.pi * t / m_per + m_phase)
            amp_mod = 1.0 + rng.uniform(0.05, 0.40, size=(n, 1)) * slow
            phase_mod = rng.uniform(0.05, 0.70, size=(n, 1)) * np.sin(
                2.0 * np.pi * t / (1.7 * m_per) - m_phase
            )
            wave = amp * amp_mod * np.sin(arg + phase_mod)
            base = amp * np.sin(arg)
            wave = np.where(mod_on[:, None], wave, base)
        else:
            wave = amp * np.sin(arg)
        acc += on * wave
    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)


def _rescale(x: np.ndarray, rng: np.random.Generator) -> np.ndarray:
    scale = np.exp(rng.uniform(np.log(0.2), np.log(40.0), size=(x.shape[0], 1)))
    shift = rng.uniform(-50.0, 50.0, size=(x.shape[0], 1))
    return x * scale + shift


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


def emit_trend_seasonal(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    hi_frac: float,
    exc_lo: float,
    exc_hi: float,
    clean_frac: float,
    clean_lo: float,
    clean_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)
    # Clean mode: very low innov so periodic reconstruction is sharp.
    clean = rng.random(n) < clean_frac
    phi = rng.uniform(0.35, 0.92, size=n)
    sigma = np.where(
        clean[:, None],
        rng.uniform(clean_lo, clean_hi, size=(n, 1)),
        rng.uniform(0.06, 0.38, 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:
    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))
    seas_shape = _harmonics(rng, n, L, depth=2)
    seas_sd = seas_shape.std(axis=1, keepdims=True)
    seas_shape = seas_shape / np.where(seas_sd < 1e-12, 1.0, seas_sd)
    seas = 1.0 + amp * seas_shape
    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:
    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:
    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."""
    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 mild weekly occurrence modulation."""
    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."""
    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))
        )
        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


# ── Grammar-of-Time productions ─────────────────────────────────────────────
# Lightweight stems only — compositions draw several stems per row, so keep
# them FFT/AR/seasonal (no per-t Python loops beyond the shared AR1 scan).


def _stem_ar_seasonal(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    seas = _harmonics(rng, n, L, depth=2)
    phi = rng.uniform(0.15, 0.90, size=n)
    sigma = rng.uniform(0.12, 0.55, 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(innov, phi))


def _stem_rff(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    return _row_standardize(emit_rff(rng, n, L, features=20))


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 _stem_ar2(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    return _row_standardize(emit_ar2(rng, n, L))


def _stem_weatherish(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """Cheap diurnal+weekly stem (no cloud path) for GoT weather phrases."""
    t = np.arange(L, dtype=np.float64)[None, :]
    a1 = rng.uniform(0.5, 2.0, size=(n, 1))
    p1 = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
    a2 = rng.uniform(0.1, 0.6, size=(n, 1))
    p2 = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
    seas = a1 * np.sin(2.0 * np.pi * t / 24.0 + p1) + a2 * np.sin(2.0 * np.pi * t / 168.0 + p2)
    phi = rng.uniform(0.5, 0.92, size=n)
    sigma = rng.uniform(0.05, 0.25, size=(n, 1))
    return _row_standardize(seas + _ar1(rng.normal(0.0, 1.0, size=(n, L)) * sigma, phi))


def _sample_stem(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    builders = (
        _stem_ar_seasonal,
        _stem_rff,
        _stem_integrated,
        _stem_ar2,
        _stem_weatherish,
    )
    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] = builder(rng, int(idx.size), L)
    return out


def emit_got_compose(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    depth: int,
    mul_frac: float,
    hi_frac: float,
    exc_lo: float,
    exc_hi: float,
) -> np.ndarray:
    """Production: Stem ⊕ Stem [⊕ Stem] — additive or multiplicative phrase."""
    depth = int(np.clip(depth, 2, 4))
    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))
        add_mask = active & (~use_mul[:, None])
        out = np.where(add_mask, out + w * stem, out)
        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)

    t = np.arange(L, dtype=np.float64)[None, :] / max(L - 1, 1)
    heavy = rng.random((n, 1)) < hi_frac
    exc = np.where(
        heavy,
        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

    punct = rng.random(n) < 0.35
    if punct.any():
        jumps = np.cumsum(_jumps(rng, n, L, rate=2.5 / L, scale=1.5), axis=1)
        out[punct] = out[punct] + jumps[punct]
    return _rescale(out, rng)


def emit_got_splice(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    n_cuts: int,
    hi_frac: float,
    exc_lo: float,
    exc_hi: float,
) -> np.ndarray:
    """Production: Stem ‖ Stem — clause boundaries splice different dynamics."""
    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)]
    cuts = np.sort(
        rng.integers(max(1, L // 8), max(2, (7 * L) // 8), size=(n, n_cuts)),
        axis=1,
    )
    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)

    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)
        near = np.abs(dist) <= blend_w
        blended = (1.0 - gate) * clauses[c] + gate * clauses[c + 1]
        out = np.where(near, blended, out)

    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)
    heavy = rng.random((n, 1)) < hi_frac
    exc = np.where(
        heavy,
        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
    return _rescale(out, rng)


def emit_got_nested(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    nest_ratio: float,
) -> np.ndarray:
    """Production: Envelope ⋉ Carrier — slow scale nests a fast carrier."""
    nest_ratio = float(np.clip(nest_ratio, 2.0, 24.0))
    t = np.arange(L, dtype=np.float64)[None, :]

    use_rff_env = rng.random(n) < 0.55
    env = np.empty((n, L), dtype=np.float64)
    rff_rows = np.nonzero(use_rff_env)[0]
    sin_rows = np.nonzero(~use_rff_env)[0]
    if rff_rows.size:
        # Long-lengthscale RFF as envelope (force smooth).
        env[rff_rows] = _row_standardize(
            emit_rff(rng, int(rff_rows.size), L, features=16)
        )
    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)

    carrier = _row_standardize(_harmonics(rng, n, L, depth=3))
    mix_ar = rng.random(n) < 0.45
    if mix_ar.any():
        ar = _row_standardize(emit_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
    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)

    sigma = rng.uniform(0.05, 0.30, 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(innov, phi)
    return _rescale(out, rng)


def _delay_batch(x: np.ndarray, lags: np.ndarray) -> np.ndarray:
    """Causal delay with edge hold, vectorised over unique lag values."""
    n, L = x.shape
    out = np.empty_like(x)
    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 emit_got_causal(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    lag_frac: float,
) -> np.ndarray:
    """Production: Driver ▷ Response — lagged temporal causal chain."""
    lag_frac = float(np.clip(lag_frac, 0.01, 0.25))
    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.45, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    resp = resp + _ar1(innov, phi)
    seas_on = rng.random((n, 1)) < 0.5
    resp = resp + seas_on * _harmonics(rng, n, L, depth=2) * rng.uniform(0.1, 0.8, size=(n, 1))
    return _rescale(resp, rng)


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


class Generator(DataGenerator):
    """Weather + Grammar-of-Time mixture. 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._clean_frac = float(cfg.get("sa_clean_frac", 0.35))
        self._clean_lo = float(cfg.get("sa_clean_lo", 0.02))
        self._clean_hi = float(cfg.get("sa_clean_hi", 0.12))
        self._got_depth = int(cfg.get("got_depth", 3))
        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._fixed = self._min_len == self._max_len

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

    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,
                clean_frac=self._clean_frac,
                clean_lo=self._clean_lo,
                clean_hi=self._clean_hi,
            ),
            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,
            lambda rng, n, L: emit_got_compose(
                rng,
                n,
                L,
                depth=self._got_depth,
                mul_frac=self._got_mul_frac,
                hi_frac=self._tr_hi,
                exc_lo=self._tr_lo,
                exc_hi=self._tr_hi_s,
            ),
            lambda rng, n, L: emit_got_splice(
                rng,
                n,
                L,
                n_cuts=self._got_splice_cuts,
                hi_frac=self._tr_hi,
                exc_lo=self._tr_lo,
                exc_hi=self._tr_hi_s,
            ),
            lambda rng, n, L: emit_got_nested(rng, n, L, nest_ratio=self._got_nest_ratio),
            lambda rng, n, L: emit_got_causal(
                rng, n, L, lag_frac=self._got_causal_lag_frac
            ),
        )

    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