"""
Noise generation: from Gaussian samples to Perlin fields.

Synthesis is the other half of noise: where PSD estimation measures what is
there, this module builds noise with exactly the properties you specify.
Every method is implemented from first principles; no exotic dependencies.

Methods
-------
Gaussian:
* ``box_muller(n, seed)``: Box-Muller transform, two uniforms → two Gaussians
* ``ziggurat(n, seed)``: Ziggurat algorithm (fast sampling from monotonic PDFs)

Coloured:
* ``colored_noise_ar(n, alpha, p, seed)``: AR(p) filter, Kasdin coefficients
* ``colored_noise_fd(n, alpha, seed)``: frequency-domain shaping
* ``voss_mccartney(n, octaves, seed)``: Voss-McCartney 1/f generator

Procedural:
* ``perlin_1d(x, seed, octaves, persistence)``: 1-D Perlin noise
* ``perlin_2d(x, y, seed, octaves, persistence)``: 2-D Perlin noise
* ``simplex_2d(x, y, seed)``: 2-D Simplex noise (lower complexity, fewer artifacts)

Verification:
* ``verify_noise(x, fs, label)``: run ``characterize()`` and print a report

References
----------
* Box, G. E. P. & Muller, M. E. (1958). A note on the generation of random
  normal deviates. *Ann. Math. Statist.*, 29(2), 610-611.
* Marsaglia, G. & Tsang, W. W. (2000). The ziggurat method for generating
  random variables. *J. Stat. Soft.*, 5(8).
* Perlin, K. (1985). An image synthesizer. *ACM SIGGRAPH*, 19(3), 287-296.
* Perlin, K. (2002). Improving noise. *ACM SIGGRAPH*, 21(3), 681-682.
* Voss, R. F. & Clarke, J. (1978). "1/f noise" in music: music from 1/f noise.
  *J. Acoust. Soc. Am.*, 63(1), 258-263.
"""

import numpy as np
from scipy.signal import lfilter


# ---------------------------------------------------------------------------
# Box-Muller: two uniforms → two independent Gaussians
# ---------------------------------------------------------------------------

def box_muller(n: int, seed: int | None = None) -> np.ndarray:
    """Generate Gaussian samples via the Box-Muller transform.

    Draws pairs of independent uniform(0,1) variates and transforms them
    into independent N(0,1) variates.  Simple, exact, and sufficient for
    most DSP work.  Slower than Ziggurat for very large N but rarely the
    bottleneck.

    Parameters
    ----------
    n : int
        Number of samples.
    seed : int or None
        RNG seed.

    Returns
    -------
    x : ndarray of shape (n,)
        Independent N(0,1) samples.
    """
    rng = np.random.default_rng(seed)
    # Round up to an even count.
    m = (n + 1) // 2
    u1 = rng.random(m)
    u2 = rng.random(m)
    r = np.sqrt(-2.0 * np.log(u1))
    theta = 2.0 * np.pi * u2
    z0 = r * np.cos(theta)
    z1 = r * np.sin(theta)
    z = np.column_stack([z0, z1]).ravel()
    return z[:n]


# ---------------------------------------------------------------------------
# Ziggurat: fast sampling from monotonic-decreasing PDFs
# ---------------------------------------------------------------------------

def ziggurat(n: int, seed: int | None = None) -> np.ndarray:
    """Generate Gaussian samples via the Ziggurat algorithm.

    The Ziggurat method (Marsaglia & Tsang, 2000) partitions the area under
    the Gaussian PDF into horizontal strips.  Most samples fall in the base
    strip (a fast rectangular rejection test); only a small fraction (~1.5%)
    requires an exponential tail evaluation.  Production RNG libraries
    (Julia, MATLAB, GNU Octave) use this algorithm.

    This implementation delegates to ``np.random.Generator.normal()`` for
    correctness: NumPy's default RNG has used Ziggurat since 2019.  The
    function exists so the algorithm has a named home in the module and can
    be discussed in the prose; the docstring and the ``box_muller()``
    implementation cover the educational ground.

    Parameters
    ----------
    n : int
        Number of samples.
    seed : int or None
        RNG seed.

    Returns
    -------
    x : ndarray of shape (n,)
        Independent N(0,1) samples.
    """
    rng = np.random.default_rng(seed)
    return rng.normal(0.0, 1.0, size=n)


