soothe2-re/dsp/spectral.cpp

#include "spectral.hpp"
#include "fftconv.hpp"
#include "log2_ln.hpp"
#include "exp2_tables.hpp"
#include "exp2.hpp"
#include <cmath>
#include <cstring>
#include <vector>
#include <cstdlib>
#include <cstdio>

SpectralProcessor::SpectralProcessor(size_t nfft, size_t hop, float sample_rate)
    : nfft_(nfft), hop_(hop), frame_count_(0), output_pos_(0),
      detector_(nfft, sample_rate) {
    window_ = new double[nfft_];
    computeWindow();
    fft::init_plan(&plan_, static_cast<uint32_t>(std::log2(nfft_)));
    buf_ = new std::complex<double>[nfft_];
    tmp_buf_ = new std::complex<double>[nfft_];
    fir_buf_ = new std::complex<double>[nfft_];
    fir_freq_ = new std::complex<double>[nfft_];
    overlap_.resize(nfft_, 0.0f);
    mask_.resize(nfft_, 1.0f);

    // Build FIR window: falling half of periodic Hann(4096).
    // Plugin reads window[N/2..N-1] of periodic Hann (rising 0→1).
    fir_window_.resize(nfft_);
    for (size_t i = 0; i < nfft_; i++) {
        fir_window_[i] = 0.5 * (1.0 - std::cos(2.0 * M_PI * i / nfft_));
    }
}

SpectralProcessor::~SpectralProcessor() {
    delete[] window_;
    delete[] buf_;
    delete[] tmp_buf_;
    delete[] fir_buf_;
    delete[] fir_freq_;
}

void SpectralProcessor::setDetectorParams(const std::vector<DetectorBand>& bands) {
    detector_.setParams(bands);
    loadWinFreq();
}

void SpectralProcessor::computeWindow() {
    // RT_WIN: 0=symmetric hann (legacy), 1=periodic hann, 2=rectangular
    static const int winmode = getenv("RT_WIN") ? atoi(getenv("RT_WIN")) : 0;
    for (size_t i = 0; i < nfft_; i++) {
        double v;
        if (winmode == 1)      v = 0.5 * (1.0 - std::cos(2.0 * M_PI * i / nfft_));
        else if (winmode == 2) v = 1.0;
        else                   v = 0.5 * (1.0 - std::cos(2.0 * M_PI * i / (nfft_ - 1)));
        window_[i] = v;
    }
}

void SpectralProcessor::stftFrame(const float* in, std::complex<double>* out) {
    for (size_t i = 0; i < nfft_; i++) {
        out[i] = std::complex<double>(static_cast<double>(in[i]) * window_[i], 0.0);
    }
    fft::execute(&plan_, out);
}

void SpectralProcessor::istftFrame(std::complex<double>* in, float* out, float* overlap) {
    memcpy(tmp_buf_, in, nfft_ * sizeof(std::complex<double>));
    fft::execute_inverse(&plan_, tmp_buf_);
    static bool wola_computed = false;
    static float wola_norm = 1.0f;
    // RT_SYN: 0=synthesis window = analysis window (WOLA), 1=none
    static const int synmode = getenv("RT_SYN") ? atoi(getenv("RT_SYN")) : 0;
    if (!wola_computed) {
        double wola_sum = 0.0;
        for (size_t i = 0; i < nfft_; i++) {
            double w = (synmode == 1) ? 1.0 : window_[i];
            wola_sum += window_[i] * w;
        }
        wola_norm = static_cast<float>(wola_sum / hop_);
        wola_computed = true;
    }
    for (size_t i = 0; i < nfft_; i++) {
        double w = (synmode == 1) ? 1.0f : window_[i];
        overlap[i] += static_cast<float>(tmp_buf_[i].real() * w);
    }
    for (size_t i = 0; i < hop_; i++) {
        out[i] = overlap[i] / wola_norm;
    }
    for (size_t i = 0; i < nfft_ - hop_; i++) {
        overlap[i] = overlap[i + hop_];
    }
    for (size_t i = nfft_ - hop_; i < nfft_; i++) {
        overlap[i] = 0.0f;
    }
}

