#include "mel_gpu.hpp"
#include "backend.hpp"
#include "ggml_graph.hpp"
#include "ggml.h"
#include <cassert>
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif
#include <cmath>
#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif
#include <vector>

namespace pk {

// CONSTANT epsilon added to std in NeMo normalize_batch (features.py: 1e-5).
// Kept identical to MelFrontend::compute so the per-feature norm matches.
static constexpr double kNormEps = 1e-5;

GpuMel::GpuMel(const ModelLoader& ml) : k_(ml) {
    // DFT basis matching pk::rfft's convention (no normalization), so the matmul
    // re/im equal the FFT's:  re[b] =  Σ_n x[n] cos(2π b n / N)
    //                         im[b] = -Σ_n x[n] sin(2π b n / N)
    // Row-major [n_bins, n_fft]: dft_cos_[b*n_fft + n], dft_sin_[b*n_fft + n].
    const int N = k_.n_fft_;
    const int B = k_.n_bins_;
    dft_cos_.assign((size_t)B * N, 0.0f);
    dft_sin_.assign((size_t)B * N, 0.0f);
    const double two_pi = 2.0 * M_PI;
    for (int b = 0; b < B; ++b) {
        for (int n = 0; n < N; ++n) {
            const double ang = two_pi * (double)b * (double)n / (double)N;
            dft_cos_[(size_t)b * N + n] = (float)std::cos(ang);
            dft_sin_[(size_t)b * N + n] = (float)(-std::sin(ang));
        }
    }
}

void GpuMel::compute(const std::vector<float>& samples,
                     std::vector<float>& feats, int& n_mels, int& T) const {
    // Only mag_power == 2 is supported (power = re^2 + im^2) — true for all
    // supported models; the DFT-matmul graph squares re/im directly.
    assert(k_.mag_power_ == 2.0f && "GpuMel requires mag_power == 2");

    if (samples.empty()) { n_mels = k_.n_mels_; T = 0; feats.clear(); return; }

    const int S        = (int)samples.size();
    const int n_fft_   = k_.n_fft_;
    const int hop_     = k_.hop_;
    const int n_mels_  = k_.n_mels_;
    const int n_bins_  = k_.n_bins_;
    const float preemph_ = k_.preemph_;

    // ----- Preemphasis (double, exactly like MelFrontend) -----
    std::vector<double> x((size_t)S);
    if (preemph_ > 0.0f && S > 0) {
        x[0] = samples[0];
        for (int t = 1; t < S; ++t)
            x[t] = (double)samples[t] - (double)preemph_ * (double)samples[t - 1];
    } else {
        for (int t = 0; t < S; ++t) x[t] = samples[t];
    }

    const int seq_len = (S > 0) ? (S / hop_) : 0;

    // ----- Center zero-pad (constant), then frame count (identical to MelFrontend) -----
    const int pad = n_fft_ / 2;
    const int padded_len = S + 2 * pad;
    const int n_frames = (padded_len >= n_fft_) ? (1 + (padded_len - n_fft_) / hop_) : 0;
    T = n_frames;
    n_mels = n_mels_;

    std::vector<double> padded((size_t)padded_len, 0.0);
    for (int j = 0; j < S; ++j) padded[pad + j] = x[j];

    feats.assign((size_t)n_mels_ * T, 0.0f);
    if (T == 0) {
        // Mirror MelFrontend: per_feature norm over zero frames is a no-op.
        return;
    }

    // ----- Windowed frame matrix Xw [T, n_fft] row-major (row t = frame t) -----
    // Xw[t*n_fft + i] = (float)(padded[t*hop + i] * (double)window_[i]). Matches
    // MelFrontend's per-frame windowing exactly.
    std::vector<float> Xw((size_t)T * n_fft_);
    for (int t = 0; t < T; ++t) {
        const int start = t * hop_;
        float* row = &Xw[(size_t)t * n_fft_];
        for (int i = 0; i < n_fft_; ++i)
            row[i] = (float)(padded[(size_t)start + i] * (double)k_.window_[i]);
    }

    // ----- ggml graph: DFT-matmul -> power -> filterbank -> log -> feat-major -----
    const float guard = k_.log_guard_;
    bool ok = pk::run_graph(0, 0,
        [&](ggml_context* ctx) -> ggml_tensor* {
            // Xw input: ne=[n_fft, T] (ne0=n_fft fastest, matches row-major Xw).
            int64_t xw_ne[2] = { n_fft_, T };
            ggml_tensor* xw = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, xw_ne,
                                  Xw.data(), (size_t)T * n_fft_ * sizeof(float));
            // DFT basis: ne=[n_fft, n_bins] (row b over n).
            int64_t basis_ne[2] = { n_fft_, n_bins_ };
            ggml_tensor* cosb = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, basis_ne,
                                    dft_cos_.data(), (size_t)n_bins_ * n_fft_ * sizeof(float));
            ggml_tensor* sinb = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, basis_ne,
                                    dft_sin_.data(), (size_t)n_bins_ * n_fft_ * sizeof(float));