def _build_ziggurat_tables(n_layers: int = 256) -> None:
    """Precompute the ziggurat layer edges and weights for N(0,1).

    Uses Marsaglia & Tsang's (2000) recurrence to compute the right edge
    x[i] of each layer.  This is a reference implementation: production
    code should use the highly optimised tables from the paper directly.

    Stores the result as module-level attributes on the ``ziggurat``
    function object.
    """
    from scipy.stats import norm
    n = n_layers
    tail_area = 1.0 / n
    R = norm.ppf(1 - tail_area / 2)

    v_tail = tail_area * R + np.exp(-R**2 / 2) / np.sqrt(2 * np.pi)

    x = np.zeros(n)
    x[0] = R
    for i in range(1, n):
        x_guess = x[i - 1] * 0.99
        for _ in range(15):
            y_i = np.exp(-x_guess**2 / 2) / np.sqrt(2 * np.pi)
            y_prev = np.exp(-x[i-1]**2 / 2) / np.sqrt(2 * np.pi)
            f = x[i - 1] * (y_i - y_prev) - v_tail
            df = -x[i - 1] * x_guess * y_i
            if abs(df) < 1e-15:
                break
            x_guess -= f / df
            if x_guess <= 0:
                x_guess = x[i-1] * 0.1
                break
        x[i] = max(x_guess, 0.0)

    we = np.exp(-x**2 / 2) / np.sqrt(2 * np.pi)
    ziggurat._x = x
    ziggurat._we = we


# ---------------------------------------------------------------------------
# Coloured noise generation
# ---------------------------------------------------------------------------

def colored_noise_ar(n: int, alpha: float = 1.0, p: int = 20,
                     seed: int | None = None) -> np.ndarray:
    """Generate 1/f^alpha noise via an AR(p) filter (Kasdin's method).

    Filters white Gaussian noise through an all-pole filter whose
    coefficients follow the Kasdin recurrence [@kasdin1995discrete]:
        theta_i = ((i - 1 - alpha/2) * theta_{i-1}) / i,  theta_0 = 1.

    Parameters
    ----------
    n : int
        Number of samples.
    alpha : float
        Spectral exponent (0=white, 1=pink, 2=brown).
    p : int
        AR model order (higher = better fit at low frequencies).
    seed : int or None
        RNG seed.

    Returns
    -------
    x : ndarray of shape (n,)
        1/f^alpha noise.
    """
    rng = np.random.default_rng(seed)
    w = rng.standard_normal(n)

    # Kasdin AR coefficients.
    theta = np.ones(p + 1)
    for i in range(1, p + 1):
        theta[i] = ((i - 1 - alpha / 2.0) * theta[i - 1]) / i

    # AR(p): x[t] = w[t] - sum_{i=1}^p theta_i * x[t-i]
    x = lfilter([1.0], theta, w)
    return x / np.std(x)


def colored_noise_fd(n: int, alpha: float = 1.0,
                     seed: int | None = None) -> np.ndarray:
    """Generate 1/f^alpha noise via frequency-domain shaping.

    Shapes white noise in the frequency domain by multiplying the spectrum
    by f^{-alpha/2}, then inverse FFT.  Simpler than the AR method but
    produces an exactly periodic sequence of length n (the last sample
    wraps to the first).

    Parameters
    ----------
    n : int
        Number of samples.
    alpha : float
        Spectral exponent.
    seed : int or None
        RNG seed.

    Returns
    -------
    x : ndarray of shape (n,)
        1/f^alpha noise (real).
    """
    rng = np.random.default_rng(seed)
    w = rng.standard_normal(n)

    W = np.fft.rfft(w)
    freqs = np.fft.rfftfreq(n)
    freqs[0] = freqs[1] * 0.5  # avoid division by zero at DC

    # Shape: multiply by f^{-alpha/2} for f > 0.
    H = np.zeros_like(freqs)
    H[1:] = freqs[1:] ** (-alpha / 2.0)
    H[0] = H[1]  # crude DC handling

    X = W * H
    x = np.fft.irfft(X, n=n)
    return x / np.std(x)


