"""Mixture-of-priors generator for Cascade.

Design goals:
* Compete on the synthetic prior, not architecture — Toto2 is fixed.
* Cover regimes common in GIFT-like held-out pools (nature, sales, econ/fin,
  energy, web/ops): seasonality, persistence, breaks, intermittency, smooth GP
  curves, and a thin slice of chaotic dynamics.
* Stay deterministic, code-only, finite, and fast (numba + chunked stream) so
  more tokens hit the trainer before the wall.

Notable knobs / behaviors:
* calendar-biased seasonal periods (7/12/24/48/168 heavy)
* heteroskedastic AR noise on the seasonal family
* multi-scale random-Fourier GP (short + long lengthscales)
* richer chaotic maps (logistic / sine / tent)
* clustered intermittent demand
* optional mild seasonality on integrated walks
* config: ``seasonal_focus``, ``het_noise_frac``
"""

from __future__ import annotations

import json
from collections.abc import Iterator
from functools import partial
from pathlib import Path

import numpy as np
from numba import njit

from cascade.interface import DataGenerator

_CHUNK = 512

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

# Defaults biased toward multi-domain real pools (nature+sales+econ+energy).
_DEFAULT_WEIGHTS: dict[str, float] = {
    "trend_seasonal_ar": 0.18,
    "regime_shift": 0.12,
    "multiplicative": 0.14,
    "ar2": 0.14,
    "integrated": 0.14,
    "threshold_ar": 0.05,
    "chaotic": 0.05,
    "rff_gp": 0.10,
    "intermittent": 0.05,
    "pulse_outlier": 0.03,
}

# Period menu + default sampling mass (Chronos-like calendar priors).
_PERIODS = np.array(
    [4, 7, 12, 24, 30, 48, 52, 96, 144, 168, 336, 504, 672, 720],
    dtype=np.float64,
)
# Uniform baseline over periods; ``seasonal_focus`` blends toward this prior.
_PERIOD_PRIOR = np.array(
    [0.02, 0.14, 0.10, 0.22, 0.04, 0.10, 0.04, 0.06, 0.04, 0.12, 0.04, 0.03, 0.03, 0.02],
    dtype=np.float64,
)
_PERIOD_PRIOR = _PERIOD_PRIOR / _PERIOD_PRIOR.sum()


# ── numba kernels ───────────────────────────────────────────────────────────


@njit(cache=False)
def _ar1_batch_jit(innov: np.ndarray, phi: np.ndarray) -> np.ndarray:
    n, L = innov.shape
    x = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        p = phi[i]
        x[i, 0] = innov[i, 0]
        for t in range(1, L):
            x[i, t] = p * x[i, t - 1] + innov[i, t]
    return x


@njit(cache=False)
def _ar2_batch_jit(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):
        aa1 = a1[i]
        aa2 = a2[i]
        x[i, 0] = innov[i, 0]
        if L > 1:
            x[i, 1] = aa1 * x[i, 0] + innov[i, 1]
        for t in range(2, L):
            x[i, t] = aa1 * x[i, t - 1] + aa2 * x[i, t - 2] + innov[i, t]
    return x


@njit(cache=False)
def _threshold_ar_jit(
    innov: np.ndarray,
    phi_hi: np.ndarray,
    phi_lo: np.ndarray,
    const_hi: np.ndarray,
    const_lo: np.ndarray,
) -> np.ndarray:
    n, L = innov.shape
    x = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        x[i, 0] = innov[i, 0]
        for t in range(1, L):
            prev = x[i, t - 1]
            if prev >= 0.0:
                phi = phi_hi[i]
                const = const_hi[i]
            else:
                phi = phi_lo[i]
                const = const_lo[i]
            v = const + phi * prev + innov[i, t]
            if v > 1e6:
                v = 1e6
            elif v < -1e6:
                v = -1e6
            x[i, t] = v
    return x


@njit(cache=False)
def _chaotic_jit(
    map_id: np.ndarray,
    r_a: np.ndarray,
    r_b: np.ndarray,
    x0: np.ndarray,
    L: int,
) -> np.ndarray:
    """map_id: 0=logistic, 1=sine, 2=tent (DynaMix-inspired breadth)."""
    n = x0.shape[0]
    x = np.empty((n, L), dtype=np.float64)
    for i in range(n):
        cur = x0[i]
        x[i, 0] = cur
        mid = map_id[i]
        a = r_a[i]
        b = r_b[i]
        for t in range(1, L):
            if mid == 1:
                cur = a * np.sin(np.pi * cur)
            elif mid == 2:
                # tent map on [0,1]; ``a`` in ~[1.1, 2.0]
                if cur < 0.5:
                    cur = a * cur
                else:
                    cur = a * (1.0 - cur)
            else:
                cur = a * cur * (1.0 - cur)
            if cur < 0.0:
                cur = 0.0
            elif cur > 1.0:
                cur = 1.0
            # unused ``b`` kept for RNG-shape stability across map types
            cur = cur + 0.0 * b
            x[i, t] = cur
    return x


