# tools/automaster_app/pipeline/saturation.py
"""Shared 4x-oversampled nonlinearity wrapper + master glue saturation.

Every nonlinear stage in the pipeline routes through oversampled() so harmonic
generation happens at 4x rate and fold-back lands above the audio band before
the decimation filter removes it.
"""
import numpy as np
from scipy.fft import next_fast_len
from scipy.signal import resample


def oversampled(fn, audio, factor=4):
    squeeze = False
    if audio.ndim == 1:
        squeeze = True
        audio = audio[:, None]
    n = audio.shape[0]
    n_fast = next_fast_len(n)
    if n_fast > n:
        pad = np.zeros((n_fast - n, audio.shape[1]), dtype=audio.dtype)
        work = np.concatenate([audio, pad], axis=0)
    else:
        work = audio
    up = resample(work, n_fast * factor, axis=0)
    out = fn(up)
    down = resample(out, n_fast, axis=0)
    result = down[:n]
    if squeeze:
        result = result[:, 0]
    return result


def glue_saturate(audio, amount, drive=1.5):
    """Master-bus glue: arctan soft saturation, parallel blend.
    amount: 0..1 blend (spec: 5-10% typical)."""
    if amount <= 0.0:
        return audio
    norm = np.arctan(drive)
    wet = oversampled(lambda u: np.arctan(u * drive) / norm, audio, factor=4)
    return wet * amount + audio * (1.0 - amount)