def voss_mccartney(n: int, octaves: int = 8,
                   seed: int | None = None) -> np.ndarray:
    """Generate 1/f (pink) noise via the Voss-McCartney algorithm.

    Maintains *octaves* independent white noise generators running at
    different update rates.  At each sample, the outputs are summed,
    giving a spectrum that is approximately 1/f up to the Nyquist
    frequency.  Simple, integer-friendly, and runs in O(n) time with
    O(octaves) state: ideal for embedded systems.

    Parameters
    ----------
    n : int
        Number of samples.
    octaves : int
        Number of octaves (more = better low-frequency fit).  Default 8.
    seed : int or None
        RNG seed.

    Returns
    -------
    x : ndarray of shape (n,)
        Approximate pink noise.
    """
    rng = np.random.default_rng(seed)
    x = np.zeros(n)
    # State: current value for each octave generator.
    state = rng.standard_normal(octaves)
    # Counter per octave, triggers an update when hitting zero.
    counter = np.ones(octaves, dtype=int)

    for i in range(n):
        for j in range(octaves):
            counter[j] -= 1
            if counter[j] == 0:
                state[j] = rng.standard_normal()
                counter[j] = 1 << j  # update period = 2^j samples
            x[i] += state[j]

    return x / np.std(x)


# ---------------------------------------------------------------------------
# Perlin noise (1-D and 2-D)
# ---------------------------------------------------------------------------

def _perlin_fade(t: np.ndarray) -> np.ndarray:
    """Perlin's improved fade function: 6t^5 - 15t^4 + 10t^3.

    Has zero first and second derivatives at t=0 and t=1, which eliminates
    discontinuities in the derivative of the noise, the spectral signature
    of the original 3t^2 - 2t^3 kernel was a faint but visible second-
    derivative discontinuity at lattice boundaries.
    """
    return t * t * t * (t * (t * 6.0 - 15.0) + 10.0)


def _perlin_grad_1d(hash_val: np.ndarray, x: np.ndarray) -> np.ndarray:
    """1-D Perlin gradient: hash mod 2 picks positive or negative unit slope. Vectorized."""
    return np.where((hash_val & 1) == 0, x, -x)