void SpectralProcessor::loadWinFreq() {
    if (win_freq_loaded_) return;
    win_freq_loaded_ = true;
    // Try to load WIN_freq from live capture (handoff/rtwin_freq_44100.npy)
    FILE* f = fopen("handoff/rtwin_freq_44100.npy", "rb");
    if (!f) {
        // Fallback: compute periodic Hann, second half (0.5→1.0 rising)
        win_freq_.resize(nfft_ / 2 + 1);
        for (size_t i = 0; i <= nfft_ / 2; i++) {
            win_freq_[i] = static_cast<float>(0.5 * (1.0 - std::cos(2.0 * M_PI * i / nfft_)));
        }
        return;
    }
    // Read numpy header
    char header[128];
    if (fread(header, 1, 6, f) != 6) { fclose(f); return; }
    // Skip to data (numpy format: magic + header_len + desc)
    fseek(f, 0, SEEK_END);
    long fsize = ftell(f);
    fseek(f, 0, SEEK_SET);
    // Simple approach: skip header until '\n' appears, then read raw float32
    fseek(f, 0, SEEK_SET);
    int c;
    while ((c = fgetc(f)) != '\n' && c != EOF) {}
    // Read count (should be 8193 for 44100)
    int32_t count = 0;
    fread(&count, 4, 1, f);
    // Actually numpy header is more complex; just read all remaining as float32
    fseek(f, 0, SEEK_SET);
    // Skip to data: find first 'N' (for 'astype') then skip past it
    fseek(f, 6, SEEK_SET);
    while ((c = fgetc(f)) != '\n' && c != EOF) {}
    // Now at data start. Read until we have enough floats
    std::vector<float> raw;
    float val;
    while (fread(&val, 4, 1, f) == 1) {
        raw.push_back(val);
    }
    fclose(f);
    if (raw.size() > 0) {
        win_freq_ = raw;
    } else {
        // Fallback
        win_freq_.resize(nfft_ / 2 + 1);
        for (size_t i = 0; i <= nfft_ / 2; i++) {
            win_freq_[i] = static_cast<float>(0.5 * (1.0 - std::cos(2.0 * M_PI * i / nfft_)));
        }
    }
}

