"""new_generator — a compositional SARIMA-backbone synthetic generator for cascade.

A pure SARIMA(p,d,q)(P,D,Q)_s prior is classic but *narrow*: linear, near-Gaussian
and (bar the explicit integration) stationary. State-of-the-art synthetic priors
for training time-series foundation models win on the dynamics SARIMA structurally
cannot produce — regime changes, volatility clustering, fat tails, non-Gaussian /
positive marginals, compositional structure and mixed series. This generator keeps
SARIMA as the *dynamical backbone* (its guaranteed-stationary recursion is a
valuable, well-behaved core) and layers the SOTA ingredients on top:

1. **Regime switching** — a series is a concatenation of 1–3 segments, each its own
   SARIMA parameterisation, stitched with level continuity. This produces the
   structural breaks in dynamics that real operational/financial series show and a
   single stationary ARMA never does.
2. **Stochastic volatility** — innovations are optionally modulated by an AR(1)
   log-variance process (vectorised), giving genuine volatility *clustering* on top
   of Gaussian or Student-t (fat-tailed) shocks.
3. **Compositional components** — additive nested-calendar seasonality, smooth
   trends, structural level shifts and sparse spikes/outliers, each scaled relative
   to the core so the mix stays balanced (KernelSynth-style composition).
4. **Nonlinear warps** — the linear core is optionally passed through an invertible
   transform (asinh tail-compression, exp/softplus positivity, signed power), so
   marginals become non-Gaussian / positive / multiplicative — the shapes real
   demand, price and count series have.
5. **TSMixup** — with some probability two independent draws are convex-combined
   (the Chronos augmentation), broadening the corpus beyond any single prior.

Everything is fully vectorised (one or a few ``scipy.signal.lfilter`` passes per
series, no per-timestep Python loops), CPU-only, and comfortably inside the
trainer's generation deadline.

Stationarity guarantee: AR/MA coefficients (regular and seasonal) are drawn as
reflection coefficients (partial autocorrelations) in (-1, 1) and mapped to lag
polynomials by the Levinson-Durbin recursion, so every root lies outside the unit
circle and the ARMA recursion is always stable. Only the explicit ``d``/``D``
integration and the warps introduce (intended) nonstationarity/nonlinearity;
magnitudes are bounded by de-meaning before integration, exponent clipping in the
warps, and a final rescale-to-target-scale + hard clip.

Determinism: the corpus is a pure function of ``(seed, n_series)``. Every per-series
sub-seed is derived from the master seed via ``np.random.SeedSequence``; no
``hash()``, wall-clock, or unseeded global RNG.
"""

from __future__ import annotations

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

import numpy as np
from scipy.signal import lfilter

from cascade.interface import DataGenerator

# Fixed stream id for the length-drawing RNG, kept distinct from any per-series
# sub-seed so lengths are order-deterministic and independent of model draws.
_LENGTH_STREAM_ID = 0x1E_2757

# Candidate seasonal periods (timesteps): intraday (4/24/48), business/weekly
# (5/7/168), monthly (12/30) and longer rhythms. A period is used only when the
# series is long enough for >= 3 full cycles.
_SEASONAL_PERIODS: tuple[int, ...] = (4, 5, 7, 12, 24, 30, 48, 96, 144, 168, 336)