def _perlin_grad_2d(hash_val: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """2-D Perlin gradient: hash mod 8 picks one of 8 directions. Vectorized."""
    h = hash_val & 7
    # Pick u based on hash: bits 00,01,100,101 → use x; others → use y.
    use_x_for_u = (h == 0) | (h == 1) | (h == 4) | (h == 5)
    u = np.where(use_x_for_u, x, y)
    v = np.where((h >= 0) & (h <= 3), y, x)  # bits: 0,1,2,3 → y; 4,5,6,7 → x
    u = np.where((h & 1) == 0, u, -u)
    v = np.where((h & 2) == 0, v, -v)
    return u + v


def _perlin_permutation(seed: int = 0, size: int = 512) -> np.ndarray:
    """Generate a reproducible permutation table for the hash function."""
    rng = np.random.default_rng(seed)
    p = np.arange(size, dtype=int)
    rng.shuffle(p)
    return np.tile(p, 2)  # doubled for wrap-free indexing


def perlin_1d(x: np.ndarray, seed: int = 0, octaves: int = 1,
              persistence: float = 0.5) -> np.ndarray:
    """1-D Perlin noise.

    Classic gradient noise: a lattice of random unit slopes, smoothly
    interpolated by the quintic fade function.  Octave summation
    (fractal Brownian motion) adds scaled copies at doubled frequencies
    for a controlled 1/f-like spectrum.

    Parameters
    ----------
    x : ndarray
        Input coordinates (1-D array).  Integer spacing gives one lattice
        cell per unit.
    seed : int
        Permutation seed.
    octaves : int
        Number of octave layers (1 = base Perlin only).
    persistence : float
        Amplitude multiplier per octave (0.5 = classic, lower = smoother).

    Returns
    -------
    y : ndarray
        Noise values at each x, in approximately [-1, 1].
    """
    perm = _perlin_permutation(seed)
    x = np.asarray(x, dtype=float)
    result = np.zeros_like(x)

    for octave in range(octaves):
        freq = 1 << octave  # 2^octave
        amp = persistence ** octave
        x_scaled = x * freq

        # Integer lattice cell.
        xi = np.floor(x_scaled).astype(int)
        xf = x_scaled - xi  # fractional part [0, 1)

        # Hash the two bounding lattice points.
        h0 = perm[xi % 512]
        h1 = perm[(xi + 1) % 512]

        # Fade the fractional coordinate.
        u = _perlin_fade(xf)

        # Gradient dot products and interpolation.
        g0 = _perlin_grad_1d(h0, xf)
        g1 = _perlin_grad_1d(h1, xf - 1.0)
        result += amp * (g0 + u * (g1 - g0))

    # Normalise to roughly [-1, 1].
    amp_sum = np.sum([persistence ** o for o in range(octaves)])
    return result / amp_sum


def perlin_2d(x: np.ndarray, y: np.ndarray, seed: int = 0, octaves: int = 1,
              persistence: float = 0.5) -> np.ndarray:
    """2-D Perlin noise.

    Parameters
    ----------
    x, y : ndarray
        Coordinate arrays (must broadcast together).  Integer spacing gives
        one lattice cell per unit.
    seed : int
        Permutation seed.
    octaves : int
        Number of octave layers.
    persistence : float
        Amplitude multiplier per octave.

    Returns
    -------
    z : ndarray
        Noise values at each (x, y), in approximately [-1, 1].
    """
    perm = _perlin_permutation(seed)
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)
    result = np.zeros_like(x)

    for octave in range(octaves):
        freq = 1 << octave
        amp = persistence ** octave
        xs = x * freq
        ys = y * freq

        # Integer lattice corners.
        xi = np.floor(xs).astype(int)
        yi = np.floor(ys).astype(int)
        xf = xs - xi
        yf = ys - yi

        # Hash the four lattice corners.
        xi_mod = xi % 512
        yi_mod = yi % 512
        h00 = perm[perm[xi_mod] + yi_mod]
        h10 = perm[perm[(xi_mod + 1) % 512] + yi_mod]
        h01 = perm[perm[xi_mod] + (yi_mod + 1) % 512]
        h11 = perm[perm[(xi_mod + 1) % 512] + (yi_mod + 1) % 512]

        # Fade the fractional coordinates.
        u = _perlin_fade(xf)
        v = _perlin_fade(yf)

        # Gradient dot products at four corners, then bilinear interpolation.
        n00 = _perlin_grad_2d(h00, xf, yf)
        n10 = _perlin_grad_2d(h10, xf - 1.0, yf)
        n01 = _perlin_grad_2d(h01, xf, yf - 1.0)
        n11 = _perlin_grad_2d(h11, xf - 1.0, yf - 1.0)

        nx0 = n00 + u * (n10 - n00)
        nx1 = n01 + u * (n11 - n01)
        result += amp * (nx0 + v * (nx1 - nx0))

    amp_sum = np.sum([persistence ** o for o in range(octaves)])
    return result / amp_sum


# ---------------------------------------------------------------------------
# Simplex noise (2-D): lower complexity, fewer directional artifacts
# ---------------------------------------------------------------------------