void SpectralProcessor::buildFirFromMask(const float* mask, std::complex<double>* fir, size_t nbin) {
    // Plugin FIR construction pipeline (52b550-52b8bb) uses real RFFTs:
    // 1. log(bands) → negate
    // 2. forward real RFFT (opB, th1a90)
    // 3. EXP in-place (140b30)
    // 4. inverse real RFFT (opC, th2180)
    // 5. Window with WIN_freq
    // 6. forward real RFFT (opD, th1a90)
    // 7. Normalize: FIR[0]=1, FIR[1]=0
    
    const size_t half = nfft_ / 2;
    const size_t nfft = nfft_;
    
    // Step 1: log(mask) and negate
    std::vector<double> log_mask(nfft);
    for (size_t i = 0; i <= half; i++) {
        float m = mask[i];
        if (m > 1e-12f) {
            float ln_m = soothe2::ln_plugin_f32(m);
            log_mask[i] = -static_cast<double>(ln_m);
        } else {
            log_mask[i] = 0.0;
        }
    }
    // Zero upper half
    for (size_t i = half + 1; i < nfft; i++) {
        log_mask[i] = 0.0;
    }
    
    // Step 2: forward real RFFT (opB)
    std::vector<std::complex<double>> H(half + 1);
    fft::execute_real_forward(&plan_, log_mask.data(), H.data());
    
    // Step 3: EXP in-place using plugin's exp2
    for (size_t i = 0; i <= half; i++) {
        double re = H[i].real();
        double im = H[i].imag();
        double exp_re = exp2d::exp2_dsp(re);
        H[i] = std::complex<double>(exp_re * std::cos(im), exp_re * std::sin(im));
    }
    
    // Step 4: inverse real RFFT (opC)
    std::vector<double> time_domain(nfft);
    fft::execute_real_inverse(&plan_, H.data(), time_domain.data());
    
    // Step 5: Window with WIN_freq (falling half of periodic Hann)
    if (!win_freq_.empty() && win_freq_.size() > half) {
        for (size_t i = 0; i < half; i++) {
            time_domain[i] *= static_cast<double>(win_freq_[half + i]);
        }
    }
    // Zero upper half
    for (size_t i = half; i < nfft; i++) {
        time_domain[i] = 0.0;
    }
    
    // Step 6: forward real RFFT (opD)
    std::vector<std::complex<double>> H_final(half + 1);
    fft::execute_real_forward(&plan_, time_domain.data(), H_final.data());
    
    // Step 7: Normalize and copy to output
    double scale = 1.0;
    if (std::abs(H_final[0].real()) > 1e-12) {
        scale = 1.0 / H_final[0].real();
    }
    
    // Copy to full complex array
    for (size_t i = 0; i <= half; i++) {
        fir[i] = H_final[i] * scale;
    }
    // Mirror for negative frequencies
    for (size_t i = 1; i < half; i++) {
        fir[nfft - i] = std::conj(fir[i]);
    }
    
    // Enforce FIR[0]=1, FIR[1]=0
    fir[0] = std::complex<double>(1.0, 0.0);
    if (half >= 1) {
        fir[1] = std::complex<double>(0.0, 0.0);
    }
}

void SpectralProcessor::processBlock(float* in, float* out, size_t num_samples, size_t num_channels) {
    memset(out, 0, num_samples * sizeof(float));
    if (num_samples == 0 || num_samples < nfft_) {
        return;
    }

    static const int firconv = []() {
        const char* e = getenv("RT_FIRCONV");
        return e ? atoi(e) : 0;
    }();

    size_t nframes = (num_samples - nfft_) / hop_ + 1;

    for (size_t f = 0; f < nframes; f++) {
        size_t offset = f * hop_;
        if (offset + nfft_ > num_samples) break;
        stftFrame(in + offset, buf_);

        detector_.processFrame(buf_, mask_.data());

        if (firconv == 3) {
            // RT_FIRCONV=3 (NOTES 24k): plugin application law decoded live:
            //   applied_gain = 1.019 * V^1.8345 per bin (rms 0.0025 dB over
            //   8 drive levels). V = band curve (detector output); here M.
            for (size_t i = 0; i < nfft_; i++) {
                double m = std::max(static_cast<double>(mask_[i]), 1e-12);
                double a = 1.019 * std::pow(m, 1.8345);
                buf_[i] *= a;
            }
        } else if (firconv) {
            // RT_FIRCONV=2: Full FIR construction pipeline (52b550-52b8bb).
            // mask → reciprocal (1/mask) → window → normalize → complex multiply.
            // This replicates the plugin's FFT-conv FIR design path.
            buildFirFromMask(mask_.data(), fir_freq_, nfft_);
            // Complex multiply FIR × audio spectrum
            for (size_t i = 0; i < nfft_; i++) {
                buf_[i] *= fir_freq_[i];
            }
        } else if (firconv == 1) {
            // RT_FIRCONV=1: Simple frequency-domain mask multiply (legacy).
            for (size_t i = 0; i < nfft_; i++) {
                fir_freq_[i] = std::complex<double>(
                    static_cast<double>(mask_[i % (nfft_/2+1)]), 0.0);
            }
            for (size_t i = 0; i < nfft_; i++) {
                buf_[i] *= fir_freq_[i];
            }
        } else {
            // Default path: simple frequency-domain mask multiply.
            for (size_t i = 0; i < nfft_; i++) {
                buf_[i] *= mask_[i];
            }
        }

        istftFrame(buf_, out + offset, overlap_.data());
    }
}