class Generator(DataGenerator):
    """SARIMA-backbone compositional corpus 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._seed = int(seed)
        self._min_len = int(cfg.get("min_length", 64))
        self._max_len = int(cfg.get("max_length", 4096))
        if not (1 <= self._min_len <= self._max_len):
            raise ValueError(f"invalid length band [{self._min_len}, {self._max_len}]")

        # SARIMA order caps (per-series orders drawn in [0, cap]).
        self._max_p = int(cfg.get("max_ar", 4))
        self._max_q = int(cfg.get("max_ma", 4))
        self._max_P = int(cfg.get("max_seasonal_ar", 2))
        self._max_Q = int(cfg.get("max_seasonal_ma", 2))

        # Integration-order weights (index i = probability mass for order i).
        self._d_weights = np.asarray(cfg.get("d_weights", [0.5, 0.38, 0.12]), dtype=np.float64)
        self._seasonal_d_weights = np.asarray(
            cfg.get("seasonal_d_weights", [0.7, 0.3]), dtype=np.float64
        )

        # Enrichment probabilities.
        self._seasonal_prob = float(cfg.get("seasonal_prob", 0.5))
        self._student_t_prob = float(cfg.get("student_t_prob", 0.35))
        self._stoch_vol_prob = float(cfg.get("stoch_vol_prob", 0.4))
        self._max_regimes = int(cfg.get("max_regimes", 3))
        self._regime_prob = float(cfg.get("regime_prob", 0.5))
        self._trend_prob = float(cfg.get("trend_prob", 0.5))
        self._calendar_prob = float(cfg.get("calendar_prob", 0.45))
        self._level_shift_prob = float(cfg.get("level_shift_prob", 0.35))
        self._spike_prob = float(cfg.get("spike_prob", 0.3))
        self._warp_prob = float(cfg.get("warp_prob", 0.45))
        self._mixup_prob = float(cfg.get("mixup_prob", 0.25))

        # Sanitisation knobs.
        self._max_abs = float(cfg.get("max_abs_value", 1.0e6))
        self._standardize = bool(cfg.get("standardize", False))

    @property
    def name(self) -> str:
        return "sarima-compositional-v2"

    # ── deterministic sub-seeding ────────────────────────────────────────────
    def _series_rng(self, index: int) -> np.random.Generator:
        return np.random.default_rng(np.random.SeedSequence([self._seed, index]))

    # ── stationary lag-polynomial construction ───────────────────────────────
    @staticmethod
    def _reflection_to_poly(rng: np.random.Generator, order: int, lo: float, hi: float) -> np.ndarray:
        """Reflection coefficients (PACF) in (lo, hi) → lag polynomial via
        Levinson-Durbin. Guarantees all roots outside the unit circle."""
        if order <= 0:
            return np.array([1.0])
        kappa = rng.uniform(lo, hi, size=order)
        phi = np.zeros(order, dtype=np.float64)
        for m in range(order):
            k = kappa[m]
            prev = phi[:m].copy()
            phi[m] = k
            if m > 0:
                phi[:m] = prev - k * prev[::-1]
        return np.concatenate(([1.0], -phi))

    @staticmethod
    def _expand_seasonal(poly: np.ndarray, period: int) -> np.ndarray:
        """Lift a lag polynomial in ``L`` to one in ``L**period``."""
        if poly.size <= 1 or period <= 1:
            return poly.copy()
        out = np.zeros((poly.size - 1) * period + 1, dtype=np.float64)
        out[0] = poly[0]
        for i in range(1, poly.size):
            out[i * period] = poly[i]
        return out

    def _choose_period(self, rng: np.random.Generator, length: int) -> int:
        if rng.random() >= self._seasonal_prob:
            return 1
        candidates = [p for p in _SEASONAL_PERIODS if p * 3 <= length]
        return int(rng.choice(candidates)) if candidates else 1

    @staticmethod
    def _weighted_order(rng: np.random.Generator, weights: np.ndarray) -> int:
        w = np.clip(weights, 0.0, None)
        total = w.sum()
        return int(rng.choice(len(w), p=w / total)) if total > 0.0 else 0

    # ── innovations: fat tails + stochastic volatility (vectorised) ──────────
    def _draw_innovations(self, rng: np.random.Generator, n: int) -> np.ndarray:
        if rng.random() < self._student_t_prob:
            df = float(rng.uniform(3.0, 12.0))
            eps = rng.standard_t(df, size=n)
        else:
            eps = rng.standard_normal(n)
        eps = eps * float(rng.uniform(0.3, 2.0))

        # Stochastic volatility: multiply by exp(0.5 * AR(1) log-variance) —
        # a vectorised lfilter draw gives volatility clustering with no Python loop.
        if rng.random() < self._stoch_vol_prob:
            phi_v = float(rng.uniform(0.9, 0.995))
            v_innov = rng.normal(0.0, float(rng.uniform(0.1, 0.4)), size=n)
            log_var = lfilter([1.0], [1.0, -phi_v], v_innov)
            log_var -= log_var.mean()
            eps = eps * np.exp(0.5 * np.clip(log_var, -6.0, 6.0))
        return eps

    # ── one stationary+integrated SARIMA segment ─────────────────────────────
    def _sarima_segment(self, rng: np.random.Generator, length: int, period: int) -> np.ndarray:
        seasonal = period > 1
        p = int(rng.integers(0, self._max_p + 1))
        q = int(rng.integers(0, self._max_q + 1))
        d = self._weighted_order(rng, self._d_weights)
        if seasonal:
            P = int(rng.integers(0, self._max_P + 1))
            Q = int(rng.integers(0, self._max_Q + 1))
            D = self._weighted_order(rng, self._seasonal_d_weights)
        else:
            P = Q = D = 0
        if p == q == P == Q == d == D == 0:
            p = 1 if rng.random() < 0.5 else 0
            q = 0 if p else 1

        ar = self._reflection_to_poly(rng, p, -0.95, 0.95)
        ma = self._reflection_to_poly(rng, q, -0.9, 0.9)
        if seasonal:
            ar = np.convolve(ar, self._expand_seasonal(self._reflection_to_poly(rng, P, -0.95, 0.95), period))
            ma = np.convolve(ma, self._expand_seasonal(self._reflection_to_poly(rng, Q, -0.9, 0.9), period))

        memory = max(ar.size, ma.size, period * max(P, D, 1))
        burn = int(min(2048, max(128, 4 * memory)))
        innov = self._draw_innovations(rng, length + burn)
        y = lfilter(ma, ar, innov)[burn:]

        for _ in range(d):
            y = np.cumsum(y - y.mean())
        if seasonal and D > 0:
            seas_int = np.zeros(period + 1, dtype=np.float64)
            seas_int[0], seas_int[period] = 1.0, -1.0
            for _ in range(D):
                y = lfilter([1.0], seas_int, y - y.mean())
        return y

    # ── regime-switching core: stitch segments with level continuity ─────────
    def _regime_core(self, rng: np.random.Generator, length: int) -> np.ndarray:
        period = self._choose_period(rng, length)
        n_regimes = 1
        if rng.random() < self._regime_prob and length >= 96:
            n_regimes = int(rng.integers(2, self._max_regimes + 1))

        if n_regimes == 1:
            return self._sarima_segment(rng, length, period)

        # Partition length into n_regimes contiguous segments (each >= 32).
        cuts = np.sort(rng.choice(np.arange(32, length - 32), size=n_regimes - 1, replace=False)) \
            if length - 64 > n_regimes else np.array([], dtype=int)
        bounds = [0, *cuts.tolist(), length]
        segs, offset = [], 0.0
        for a, b in zip(bounds[:-1], bounds[1:], strict=True):
            seg_len = b - a
            if seg_len <= 0:
                continue
            # Occasionally re-roll seasonality per regime for richer breaks.
            seg_period = period if rng.random() < 0.7 else self._choose_period(rng, seg_len)
            core = self._sarima_segment(rng, seg_len, seg_period)
            core = core - core[0] + offset
            segs.append(core)
            offset = core[-1]
        x = np.concatenate(segs) if segs else self._sarima_segment(rng, length, period)
        return x[:length]

    # ── additive compositional components (scaled to the core) ───────────────
    def _add_components(self, rng: np.random.Generator, x: np.ndarray, length: int) -> np.ndarray:
        t = np.arange(length, dtype=np.float64)
        scale = x.std()
        if scale <= 1e-12:
            scale = 1.0

        # Nested-calendar seasonality (short period nested in ~7x/~30x multiples).
        if rng.random() < self._calendar_prob:
            base_period = rng.uniform(4.0, 48.0)
            for ratio in (1.0, 7.0, 30.0):
                period = base_period * ratio
                if period >= length * 1.5:
                    continue
                amp = scale * rng.uniform(0.2, 1.2) / ratio
                phase = rng.uniform(0.0, 2.0 * np.pi)
                x = x + amp * np.sin(2.0 * np.pi * t / period + phase)

        # Smooth deterministic trend (linear + occasional curvature).
        if rng.random() < self._trend_prob:
            u = t / max(1.0, length - 1)
            x = x + scale * rng.normal(0.0, 1.0) * u
            if rng.random() < 0.4:
                x = x + scale * rng.normal(0.0, 0.7) * (u - 0.5) ** 2

        # Structural level shifts.
        if rng.random() < self._level_shift_prob:
            for _ in range(int(rng.integers(1, 4))):
                at = int(rng.integers(length // 10, max(length // 10 + 1, 9 * length // 10)))
                x = x.copy()
                x[at:] += scale * rng.normal(0.0, 1.0)

        # Sparse spikes / outliers.
        if rng.random() < self._spike_prob:
            n_spikes = int(rng.integers(1, max(2, length // 200 + 2)))
            idx = rng.integers(0, length, size=n_spikes)
            x = x.copy()
            x[idx] += scale * rng.normal(0.0, 4.0, size=n_spikes)
        return x

    # ── optional invertible nonlinear warp → non-Gaussian / positive marginals ─
    def _maybe_warp(self, rng: np.random.Generator, x: np.ndarray) -> np.ndarray:
        if rng.random() >= self._warp_prob:
            return x
        mu, sd = x.mean(), x.std()
        z = (x - mu) / sd if sd > 1e-12 else x - mu
        kind = rng.integers(0, 4)
        if kind == 0:  # tail compression
            return np.arcsinh(rng.uniform(0.5, 3.0) * z)
        if kind == 1:  # log-normal-like positivity / multiplicative
            return np.exp(np.clip(rng.uniform(0.2, 1.0) * z, -10.0, 10.0))
        if kind == 2:  # softplus positivity (demand/count-like)
            return np.log1p(np.exp(np.clip(rng.uniform(0.5, 1.5) * z, -20.0, 20.0)))
        gamma = rng.uniform(1.5, 3.0)  # signed power (peaky)
        return np.sign(z) * np.abs(z) ** gamma

    # ── full single draw ─────────────────────────────────────────────────────
    def _draw_one(self, rng: np.random.Generator, length: int) -> np.ndarray:
        x = self._regime_core(rng, length)
        x = self._add_components(rng, x, length)
        x = self._maybe_warp(rng, x)
        return x

    # ── scaling + sanitisation ───────────────────────────────────────────────
    def _finalize(self, rng: np.random.Generator, y: np.ndarray, length: int) -> np.ndarray:
        x = np.asarray(y, dtype=np.float64).ravel()
        if x.size != length:
            if x.size > length:
                x = x[:length]
            else:
                x = np.concatenate([x, np.full(length - x.size, x[-1] if x.size else 0.0)])
        if not np.isfinite(x).all():
            x = np.nan_to_num(x, nan=0.0, posinf=self._max_abs, neginf=-self._max_abs)

        std = x.std()
        if std > 1e-12:
            x = x / std * float(rng.lognormal(mean=0.0, sigma=1.2))
        x = x + rng.normal(0.0, 3.0)

        if self._standardize:
            std2 = x.std()
            if std2 > 1e-12:
                x = (x - x.mean()) / std2

        np.clip(x, -self._max_abs, self._max_abs, out=x)
        if not np.isfinite(x).all():
            x = rng.standard_normal(length)
        return np.ascontiguousarray(x, dtype=np.float64)

    # ── entrypoint ───────────────────────────────────────────────────────────
    def generate(self, n_series: int) -> Iterator[np.ndarray]:
        if n_series <= 0:
            return
        len_rng = np.random.default_rng(np.random.SeedSequence([self._seed, _LENGTH_STREAM_ID]))
        for i in range(n_series):
            length = int(len_rng.integers(self._min_len, self._max_len + 1))
            rng = self._series_rng(i + 1)
            x = self._draw_one(rng, length)
            # TSMixup: convex-combine two independent draws (Chronos augmentation).
            if rng.random() < self._mixup_prob:
                x2 = self._draw_one(rng, length)
                w = float(rng.uniform(0.2, 0.8))
                sx, s2 = x.std() or 1.0, x2.std() or 1.0
                x = w * (x / sx) + (1.0 - w) * (x2 / s2)
            yield self._finalize(rng, x, length)