@njit(cache=False)
def _apply_holds_jit(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]


class Generator(DataGenerator):
    """Mixture-of-priors 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", 3.0))
        self._gr_exc_lo = float(cfg.get("gr_exc_lo", 0.3))
        self._gr_exc_hi = float(cfg.get("gr_exc_hi", 2.0))
        # 0 = uniform periods; 1 = full calendar prior (7/12/24/48/168-heavy).
        self._seasonal_focus = float(np.clip(cfg.get("seasonal_focus", 0.7), 0.0, 1.0))
        # Fraction of seasonal series that get slow heteroskedastic vol.
        self._het_noise_frac = float(np.clip(cfg.get("het_noise_frac", 0.45), 0.0, 1.0))
        self._fixed_len = self._min_len == self._max_len

    @property
    def name(self) -> str:
        return str(self._cfg.get("name", "persist-heavy"))

    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
        period_p = (1.0 - self._seasonal_focus) * (np.ones_like(_PERIOD_PRIOR) / len(_PERIOD_PRIOR))
        period_p = period_p + self._seasonal_focus * _PERIOD_PRIOR
        period_p = period_p / period_p.sum()

        builders = (
            partial(
                _trend_seasonal_ar,
                hi_frac=self._tr_hi_frac,
                exc_lo=self._tr_exc_lo,
                exc_hi=self._tr_exc_hi,
                period_p=period_p,
                het_frac=self._het_noise_frac,
            ),
            partial(_regime_shift, period_p=period_p),
            partial(
                _multiplicative,
                hi_frac=self._tr_hi_frac,
                exc_lo=self._gr_exc_lo,
                exc_hi=self._gr_exc_hi,
            ),
            _ar2,
            partial(_integrated, period_p=period_p),
            _threshold_ar,
            _chaotic,
            _rff_gp,
            _intermittent,
            partial(_pulse_outlier, period_p=period_p),
        )
        produced = 0
        while produced < n_series:
            # Full chunk every time → series i is a pure function of (seed, i).
            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 = _sanitize(builders[fam](rng, int(idx.size), max_len))
                if self._fixed_len:
                    for row, series_i in enumerate(idx):
                        chunk[int(series_i)] = np.ascontiguousarray(block[row], dtype=np.float64)
                else:
                    for row, series_i in enumerate(idx):
                        L = int(lengths[series_i])
                        chunk[int(series_i)] = np.ascontiguousarray(
                            block[row, :L], dtype=np.float64
                        )
            take = min(_CHUNK, n_series - produced)
            for arr in chunk[:take]:
                if arr is None:  # pragma: no cover
                    raise RuntimeError("internal: unfilled series slot")
                yield arr
            produced += take


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


def _ar1_batch(innov: np.ndarray, phi: np.ndarray) -> np.ndarray:
    return _ar1_batch_jit(np.ascontiguousarray(innov), np.ascontiguousarray(phi.reshape(-1)))


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


def _seasonal(
    rng: np.random.Generator,
    n: int,
    L: int,
    k_max: int = 3,
    period_p: np.ndarray | None = None,
) -> np.ndarray:
    t = np.arange(L, dtype=np.float64)[None, :]
    p = _PERIOD_PRIOR if period_p is None else period_p
    k = rng.integers(1, k_max + 1, size=n)
    out = np.zeros((n, L), dtype=np.float64)
    for j in range(k_max):
        active = (k > j).astype(np.float64)[:, None]
        per = rng.choice(_PERIODS, size=n, p=p)[:, None]
        amp = rng.uniform(0.2, 2.0, size=n)[:, None]
        phase = rng.uniform(0.0, 2.0 * np.pi, size=n)[:, None]
        out += active * amp * np.sin(2.0 * np.pi * t / per + phase)
    return out


def _sparse_jumps(rng: np.random.Generator, n: int, L: int, rate: float, scale) -> np.ndarray:
    mask = 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]
    jumps = mask * mag * s
    jumps[:, 0] = 0.0
    return jumps


# ── family builders ─────────────────────────────────────────────────────────


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,
    period_p: np.ndarray | None = None,
    het_frac: float = 0.45,
) -> 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, period_p=period_p)
    phi = rng.uniform(0.0, 0.85, size=n)
    sigma = rng.uniform(0.1, 0.6, size=(n, 1))
    innov = rng.normal(0.0, 1.0, size=(n, L)) * sigma
    # Heteroskedastic slow vol (GARCH-lite). Always draw vol for all n so the
    # RNG stream (and series i) does not depend on how many rows flipped het.
    log_vol = np.cumsum(_sparse_jumps(rng, n, L, rate=2.0 / L, scale=0.35), axis=1)
    vol = np.exp(np.clip(log_vol, -2.0, 2.0))
    het = (rng.random(n) < het_frac)[:, None]
    innov = innov * np.where(het, vol, 1.0)
    return series + _ar1_batch(innov, phi)


def _regime_shift(
    rng: np.random.Generator,
    n: int,
    L: int,
    *,
    period_p: np.ndarray | None = None,
) -> 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, period_p=period_p) * rng.uniform(0.0, 1.0, size=(n, 1))
    return level + 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))
    # Sales/energy-friendly periods (incl. 48 for sub-daily energy).
    seas = 1.0 + amp * np.sin(
        2.0 * np.pi * t / rng.choice([7.0, 12.0, 24.0, 48.0, 52.0, 168.0], size=n)[:, None]
        + rng.uniform(0.0, 2 * np.pi, size=(n, 1))
    )
    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,
    *,
    period_p: np.ndarray | None = None,
) -> 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]
    y = np.where(o2, walk2 / max(L, 1) ** 0.5, walk)
    # Mild seasonal overlay on ~40% — always draw seas/amp for all n (determinism).
    seas = _seasonal(rng, n, L, k_max=2, period_p=period_p)
    amp = rng.uniform(0.1, 0.8, size=(n, 1))
    mask = (rng.random(n) < 0.40)[:, None]
    return y + np.where(mask, seas * amp, 0.0)


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
    return _threshold_ar_jit(
        np.ascontiguousarray(innov),
        np.ascontiguousarray(phi_hi),
        np.ascontiguousarray(phi_lo),
        np.ascontiguousarray(const_hi),
        np.ascontiguousarray(const_lo),
    )


def _chaotic(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    # 0 logistic, 1 sine, 2 tent — small curated chaotic prior (DynaMix lesson).
    # Draw all parameter banks for every series (fixed RNG shape), then select.
    map_id = rng.integers(0, 3, size=n)
    r_log = rng.uniform(3.6, 4.0, size=n)
    r_sin = rng.uniform(0.85, 1.0, size=n)
    r_tent = rng.uniform(1.2, 1.99, size=n)
    r_b = rng.uniform(0.0, 1.0, size=n)  # reserved; keeps draw order stable
    r_a = np.where(map_id == 0, r_log, np.where(map_id == 1, r_sin, r_tent))
    x0 = rng.uniform(0.05, 0.95, size=n)
    return _chaotic_jit(
        np.ascontiguousarray(map_id),
        np.ascontiguousarray(r_a),
        np.ascontiguousarray(r_b),
        np.ascontiguousarray(x0),
        L,
    )


def _rff_gp(rng: np.random.Generator, n: int, L: int, K: int = 48) -> np.ndarray:
    """Multi-scale RFF ≈ mixture of short/long stationary GPs (Chronos-2 style)."""
    t = np.arange(L, dtype=np.float64)[None, :]
    # 50/50 short (local wiggles) vs long (smooth trends) lengthscales.
    short = rng.random((n, 1)) < 0.5
    lengthscale = np.where(
        short,
        rng.uniform(8.0, 40.0, size=(n, 1)),
        rng.uniform(40.0, 220.0, size=(n, 1)),
    )
    acc = np.zeros((n, L), dtype=np.float64)
    scale = np.sqrt(2.0 / K)
    for _ in range(K):
        w = rng.normal(0.0, 1.0, size=(n, 1)) / lengthscale
        b = rng.uniform(0.0, 2.0 * np.pi, size=(n, 1))
        np.add(acc, np.cos(w * t + b), out=acc)
    amp = rng.uniform(0.5, 2.0, size=(n, 1))
    return scale * amp * acc


def _intermittent(rng: np.random.Generator, n: int, L: int) -> np.ndarray:
    """Zero-inflated demand with optional burst clustering (sales-like)."""
    p = rng.uniform(0.05, 0.4, size=(n, 1))
    occur = (rng.random((n, L)) < p).astype(np.float64)
    # Cluster: with 35% probability, widen hits by OR-ing a shifted copy.
    cluster = rng.random(n) < 0.35
    if np.any(cluster):
        shifted = np.zeros_like(occur)
        shifted[:, 1:] = occur[:, :-1]
        occur[cluster] = np.maximum(occur[cluster], shifted[cluster])
    magnitude = rng.gamma(shape=2.0, scale=1.0, size=(n, L)) * rng.uniform(1.0, 10.0, size=(n, 1))
    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,
    *,
    period_p: np.ndarray | None = None,
) -> np.ndarray:
    base = _rff_gp(rng, n, L, K=24) * rng.uniform(0.5, 2.0, size=(n, 1))
    base += _seasonal(rng, n, L, k_max=1, period_p=period_p) * rng.uniform(0.0, 1.0, size=(n, 1))
    pulses = _sparse_jumps(rng, n, L, rate=5.0 / L, scale=rng.uniform(3.0, 8.0, size=n))
    series = base + pulses
    hold = rng.random((n, L)) < (2.0 / L)
    hold[:, 0] = False
    _apply_holds_jit(series, hold)
    return series


def _sanitize(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)