"""
AOS Post-DSP: Pure tensor math operations on output PCM.
Separate module for clarity. Called by app.py run_inference().
"""
import numpy as np
import time

def volume_clamp(pcm, gain):
    """Element-wise hard clamp: pcm * gain clipped to [-1.0, 1.0]"""
    if gain >= 1.0:
        return pcm
    return np.clip(pcm * gain, -1.0, 1.0)

def pitch_flatten(pcm, flattening, sample_rate):
    """STFT-based F0 variance attenuation.
    Pure tensor math: smooth magnitude spectrogram toward time-mean.
    """
    if flattening <= 0.0:
        return pcm
    import torch
    
    t0 = time.perf_counter()
    pcm_t = torch.from_numpy(pcm).float()
    
    n_fft = min(2048, max(512, len(pcm_t) // 8))
    hop_length = n_fft // 4
    window = torch.hann_window(n_fft)
    
    stft = torch.stft(pcm_t.unsqueeze(0), n_fft, hop_length,
                       window=window, return_complex=True)
    
    mag = stft.abs()
    phase = stft.angle()
    
    mag_mean = mag.mean(dim=-1, keepdim=True)
    mag_new = mag_mean + (1.0 - min(flattening, 0.99)) * (mag - mag_mean)
    
    stft_new = mag_new * torch.exp(1j * phase)
    pcm_out = torch.istft(stft_new, n_fft, hop_length,
                           window=window, length=len(pcm_t))
    
    elapsed = (time.perf_counter() - t0) * 1000
    print(f"[AOS] pitch_flatten took {elapsed:.1f}ms")
    return pcm_out.squeeze(0).numpy()