def simplex_2d(x: np.ndarray, y: np.ndarray, seed: int = 0) -> np.ndarray:
    """2-D Simplex noise: lower complexity than Perlin with fewer
    directional artifacts.

    Simplex noise works on a triangular (simplex) grid rather than a
    square lattice, which reduces the number of gradient contributions
    per point and eliminates the axis-aligned visual bias Perlin noise
    exhibits.  This is a simplified implementation for educational use.

    Parameters
    ----------
    x, y : ndarray
        Coordinate arrays (must broadcast).
    seed : int
        Permutation seed.

    Returns
    -------
    z : ndarray
        Noise values at each (x, y), approximately in [-1, 1].
    """
    perm = _perlin_permutation(seed)
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)

    # Skew input into simplex grid.
    F2 = 0.5 * (np.sqrt(3.0) - 1.0)
    G2 = (3.0 - np.sqrt(3.0)) / 6.0
    s = (x + y) * F2
    xi = np.floor(x + s).astype(int)
    yi = np.floor(y + s).astype(int)

    t = (xi + yi).astype(float) * G2
    x0 = x - (xi - t)
    y0 = y - (yi - t)

    # Determine which simplex we are in.
    i1 = (x0 > y0).astype(int)
    j1 = 1 - i1

    # Offsets for the three simplex corners.
    x1 = x0 - i1 + G2
    y1 = y0 - j1 + G2
    x2 = x0 - 1.0 + 2.0 * G2
    y2 = y0 - 1.0 + 2.0 * G2

    # Hash the corners.
    xi_mod = xi % 512
    yi_mod = yi % 512
    h0 = perm[perm[xi_mod] + yi_mod]
    h1 = perm[perm[(xi_mod + i1) % 512] + (yi_mod + j1) % 512]
    h2 = perm[perm[(xi_mod + 1) % 512] + (yi_mod + 1) % 512]

    # Contribution from each corner (radial kernel).
    def _contrib(h, dx, dy):
        t0 = 0.5 - dx*dx - dy*dy
        mask = t0 > 0
        t02 = t0 * t0
        grad = _perlin_grad_2d(h, dx, dy)
        return np.where(mask, t02 * t02 * grad, 0.0)

    n0 = _contrib(h0, x0, y0)
    n1 = _contrib(h1, x1, y1)
    n2 = _contrib(h2, x2, y2)

    # Scale to roughly [-1, 1].
    return (n0 + n1 + n2) * 70.0


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------

def verify_noise(x: np.ndarray, fs: float = 1.0, label: str = 'noise') -> dict:
    """Characterise generated noise and print a verification report.

    Delegates to ``random_processes.characterize()``.  Returns the info
    dict for further programmatic use.

    Parameters
    ----------
    x : ndarray
        Noise samples.
    fs : float
        Sample rate.
    label : str
        Human-readable label for the printout.

    Returns
    -------
    info : dict
        From ``random_processes.characterize()``.
    """
    from scipy.stats import kurtosis, skew
    from scipy.signal import welch

    info = {}
    info['mean'] = float(np.mean(x))
    info['std'] = float(np.std(x))
    info['skewness'] = float(skew(x))
    info['excess_kurtosis'] = float(kurtosis(x))

    xc = x - info['mean']
    acf = np.correlate(xc, xc, mode='full')
    acf = acf / acf[len(xc) - 1]
    mid = len(xc) - 1
    for lag, key in [(1, 'acf_lag1'), (2, 'acf_lag2')]:
        info[key] = float(acf[mid + lag]) if lag < len(x) else np.nan

    nperseg = min(256, len(x) // 4)
    if nperseg >= 16:
        f, psd = welch(x, fs, nperseg=nperseg)
        mask = f > 0
        if mask.sum() > 4:
            coeffs = np.polyfit(np.log10(f[mask]), np.log10(psd[mask]), 1)
            info['psd_slope'] = float(coeffs[0])
            pred = np.polyval(coeffs, np.log10(f[mask]))
            ss_res = np.sum((np.log10(psd[mask]) - pred)**2)
            ss_tot = np.sum((np.log10(psd[mask]) - np.mean(np.log10(psd[mask])))**2)
            info['psd_r2'] = float(1 - ss_res/ss_tot) if ss_tot > 0 else 0.0
        else:
            info['psd_slope'] = np.nan; info['psd_r2'] = np.nan
    else:
        info['psd_slope'] = np.nan; info['psd_r2'] = np.nan

    print(f"\n=== {label} ===")
    print(f"  Samples: {len(x)}")
    print(f"  Mean:    {info['mean']:.4f}")
    print(f"  Std:     {info['std']:.4f}")
    print(f"  Skew:    {info['skewness']:.3f}")
    print(f"  Kurt:    {info['excess_kurtosis']:.3f}")
    print(f"  ACF(1):  {info['acf_lag1']:.4f}")
    print(f"  ACF(2):  {info['acf_lag2']:.4f}")
    print(f"  PSD slope: {info['psd_slope']:.2f}  (R²={info['psd_r2']:.2f})")

    return info