// Implementation kazeia_mel — voir kazeia_mel.h pour l'API publique.
#include "kazeia_mel.h"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <fstream>

namespace kazeia_mel {

// ============================================================================
// Configurations pré-faites
// ============================================================================
MelConfig config_qwen3_tts_speaker() {
    MelConfig c;
    c.sample_rate    = 24000;
    c.n_fft          = 1024;
    c.hop_size       = 256;
    c.win_size       = 1024;
    c.n_mels         = 128;
    c.pad_amount     = (1024 - 256) / 2;     // 384, reflect (TTS convention)
    c.fixed_n_frames = 0;
    c.clip_eps       = 1e-5f;
    c.window         = WindowKind::HANN_SYMMETRIC;
    c.comp           = Compression::LOG_NATURAL;
    c.post           = PostNorm::NONE;
    return c;
}

MelConfig config_whisper() {
    MelConfig c;
    c.sample_rate    = 16000;
    c.n_fft          = 400;
    c.hop_size       = 160;
    c.win_size       = 400;
    c.n_mels         = 80;
    c.pad_amount     = 200;                  // n_fft/2 (HF whisper)
    c.fixed_n_frames = 3000;                 // 30s padding fixe
    c.clip_eps       = 1e-10f;
    c.window         = WindowKind::HANN_PERIODIC;
    c.comp           = Compression::LOG10;
    c.post           = PostNorm::WHISPER;
    return c;
}

// ============================================================================
// Helpers internes
// ============================================================================

// FFT radix-2 in-place (n doit être puissance de 2). Si n_fft pas puissance de 2
// (cas Whisper N_FFT=400), on zero-pad jusqu'à la puissance de 2 supérieure.
static void fft_radix2(float* real, float* imag, int n) {
    int j = 0;
    for (int i = 1; i < n; i++) {
        int bit = n >> 1;
        while (j & bit) { j ^= bit; bit >>= 1; }
        j ^= bit;
        if (i < j) { std::swap(real[i], real[j]); std::swap(imag[i], imag[j]); }
    }
    for (int len = 2; len <= n; len <<= 1) {
        int half = len / 2;
        double angle = -2.0 * M_PI / len;
        float wR = (float)std::cos(angle), wI = (float)std::sin(angle);
        for (int i = 0; i < n; i += len) {
            float cR = 1.f, cI = 0.f;
            for (int k = 0; k < half; k++) {
                float tR = cR * real[i+k+half] - cI * imag[i+k+half];
                float tI = cR * imag[i+k+half] + cI * real[i+k+half];
                real[i+k+half] = real[i+k] - tR;
                imag[i+k+half] = imag[i+k] - tI;
                real[i+k]      += tR;
                imag[i+k]      += tI;
                float ncR = cR * wR - cI * wI;
                cI = cR * wI + cI * wR;
                cR = ncR;
            }
        }
    }
}

static int next_pow2(int n) {
    int p = 1; while (p < n) p <<= 1; return p;
}

static std::vector<float> make_window(int n, WindowKind kind) {
    std::vector<float> w(n);
    if (kind == WindowKind::HANN_SYMMETRIC) {
        // torch.hann_window(N) default = symmetric : 0.5 * (1 - cos(2*pi*n/(N-1)))
        for (int i = 0; i < n; i++)
            w[i] = 0.5f * (1.0f - std::cos(2.0f * (float)M_PI * (float)i / (float)(n - 1)));
    } else {
        // Periodic (HF whisper, numpy default) : 0.5 * (1 - cos(2*pi*n/N))
        for (int i = 0; i < n; i++)
            w[i] = 0.5f * (1.0f - std::cos(2.0f * (float)M_PI * (float)i / (float)n));
    }
    return w;
}

// Pad reflect : miroir interne (n'inclut PAS le sample bord). Équivalent torch reflect.
static std::vector<float> pad_reflect(const std::vector<float>& x, int padding) {
    const int n = (int)x.size();
    std::vector<float> out((size_t)n + 2 * padding);
    for (int i = 0; i < padding; i++)         out[i]               = x[padding - i];
    for (int i = 0; i < n; i++)               out[padding + i]     = x[i];
    for (int i = 0; i < padding; i++)         out[padding + n + i] = x[n - 2 - i];
    return out;
}

// ============================================================================
// load_mel_basis : fichier binaire f32 [n_mels, n_fft/2+1]
// ============================================================================
bool load_mel_basis(const char * path, const MelConfig & cfg, std::vector<float> & out) {
    const int fft_bins = cfg.n_fft / 2 + 1;
    const size_t expected = (size_t)cfg.n_mels * fft_bins;
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f) { fprintf(stderr, "kazeia_mel::load_mel_basis: open %s FAIL\n", path); return false; }
    size_t n = (size_t)f.tellg() / sizeof(float);
    if (n != expected) {
        fprintf(stderr, "kazeia_mel::load_mel_basis: %s contient %zu f32, attendu %zu (=%d mels x %d bins)\n",
                path, n, expected, cfg.n_mels, fft_bins);
        return false;
    }
    f.seekg(0);
    out.assign(expected, 0.0f);
    f.read((char*)out.data(), expected * sizeof(float));
    return (bool)f;
}

