//
//  audio.cpp
//  MNN
//
//  Created by MNN on 2024/11/15.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "audio/audio.hpp"
#include <MNN/expr/MathOp.hpp>
#include <MNN/expr/NeuralNetWorkOp.hpp>
#include "MNN_generated.h"
#include <cmath>
#include <algorithm>
#include <complex>
#include <fstream>
#include <iostream>
#include <limits>
#ifndef M_PI
#define M_PI 3.141592654
#endif
#ifdef _MSC_VER
#define NOMINMAX
#include <intrin.h>
#include <windows.h>
#endif


namespace MNN {
namespace AUDIO {
#ifdef _MSC_VER
inline uint32_t mnn_clz( uint32_t value ) {
    DWORD leading_zero = 0;
    if (_BitScanReverse(&leading_zero, value)) {
        return 31 - leading_zero;
    }else {
         // Same remarks as above
         return 32;
    }
}
#else
inline uint32_t mnn_clz( uint32_t value ) {
    return __builtin_clz(value);
}
#endif
struct WaveHeader {
    void SeekToDataChunk(std::istream &is) {
        //                              a t a d
        while (is && subchunk2_id != 0x61746164) {
            is.seekg(subchunk2_size, std::istream::cur);
            is.read(reinterpret_cast<char *>(&subchunk2_id), sizeof(int32_t));
            is.read(reinterpret_cast<char *>(&subchunk2_size), sizeof(int32_t));
        }
    }
    int32_t chunk_id = 0x46464952; // "RIFF"
    int32_t chunk_size;
    int32_t format         = 0x45564157; // "WAVE"
    int32_t subchunk1_id   = 0x20746d66; // "fmt "
    int32_t subchunk1_size = 16;         // PCM
    int16_t audio_format   = 1;          // PCM = 1
    int16_t num_channels   = 1;          // Mono
    int32_t sample_rate;
    int32_t byte_rate;
    int16_t block_align;
    int16_t bits_per_sample = 16;
    int32_t subchunk2_id    = 0x61746164; // "data"
    int32_t subchunk2_size;
};

std::pair<VARP, int> load(const std::string &filename, int sr, int frame_offset, int num_frames) {
    std::ifstream is(filename, std::ifstream::binary);
    auto ret = std::make_pair<VARP, int>(nullptr, 0);
    if (!is) {
        MNN_ERROR("Failed to open file: %s\n", filename.c_str());
        return ret;
    }
    WaveHeader header{};
    is.read(reinterpret_cast<char *>(&header.chunk_id), sizeof(header.chunk_id));
    if (header.chunk_id != 0x46464952) { // "RIFF"
        MNN_ERROR("Expected chunk_id RIFF. Given: 0x%08x\n", header.chunk_id);
        return ret;
    }

    is.read(reinterpret_cast<char *>(&header.chunk_size), sizeof(header.chunk_size));
    is.read(reinterpret_cast<char *>(&header.format), sizeof(header.format));
    if (header.format != 0x45564157) { // "WAVE"
        MNN_ERROR("Expected format WAVE. Given: 0x%08x\n", header.format);
        return ret;
    }

    is.read(reinterpret_cast<char *>(&header.subchunk1_id), sizeof(header.subchunk1_id));
    is.read(reinterpret_cast<char *>(&header.subchunk1_size), sizeof(header.subchunk1_size));

    if (header.subchunk1_id == 0x4b4e554a) { // "JUNK"
        is.seekg(header.subchunk1_size, std::istream::cur);
        is.read(reinterpret_cast<char *>(&header.subchunk1_id), sizeof(header.subchunk1_id));
        is.read(reinterpret_cast<char *>(&header.subchunk1_size), sizeof(header.subchunk1_size));
    }

    if (header.subchunk1_id != 0x20746d66) { // "fmt "
        MNN_ERROR("Expected subchunk1_id 'fmt '. Given: 0x%08x\n", header.subchunk1_id);
        return ret;
    }

    if (header.subchunk1_size != 16 && header.subchunk1_size != 18) {
        MNN_ERROR("Expected subchunk1_size 16 or 18. Given: %d\n", header.subchunk1_size);
        return ret;
    }

    is.read(reinterpret_cast<char *>(&header.audio_format), sizeof(header.audio_format));
    if (header.audio_format != 1 && header.audio_format != 3) {
        MNN_ERROR("Unsupported audio_format: %d. Only PCM(1) and IEEE Float(3) supported.\n", header.audio_format);
        return ret;
    }

    is.read(reinterpret_cast<char *>(&header.num_channels), sizeof(header.num_channels));
    if (header.num_channels != 1) {
        MNN_ERROR("Warning: %d channels found. Only the first channel will be used.\n", header.num_channels);
    }

    is.read(reinterpret_cast<char *>(&header.sample_rate), sizeof(header.sample_rate));
    is.read(reinterpret_cast<char *>(&header.byte_rate), sizeof(header.byte_rate));
    is.read(reinterpret_cast<char *>(&header.block_align), sizeof(header.block_align));
    is.read(reinterpret_cast<char *>(&header.bits_per_sample), sizeof(header.bits_per_sample));

    if (header.byte_rate != (header.sample_rate * header.num_channels * header.bits_per_sample / 8)) {
        MNN_ERROR("Incorrect byte rate: %d. Expected: %d\n", header.byte_rate,
                  header.sample_rate * header.num_channels * header.bits_per_sample / 8);
        return ret;
    }

    if (header.block_align != (header.num_channels * header.bits_per_sample / 8)) {
        MNN_ERROR("Incorrect block align: %d. Expected: %d\n", header.block_align,
                  header.num_channels * header.bits_per_sample / 8);
        return ret;
    }

    if (header.bits_per_sample != 8 && header.bits_per_sample != 16 && header.bits_per_sample != 32) {
        MNN_ERROR("Unsupported bits_per_sample: %d. Only 8, 16, or 32 bits per sample supported.\n",
                  header.bits_per_sample);
        return ret;
    }

    if (header.subchunk1_size == 18) {
        int16_t extra_size;
        is.read(reinterpret_cast<char *>(&extra_size), sizeof(int16_t));
        if (extra_size != 0) {
            MNN_ERROR("Unexpected extra size: %d. Expected 0.\n", extra_size);
            return ret;
        }
    }

    is.read(reinterpret_cast<char *>(&header.subchunk2_id), sizeof(header.subchunk2_id));
    is.read(reinterpret_cast<char *>(&header.subchunk2_size), sizeof(header.subchunk2_size));
    header.SeekToDataChunk(is);

    if (!is) {
        MNN_ERROR("Could not locate data chunk.\n");
        return ret;
    }

    int total_frames = header.subchunk2_size / header.block_align;
    if (frame_offset < 0 || frame_offset >= total_frames) {
        MNN_ERROR("Frame offset out of range.\n");
        return ret;
    }

    if (num_frames <= 0 || frame_offset + num_frames > total_frames) {
        num_frames = total_frames - frame_offset;
    }

    is.seekg(frame_offset * header.block_align, std::istream::cur);

    ret.first    = _Input({num_frames}, NHWC);
    ret.second   = header.sample_rate;
    auto ans_ptr = ret.first->writeMap<float>();
    if (header.bits_per_sample == 16 && header.audio_format == 1) {
        std::vector<int16_t> samples(num_frames * header.num_channels);
        is.read(reinterpret_cast<char *>(samples.data()), num_frames * header.block_align);
        if (!is) {
            MNN_ERROR("Failed to read audio data.\n");
            return ret;
        }
        for (int i = 0; i < num_frames; ++i) {
            ans_ptr[i] = samples[i * header.num_channels] / 32768.f;
        }
    } else if (header.bits_per_sample == 8 && header.audio_format == 1) {
        std::vector<uint8_t> samples(num_frames * header.num_channels);
        is.read(reinterpret_cast<char *>(samples.data()), num_frames * header.block_align);
        if (!is) {
            MNN_ERROR("Failed to read audio data.\n");
            return ret;
        }
        for (int i = 0; i < num_frames; ++i) {
            ans_ptr[i] = static_cast<float>(samples[i * header.num_channels]) / 128.f - 1.f;
        }
    } else if (header.bits_per_sample == 32 && header.audio_format == 1) {
        std::vector<int32_t> samples(num_frames * header.num_channels);
        is.read(reinterpret_cast<char *>(samples.data()), num_frames * header.block_align);
        if (!is) {
            MNN_ERROR("Failed to read audio data.\n");
            return ret;
        }
        for (int i = 0; i < num_frames; ++i) {
            ans_ptr[i] = static_cast<float>(samples[i * header.num_channels]) / static_cast<float>(INT32_MAX);
        }
    } else if (header.bits_per_sample == 32 && header.audio_format == 3) {
        std::vector<float> samples(num_frames * header.num_channels);
        is.read(reinterpret_cast<char *>(samples.data()), num_frames * header.block_align);
        if (!is) {
            MNN_ERROR("Failed to read audio data.\n");
            return ret;
        }
        for (int i = 0; i < num_frames; ++i) {
            ans_ptr[i] = samples[i * header.num_channels];
        }
    } else {
        MNN_ERROR("Unsupported bits per sample: %d or audio format: %d.\n", header.bits_per_sample,
                  header.audio_format);
        return ret;
    }

    if (sr > 0 && sr != ret.second) {
        // resample
        float resample_ratio    = static_cast<float>(sr) / header.sample_rate;
        int resample_num_frames = static_cast<int>(num_frames * resample_ratio);
        auto resampled_data     = _Input({resample_num_frames}, NHWC);
        auto src                = ret.first->readMap<float>();
        auto dst                = resampled_data->writeMap<float>();
        for (int i = 0; i < resample_num_frames; ++i) {
            float interp_index = i / resample_ratio;
            int low_index      = static_cast<int>(interp_index);
            int high_index     = std::min(low_index + 1, num_frames - 1);
            float frac         = interp_index - low_index;
            dst[i]             = (1 - frac) * src[low_index] + frac * src[high_index];
        }
        ret.first  = resampled_data;
        ret.second = sr;
    }
    return ret;
}

bool save(const std::string &filename, VARP audio, int sample_rate) {
    std::ofstream os(filename, std::ios::binary);
    if (!os) {
        MNN_ERROR("Failed to open file for writing: %s\n", filename.c_str());
        return false;
    }

    auto audio_size = audio->getInfo()->size;
    auto audio_ptr  = audio->readMap<float>();
    WaveHeader header;
    header.num_channels   = 1;
    header.sample_rate    = sample_rate;
    header.byte_rate      = sample_rate * header.num_channels * (header.bits_per_sample / 8);
    header.block_align    = header.num_channels * (header.bits_per_sample / 8);
    header.subchunk2_size = audio_size * (header.bits_per_sample / 8);
    header.chunk_size     = 36 + header.subchunk2_size;

    os.write(reinterpret_cast<const char *>(&header), sizeof(WaveHeader));

    // Convert float samples to int16 and write to file
    for (int i = 0; i < audio_size; i++) {
        float sample       = audio_ptr[i];
        int16_t int_sample = static_cast<int16_t>(std::max(-1.0f, std::min(1.0f, sample)) * 32767);
        os.write(reinterpret_cast<const char *>(&int_sample), sizeof(int16_t));
    }

    if (!os) {
        MNN_ERROR("Failed to write audio data to file.\n");
        return false;
    }

    os.close();
    return true;
}

template <typename T>
static inline VARP _var(std::vector<T> vec, const std::vector<int> &dims) {
    return _Const(vec.data(), dims, NHWC, halide_type_of<T>());
}

unsigned int next_power_of_2(unsigned int x) {
    if (x == 0)
        return 1;
    if ((x & (x - 1)) == 0)
        return x;
    return 1U << (32 - mnn_clz(x));
}

VARP hamming_window(int n_fft, bool periodic, float alpha, float beta) {
    auto window     = _Input({n_fft}, NHWC);
    auto window_ptr = window->writeMap<float>();
    int N           = periodic ? n_fft : n_fft - 1;
    for (int n = 0; n < n_fft; ++n) {
        window_ptr[n] = alpha - beta * std::cos(2.0 * M_PI * n / N);
    }
    return window;
}

VARP hann_window(int n_fft, bool periodic) {
    auto window     = _Input({n_fft}, NHWC);
    auto window_ptr = window->writeMap<float>();
    int N           = periodic ? n_fft : n_fft - 1;
    for (int n = 0; n < n_fft; ++n) {
        window_ptr[n] = 0.5 * (1 - std::cos(2 * M_PI * n / N));
    }
    return window;
}

float hz_to_mel(float freq, bool htk) {
    if (htk) {
        return 2595 * std::log10(1 + freq / 700);
    } else {
        constexpr float f_min = 0.0, f_sp = 200.0 / 3.0, min_log_hz = 1000.0;
        constexpr float logstep     = 0.06875177742094912;
        constexpr float min_log_mel = (min_log_hz - f_min) / f_sp;
        float mels                  = (freq - f_min) / f_sp;
        if (freq >= min_log_hz) {
            mels = min_log_mel + std::log(freq / min_log_hz) / logstep;
        }
        return mels;
    }
}

float mel_to_hz(float mel, bool htk) {
    if (htk) {
        return 700 * (std::pow(10, mel / 2595.0) - 1);
    } else {
        constexpr float f_min = 0.0f, f_sp = 200.0f / 3, min_log_hz = 1000.0f;
        constexpr float logstep     = 0.06875177742094912;
        constexpr float min_log_mel = (min_log_hz - f_min) / f_sp;
        float freq                  = f_min + f_sp * mel;
        if (mel >= min_log_mel) {
            freq = min_log_hz * std::exp(logstep * (mel - min_log_mel));
        }
        return freq;
    }
}

VARP melscale_fbanks(const MelscaleParams *params) {
    int n_mels = 128, n_fft = 400, sample_rate = 16000;
    bool htk = true, norm = false;
    float f_min = 0.0, f_max = 0.0;
    if (params != nullptr) {
        n_mels      = params->n_mels;
        n_fft       = params->n_fft;
        sample_rate = params->sample_rate;
        htk         = params->htk;
        norm        = params->norm;
        f_min       = params->f_min;
        f_max       = params->f_max;
    }
    int n_freqs   = n_fft / 2 + 1;
    float nyquist = 0.5 * sample_rate;
    std::vector<float> all_freqs(n_freqs);
    for (int i = 0; i < n_freqs; ++i) {
        all_freqs[i] = i * nyquist / (n_freqs - 1);
    }
    f_max         = f_max <= 0.0 ? nyquist : f_max;
    float m_min   = hz_to_mel(f_min, htk);
    float m_max   = hz_to_mel(f_max, htk);
    float m_delta = (m_max - m_min) / (n_mels + 1);

    auto bins     = _Input({n_mels, n_freqs}, NHWC);
    auto bins_ptr = bins->writeMap<float>();
    for (int n = 0; n < n_mels; ++n) {
        float left  = mel_to_hz(m_min + m_delta * (n + 0), htk);
        float curr  = mel_to_hz(m_min + m_delta * (n + 1), htk);
        float right = mel_to_hz(m_min + m_delta * (n + 2), htk);
        float enorm = (htk && norm) ? 1.0 : 2.0 / (right - left);
        for (int k = 0; k < n_freqs; ++k) {
            float val = 0.f, f_k = all_freqs[k];
            if (f_k >= left && f_k <= curr) {
                val = (f_k - left) / (curr - left);
            } else if (f_k > curr && f_k <= right) {
                val = (right - f_k) / (right - curr);
            }
            bins_ptr[n * n_freqs + k] = val * enorm;
        }
    }
    return bins;
}

VARP spectrogram(VARP waveform, const SpectrogramParams *params) {
    int pad_left = 0, pad_right = 0, pad_mode = REFLECT;
    int n_fft = 400, hop_length = 0, win_length = 0, window_type = HANNING;
    bool center = false, normalized = false;
    float power = 2.0;
    if (params) {
        pad_left    = params->pad_left;
        pad_right   = params->pad_right;
        center      = params->center;
        pad_mode    = params->pad_mode;
        n_fft       = params->n_fft;
        hop_length  = params->hop_length;
        win_length  = params->win_length;
        window_type = params->window_type;
        normalized  = params->normalized;
        power       = params->power;
    }
    if (pad_left > 1 || pad_right > 1) {
        waveform = MNN::Express::_Pad(waveform, _var<int>({pad_left, pad_right}, {2}), MNN::Express::CONSTANT);
    }
    if (center) {
        waveform = MNN::Express::_Pad(waveform, _var<int>({n_fft / 2, n_fft / 2}, {2}), static_cast<MNN::Express::PadValueMode>(pad_mode));
    }
    waveform = _Reshape(waveform, {1, -1, 1});
    hop_length = hop_length ? hop_length : n_fft / 2;
    win_length = win_length ? win_length : n_fft;
    VARP window;
    switch (window_type) {
        case HANNING:
            window = hann_window(win_length);
            break;
        case HAMMING:
            window = hamming_window(win_length);
            break;
        default:
            window = hann_window(win_length);
            break;
    }
    std::unique_ptr<OpT> op(new OpT);
    op->type      = OpType_Stft;
    op->main.type = OpParameter_StftParam;
    auto param = new StftParamT;
    param->abs = true;
    op->main.value = param;
    EXPRP stftexpr = Expr::create(std::move(op), {waveform, _Scalar<int>(hop_length), window});
    int frame_size = win_length > 0 ? win_length : n_fft;
    int nstfts = ((waveform->getInfo()->dim[1] - frame_size) / hop_length) + 1;
    int dft_unique_bins = n_fft / 2 + 1;
    auto specgram = MNN::Express::Variable::create(stftexpr);
    specgram = _Square(specgram);
    auto startsDims = std::vector<int>{0, 0, 0, 0};
    auto starts1Dims = std::vector<int>{0, 0, 0, 1};
    auto sizeDims = std::vector<int>{1, nstfts, dft_unique_bins, 1};
    auto startVar = _Const(startsDims.data(), {4}, NCHW, halide_type_of<int>());
    auto start1Var = _Const(starts1Dims.data(), {4}, NCHW, halide_type_of<int>());
    auto sizeVar = _Const(sizeDims.data(), {4}, NCHW, halide_type_of<int>());
    auto specgramReal = _Slice(specgram, startVar, sizeVar);
    auto specgramVirt = _Slice(specgram, start1Var, sizeVar);
    specgram = specgramReal + specgramVirt;
    specgram = _Reshape(specgram, {nstfts, dft_unique_bins});
    if (normalized) {
        float window_norm = 1.0f / _ReduceSum(_Square(window))->readMap<float>()[0];
        specgram = specgram * _Scalar<float>(window_norm);
    }
    if (power == 1.0f) {
        specgram = _Sqrt(specgram);
    } else if (power != 2.0f) {
        specgram = _Pow(specgram, _Scalar<float>(power / 2.0f));
    }
    return specgram;
}

VARP mel_spectrogram(VARP waveform, const MelscaleParams *mel_params, const SpectrogramParams *spec_params) {
    auto banks        = melscale_fbanks(mel_params);
    auto specgram     = spectrogram(waveform, spec_params);
    auto mel_specgram = _MatMul(specgram, banks, false, true);
    return mel_specgram;
}

VARP fbank(VARP waveform, int sampling_rate, int n_mels, int n_fft, int hop_length, float dither, float preemphasis) {
    int wav_len      = waveform->getInfo()->size;
    int frame_num    = (wav_len - n_fft) / hop_length + 1;
    if (frame_num <= 0 || wav_len < n_fft) {
        return nullptr; // frame_num is zero
    }
    // get_strided: sizes: [m, n_fft], strides: [windows_shift, 1]
    int m                           = 1 + (wav_len - n_fft) / hop_length;
    std::vector<int> strided_region = {
        0, // src offset
        wav_len,
        hop_length,
        1, // src strides
        0, // dst offset
        m * n_fft,
        n_fft,
        1, // dst strides
        1,
        m,
        n_fft // dst sizes
    };
    auto strided_wav = _Raster({waveform}, strided_region, {m, n_fft});
    auto wav_dim     = strided_wav->getInfo()->dim;
    // add_dither
    if (dither > 0.f) {
        auto rand_dither = _RandomUnifom(_var<int>(wav_dim, {static_cast<int>(wav_dim.size())}),
                                         halide_type_of<float>(), -dither, dither);
        strided_wav      = strided_wav + rand_dither;
    }
    // subtract each row/frame by its mean
    {
        auto row_means   = _ReduceMean(strided_wav, {-1}, true);
        strided_wav      = strided_wav - row_means;
    }
    if (preemphasis != 0.f) {
        std::vector<int> offset_region          = {
            // region 0
            0,                               // src offset
            m * n_fft, n_fft, 1, // src strides
            0,                               // dst offset
            m * n_fft, n_fft, 1, // dst strides
            1, m, 1,                         // dst sizes
            // region 1
            0,                               // src offset
            m * n_fft, n_fft, 1, // src strides
            1,                               // dst offset
            m * n_fft, n_fft, 1, // dst strides
            1, m, n_fft - 1            // dst sizes
        };
        auto offset_strided_wav = _Raster({strided_wav, strided_wav}, offset_region, {m, n_fft});
        strided_wav             = strided_wav - _Scalar<float>(preemphasis) * offset_strided_wav;
    }
    int padded_n_fft = next_power_of_2(n_fft);
    MelscaleParams mel_params;
    mel_params.n_mels      = n_mels;
    mel_params.n_fft       = padded_n_fft;
    mel_params.sample_rate = sampling_rate;
    mel_params.f_min       = 20.0;
    SpectrogramParams spec_params;
    spec_params.n_fft      = padded_n_fft;
    spec_params.hop_length = n_fft;
    auto mel_energies      = mel_spectrogram(strided_wav, &mel_params, &spec_params);
    mel_energies           = _Log(mel_energies);
    return mel_energies;
}

VARP conformer_fbank(VARP waveform, int sample_rate, int n_mels, int n_fft, int hop_length, int win_length,
                     float preemphasis) {
    if (waveform == nullptr || waveform->getInfo() == nullptr) {
        return nullptr;
    }
    int wav_len = waveform->getInfo()->size;
    if (wav_len <= 0) {
        return nullptr;
    }
    // 1. Preemphasis: x_new[0] = x[0], x_new[n] = x[n] - preemph * x[n-1]
    if (preemphasis != 0.f) {
        float zero_val = 0.f;
        auto zero = _Const(&zero_val, {1}, NHWC, halide_type_of<float>());
        auto x_prev = _Slice(waveform, _var<int>({0}, {1}), _var<int>({wav_len - 1}, {1}));
        auto shifted = _Concat({zero, x_prev}, 0);
        waveform = waveform - _Scalar<float>(preemphasis) * shifted;
    }
    // 2. Center padding (constant/zero, matching NeMo's pad_mode="constant")
    waveform = MNN::Express::_Pad(waveform, _var<int>({n_fft / 2, n_fft / 2}, {2}), MNN::Express::CONSTANT);
    waveform = _Reshape(waveform, {1, -1, 1});
    // 3. Create zero-padded window: hann(win_length) padded to n_fft
    if (win_length <= 0) win_length = n_fft;
    auto window = hann_window(win_length);
    if (win_length < n_fft) {
        int pad_left_w = (n_fft - win_length) / 2;
        int pad_right_w = n_fft - win_length - pad_left_w;
        window = MNN::Express::_Pad(window, _var<int>({pad_left_w, pad_right_w}, {2}), MNN::Express::CONSTANT);
    }
    // 4. STFT
    {
        std::unique_ptr<OpT> op(new OpT);
        op->type      = OpType_Stft;
        op->main.type = OpParameter_StftParam;
        auto param = new StftParamT;
        param->abs = true;
        op->main.value = param;
        EXPRP stftexpr = Expr::create(std::move(op), {waveform, _Scalar<int>(hop_length), window});
        waveform = MNN::Express::Variable::create(stftexpr);
    }
    int padded_len = wav_len + n_fft;
    int nstfts = ((padded_len - n_fft) / hop_length) + 1;
    int dft_unique_bins = n_fft / 2 + 1;
    // Power spectrum
    auto specgram = _Square(waveform);
    auto startsDims = std::vector<int>{0, 0, 0, 0};
    auto starts1Dims = std::vector<int>{0, 0, 0, 1};
    auto sizeDims = std::vector<int>{1, nstfts, dft_unique_bins, 1};
    auto startVar = _Const(startsDims.data(), {4}, NCHW, halide_type_of<int>());
    auto start1Var = _Const(starts1Dims.data(), {4}, NCHW, halide_type_of<int>());
    auto sizeVar = _Const(sizeDims.data(), {4}, NCHW, halide_type_of<int>());
    auto specgramReal = _Slice(specgram, startVar, sizeVar);
    auto specgramVirt = _Slice(specgram, start1Var, sizeVar);
    specgram = specgramReal + specgramVirt;
    specgram = _Reshape(specgram, {nstfts, dft_unique_bins}); // [T, n_fft/2+1]
    // 5. Mel filterbank (slaney norm)
    MelscaleParams mel_params;
    mel_params.n_mels      = n_mels;
    mel_params.n_fft       = n_fft;
    mel_params.sample_rate = sample_rate;
    mel_params.htk         = false;
    mel_params.norm        = true;
    auto banks = melscale_fbanks(&mel_params);
    auto mel_specgram = _MatMul(specgram, banks, false, true); // [T, n_mels]
    // 6. Log with guard value
    float log_guard = 5.96046e-8f; // 2^-24
    mel_specgram = _Log(mel_specgram + _Scalar<float>(log_guard));
    // 7. Per-feature normalization using valid frame count
    //    Valid frames = floor((wav_len + pad_amount - n_fft) / hop_length)
    //    where pad_amount = n_fft (from center padding n_fft/2 on each side)
    int total_frames = mel_specgram->getInfo()->dim[0];
    int valid_frames = wav_len / hop_length; // floor division
    if (valid_frames <= 1 || total_frames <= 1) {
        return nullptr;
    }
    // Mask frames beyond valid_frames to zero
    if (valid_frames < total_frames) {
        // Zero out invalid frames by multiplying with a mask
        auto frame_idx = _Cast<float>(_Range(_Scalar<int>(0), _Scalar<int>(total_frames), _Scalar<int>(1)));
        auto valid_mask = _Cast<float>(_Less(frame_idx, _Scalar<float>((float)valid_frames)));
        valid_mask = _Unsqueeze(valid_mask, {1}); // [T, 1] for broadcasting
        mel_specgram = mel_specgram * valid_mask;
    }
    // Compute mean and std using only valid frames
    auto frame_sum = _ReduceSum(mel_specgram, {0}, true); // [1, n_mels]
    auto mean = frame_sum * _Scalar<float>(1.0f / valid_frames);
    auto diff = mel_specgram - mean;
    // Zero out invalid frames in diff too
    if (valid_frames < total_frames) {
        auto frame_idx2 = _Cast<float>(_Range(_Scalar<int>(0), _Scalar<int>(total_frames), _Scalar<int>(1)));
        auto valid_mask2 = _Cast<float>(_Less(frame_idx2, _Scalar<float>((float)valid_frames)));
        valid_mask2 = _Unsqueeze(valid_mask2, {1});
        diff = diff * valid_mask2;
    }
    auto diff_sq_sum = _ReduceSum(diff * diff, {0}, true);
    // Bessel's correction: std = sqrt(sum / (N-1) + eps)
    auto std_val = _Sqrt(diff_sq_sum * _Scalar<float>(1.0f / (valid_frames - 1)) + _Scalar<float>(1e-5f));
    mel_specgram = diff / std_val;
    // 8. Truncate to valid frames: floor(wav_len / hop_length)
    //    (NeMo preprocessor returns this as seq_len, extra frames are padding)
    if (valid_frames < total_frames) {
        mel_specgram = _Slice(mel_specgram, _var<int>({0, 0}, {2}), _var<int>({valid_frames, n_mels}, {2}));
    }
    // 9. Transpose to [1, n_mels, T]
    mel_specgram = _Unsqueeze(mel_specgram, {0, 1});
    mel_specgram = _Convert(mel_specgram, NCHW);
    mel_specgram = _Squeeze(mel_specgram, {2});
    return mel_specgram;
}

VARP whisper_fbank(VARP waveform, int sample_rate, int n_mels, int n_fft, int hop_length, int chunk_len) {
    int n_samples = chunk_len * sample_rate;
    int pad_right = n_samples - waveform->getInfo()->size;
    pad_right     = pad_right > 0 ? pad_right : 0;
    MelscaleParams mel_params;
    mel_params.n_mels      = n_mels;
    mel_params.n_fft       = n_fft;
    mel_params.sample_rate = sample_rate;
    mel_params.htk         = false;
    mel_params.norm        = true;
    SpectrogramParams spec_params;
    spec_params.pad_right  = pad_right;
    spec_params.n_fft      = n_fft;
    spec_params.hop_length = hop_length;
    spec_params.center     = true;
    auto mel_specgram      = mel_spectrogram(waveform, &mel_params, &spec_params);
    mel_specgram =
        _Slice(mel_specgram, _var<int>({0, 0}, {2}), _var<int>({mel_specgram->getInfo()->dim[0] - 1, -1}, {2}));
    auto log_specgram = _Log(_Maximum(mel_specgram, _Scalar<float>(1e-10f))) / _Log(_Scalar<float>(10.0));
    log_specgram      = _Maximum(log_specgram, _ReduceMax(log_specgram) - _Scalar<float>(8.0));
    log_specgram      = (log_specgram + _Scalar<float>(4.0)) / _Scalar<float>(4.0);
    // NHWC -> NCHW
    log_specgram = _Unsqueeze(log_specgram, {0, 1});
    log_specgram = _Convert(log_specgram, NCHW);
    log_specgram = _Squeeze(log_specgram, {2});
    return log_specgram;
}

VARP usm_fbank(VARP waveform, int sample_rate, int n_mels, int n_fft, int hop_length, int frame_length, float mel_floor) {
    // USM-style mel spectrogram (Gemma4 audio encoder)
    // Key differences from whisper: magnitude spectrum, semicausal padding, log(mel + floor)
    // 1. Semicausal padding
    int pad_left = frame_length / 2;
    MelscaleParams mel_params;
    mel_params.n_mels      = n_mels;
    mel_params.n_fft       = n_fft;
    mel_params.sample_rate = sample_rate;
    mel_params.htk         = true;
    mel_params.norm        = true;  // htk+norm → enorm=1.0 (no area normalization), matching HF mel_filter_bank(norm=None)
    // 2. Spectrogram with zero-padded window (frame_length -> n_fft for FFT)
    // The STFT uses window length as frame size, so we zero-pad the window to n_fft
    SpectrogramParams spec_params;
    spec_params.pad_left   = pad_left;
    spec_params.n_fft      = n_fft;
    spec_params.hop_length = hop_length;
    spec_params.win_length = n_fft;  // Use n_fft as window size (zero-padded)
    spec_params.power      = 1.0f;   // magnitude spectrum
    // Override default window: create hann(frame_length) zero-padded to n_fft
    auto banks = melscale_fbanks(&mel_params);
    // Pad waveform
    waveform = MNN::Express::_Pad(waveform, _var<int>({pad_left, 0}, {2}), MNN::Express::CONSTANT);
    waveform = _Reshape(waveform, {1, -1, 1});
    // Create periodic hann window (matching HF window_function('hann', periodic=True)), zero-padded to n_fft
    auto hann = hann_window(frame_length, true);
    auto padded_window = MNN::Express::_Pad(hann, _var<int>({0, n_fft - frame_length}, {2}), MNN::Express::CONSTANT);
    // STFT
    std::unique_ptr<OpT> op(new OpT);
    op->type      = OpType_Stft;
    op->main.type = OpParameter_StftParam;
    auto param = new StftParamT;
    param->abs = true;
    op->main.value = param;
    EXPRP stftexpr = Expr::create(std::move(op), {waveform, _Scalar<int>(hop_length), padded_window});
    int nstfts = ((waveform->getInfo()->dim[1] - n_fft) / hop_length) + 1;
    int dft_unique_bins = n_fft / 2 + 1;
    auto specgram = MNN::Express::Variable::create(stftexpr);
    // STFT output is complex: [1, nstfts, dft_unique_bins, 2], compute magnitude
    specgram = _Square(specgram);
    auto startsDims = std::vector<int>{0, 0, 0, 0};
    auto starts1Dims = std::vector<int>{0, 0, 0, 1};
    auto sizeDims = std::vector<int>{1, nstfts, dft_unique_bins, 1};
    auto startVar = _Const(startsDims.data(), {4}, NCHW, halide_type_of<int>());
    auto start1Var = _Const(starts1Dims.data(), {4}, NCHW, halide_type_of<int>());
    auto sizeVar = _Const(sizeDims.data(), {4}, NCHW, halide_type_of<int>());
    auto specgramReal = _Slice(specgram, startVar, sizeVar);
    auto specgramVirt = _Slice(specgram, start1Var, sizeVar);
    specgram = _Sqrt(specgramReal + specgramVirt);  // magnitude = sqrt(real^2 + imag^2)
    specgram = _Reshape(specgram, {nstfts, dft_unique_bins});
    // 3. Mel filterbank
    auto mel_specgram = _MatMul(specgram, banks, false, true);
    // 4. log(mel + mel_floor)
    auto log_specgram = _Log(mel_specgram + _Scalar<float>(mel_floor));
    // Output: [T, n_mels] -> [1, T, n_mels]
    log_specgram = _Unsqueeze(log_specgram, {0});
    return log_specgram;
}

} // namespace AUDIO
} // namespace MNN