            // re[t][b] = Σ_n cos(b,n) * Xw[t][n]; im likewise. -> ne=[n_bins, T].
            ggml_tensor* re = ggml_mul_mat(ctx, cosb, xw);
            ggml_tensor* im = ggml_mul_mat(ctx, sinb, xw);

            // power = re^2 + im^2  (mag_power == 2). -> ne=[n_bins, T].
            ggml_tensor* power = ggml_add(ctx,
                                          ggml_mul(ctx, re, re),
                                          ggml_mul(ctx, im, im));

            // Filterbank: ne=[n_bins, n_mels] (row m = fb_[m*n_bins + b]).
            int64_t fb_ne[2] = { n_bins_, n_mels_ };
            ggml_tensor* fb = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, fb_ne,
                                  k_.fb_.data(), (size_t)n_mels_ * n_bins_ * sizeof(float));
            // mel[t][m] = Σ_b fb(m,b) * power[t][b]. -> ne=[n_mels, T].
            ggml_tensor* mel = ggml_mul_mat(ctx, fb, power);

            // log(mel + log_guard). add via a [1] broadcast tensor (ggml_add1 is
            // deprecated); ne all 1 -> ggml_can_repeat into [n_mels, T].
            int64_t g_ne[1] = { 1 };
            ggml_tensor* g = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, g_ne,
                                 &guard, sizeof(float));
            mel = ggml_add(ctx, mel, g);
            ggml_tensor* lm = ggml_log(ctx, mel);   // ne=[n_mels, T]

            // Transpose to feat-major: out (t,m) at m*T + t. -> ne=[T, n_mels].
            ggml_tensor* out_t = ggml_cont(ctx, ggml_transpose(ctx, lm));
            return out_t;
        }, feats);
    assert(ok && "GpuMel graph failed");
    (void)ok;

    // run_graph wrote n_mels_*T f32 in feat-major order (feats[m*T + t]).

    // ----- per-feature normalization (verbatim from MelFrontend::compute) -----
    if (k_.per_feature_ && T > 0) {
        const int valid = std::min(seq_len, T);
        for (int m = 0; m < n_mels_; ++m) {
            float* row = &feats[(size_t)m * T];
            double mean = 0.0;
            if (valid > 0) {
                for (int t = 0; t < valid; ++t) mean += row[t];
                mean /= (double)valid;
            }
            double var = 0.0;
            if (valid > 1) {
                for (int t = 0; t < valid; ++t) {
                    const double d = (double)row[t] - mean;
                    var += d * d;
                }
                var /= (double)(valid - 1); // unbiased (ddof=1)
            }
            double sd = std::sqrt(var);
            sd += kNormEps;
            for (int t = 0; t < T; ++t) {
                if (t < valid) row[t] = (float)(((double)row[t] - mean) / sd);
                else           row[t] = 0.0f; // masked beyond seq_len
            }
        }
    }
}

} // namespace pk