// ============================================================================
// compute : STFT (radix-2 zero-pad) + projection mel + log + post-norm
// ============================================================================
std::vector<float> compute(const std::vector<float> & wav_in,
                            const std::vector<float> & mel_basis,
                            const MelConfig & cfg,
                            int & out_T) {
    const int fft_bins = cfg.n_fft / 2 + 1;
    const int pad_amount = (cfg.pad_amount >= 0) ? cfg.pad_amount : (cfg.n_fft - cfg.hop_size) / 2;

    // Optionnel : padding audio à T_fixe (Whisper 30s)
    std::vector<float> wav = wav_in;
    if (cfg.fixed_n_frames > 0) {
        int target_samples = cfg.fixed_n_frames * cfg.hop_size;
        if ((int)wav.size() > target_samples) wav.resize(target_samples);
        else if ((int)wav.size() < target_samples) wav.resize(target_samples, 0.0f);
    }

    auto pad  = pad_reflect(wav, pad_amount);
    auto hann = make_window(cfg.win_size, cfg.window);
    const int len_pad = (int)pad.size();
    if (len_pad < cfg.n_fft) {
        fprintf(stderr, "kazeia_mel::compute: wav trop court (%d < %d)\n", len_pad, cfg.n_fft);
        out_T = 0; return {};
    }

    // T_compute : sans fixed_n_frames -> calcule autant que possible (center=False).
    // Avec fixed_n_frames -> exactement N_FRAMES (le padding ci-dessus garantit la place).
    int T;
    if (cfg.fixed_n_frames > 0) T = cfg.fixed_n_frames;
    else                         T = (len_pad - cfg.n_fft) / cfg.hop_size + 1;
    out_T = T;

    // IMPORTANT : si n_fft n'est pas une puissance de 2, on DOIT faire une DFT directe
    // (pas un FFT zero-padded), sinon les bins du FFT sont à des fréquences décalées
    // par rapport à celles attendues par mel_basis (cf bug "empty output" sur certains
    // audios — Whisper N_FFT=400 -> DFT exigée, mel_filters calculé pour fs/400=40Hz/bin).
    const int fftN = next_pow2(cfg.n_fft);
    const bool is_pow2 = (fftN == cfg.n_fft);

    std::vector<float> spec((size_t)fft_bins * T, 0.0f);

    if (is_pow2) {
        // FAST PATH : FFT radix-2 (TTS speaker encoder N_FFT=1024)
        std::vector<float> re(fftN), im(fftN);
        for (int t = 0; t < T; t++) {
            const int off = t * cfg.hop_size;
            for (int i = 0; i < cfg.win_size; i++) { re[i] = pad[off + i] * hann[i]; im[i] = 0.0f; }
            for (int i = cfg.win_size; i < fftN; i++) { re[i] = 0.0f; im[i] = 0.0f; }
            fft_radix2(re.data(), im.data(), fftN);
            for (int k = 0; k < fft_bins; k++)
                spec[(size_t)k * T + t] = re[k]*re[k] + im[k]*im[k];
        }
    } else {
        // CORRECT PATH : DFT directe N=n_fft (Whisper N_FFT=400)
        // Precompute twiddles cos/sin[k=0..fft_bins-1][n=0..N-1]. ~256 KB pour Whisper.
        const int N = cfg.n_fft;
        std::vector<float> cos_tbl((size_t)fft_bins * N), sin_tbl((size_t)fft_bins * N);
        for (int k = 0; k < fft_bins; k++) {
            for (int n = 0; n < N; n++) {
                float angle = -2.0f * (float)M_PI * (float)k * (float)n / (float)N;
                cos_tbl[(size_t)k * N + n] = std::cos(angle);
                sin_tbl[(size_t)k * N + n] = std::sin(angle);
            }
        }
        std::vector<float> windowed((size_t)N);
        for (int t = 0; t < T; t++) {
            const int off = t * cfg.hop_size;
            for (int i = 0; i < N; i++) windowed[i] = pad[off + i] * hann[i];
            for (int k = 0; k < fft_bins; k++) {
                const float * ct = &cos_tbl[(size_t)k * N];
                const float * st = &sin_tbl[(size_t)k * N];
                float re = 0.0f, im = 0.0f;
                for (int n = 0; n < N; n++) { re += windowed[n] * ct[n]; im += windowed[n] * st[n]; }
                spec[(size_t)k * T + t] = re*re + im*im;
            }
        }
    }

    // mel = basis [n_mels, fft_bins] @ spec [fft_bins, T] -> [n_mels, T]
    std::vector<float> mel((size_t)cfg.n_mels * T, 0.0f);
    for (int m = 0; m < cfg.n_mels; m++) {
        for (int t = 0; t < T; t++) {
            float s = 0.0f;
            for (int k = 0; k < fft_bins; k++) {
                s += mel_basis[m * fft_bins + k] * spec[(size_t)k * T + t];
            }
            mel[(size_t)m * T + t] = std::max(s, cfg.clip_eps);
        }
    }

    // Compression (log naturel ou log10) + sqrt si TTS (variante: TTS prend sqrt magnitude
    // AVANT projection mel ; Whisper prend power direct projetée puis log10).
    // Distinction empirique du speaker encoder Qwen3-TTS : il calcule magnitude = sqrt(re²+im²+1e-9)
    // puis projette puis log. Pour rester strictement bit-correct avec le code spk d'origine,
    // on adapte si compression == LOG_NATURAL : sqrt avant log.
    if (cfg.comp == Compression::LOG_NATURAL) {
        // TTS : sqrt(power+1e-9) puis mel_proj puis log(max(., eps))
        // Refait correctement : revert la projection power, refais en magnitude.
        // Implémenter ainsi pour bit-match :
        std::vector<float> mag_spec((size_t)fft_bins * T);
        for (int k = 0; k < fft_bins; k++)
            for (int t = 0; t < T; t++)
                mag_spec[(size_t)k * T + t] = std::sqrt(spec[(size_t)k * T + t] + 1e-9f);
        for (int m = 0; m < cfg.n_mels; m++) {
            for (int t = 0; t < T; t++) {
                float s = 0.0f;
                for (int k = 0; k < fft_bins; k++)
                    s += mel_basis[m * fft_bins + k] * mag_spec[(size_t)k * T + t];
                mel[(size_t)m * T + t] = std::log(std::max(s, cfg.clip_eps));
            }
        }
    } else {  // LOG10 (Whisper)
        for (auto & v : mel) v = std::log10(v);
    }

    // Post-norm (Whisper : clamp à max-8, puis (x+4)/4)
    if (cfg.post == PostNorm::WHISPER) {
        float maxv = *std::max_element(mel.begin(), mel.end());
        float floor_v = maxv - 8.0f;
        for (auto & v : mel) {
            if (v < floor_v) v = floor_v;
            v = (v + 4.0f) * 0.25f;
        }
    }

    return mel;
}

}  // namespace kazeia_mel