#include "acoustic_analyzer/core/feature_extractor.h"
#include "acoustic_analyzer/core/fft_utils.h"
#include <cmath>
#include <algorithm>
#include <numeric>
#include <fstream>
#include <memory>

namespace acoustic {

struct FeatureExtractor::Impl {
    int sample_rate;
    int n_fft;
    int hop_length;
    int n_mels;
    float f_min;
    float f_max;
    float preemphasis;
    std::vector<float> window;
    std::vector<float> mel_filter_bank;
    bool mel_loaded = false;

    Impl(int sr, int fft, int hop, int mels, float fmin, float fmax, float pre)
        : sample_rate(sr), n_fft(fft), hop_length(hop), n_mels(mels),
          f_min(fmin), f_max(fmax), preemphasis(pre) {
        window.resize(n_fft);
        for (int i = 0; i < n_fft; ++i) {
            window[i] = 0.5f - 0.5f * std::cos(2.0f * static_cast<float>(M_PI) * i / (n_fft - 1));
        }
        int n_fft_bins = n_fft / 2 + 1;
        mel_filter_bank.resize(n_mels * n_fft_bins);
        // Try to load from models/mel_filter_bank.bin
        std::ifstream f("models/mel_filter_bank.bin", std::ios::binary);
        if (f) {
            f.read(reinterpret_cast<char*>(mel_filter_bank.data()),
                   mel_filter_bank.size() * sizeof(float));
            if (f.gcount() == static_cast<std::streamsize>(mel_filter_bank.size() * sizeof(float))) {
                mel_loaded = true;
            }
        }
        if (!mel_loaded) {
            GenerateMelFilterBank();
            mel_loaded = true;
        }
    }

    void GenerateMelFilterBank() {
        int n_fft_bins = n_fft / 2 + 1;
        float mel_min = 2595.0f * std::log10(1.0f + f_min / 700.0f);
        float mel_max = 2595.0f * std::log10(1.0f + f_max / 700.0f);
        std::vector<float> mel_points(n_mels + 2);
        for (int i = 0; i < n_mels + 2; ++i) {
            float mel = mel_min + i * (mel_max - mel_min) / (n_mels + 1);
            mel_points[i] = 700.0f * (std::pow(10.0f, mel / 2595.0f) - 1.0f);
        }
        std::vector<float> bin_points(n_mels + 2);
        for (int i = 0; i < n_mels + 2; ++i) {
            bin_points[i] = std::floor((n_fft + 1) * mel_points[i] / sample_rate);
        }
        for (int i = 0; i < n_mels; ++i) {
            for (int j = 0; j < n_fft_bins; ++j) {
                float val = 0.0f;
                float left = bin_points[i];
                float center = bin_points[i + 1];
                float right = bin_points[i + 2];
                if (j >= left && j <= center && center > left) {
                    val = (j - left) / (center - left);
                } else if (j >= center && j <= right && right > center) {
                    val = (right - j) / (right - center);
                }
                mel_filter_bank[i * n_fft_bins + j] = val;
            }
        }
    }

    Eigen::MatrixXf ComputeMelSpec(const std::vector<float>& audio) {
        std::size_t n_samples = audio.size();
        int n_fft_bins = n_fft / 2 + 1;
        int n_frames = static_cast<int>((static_cast<int>(n_samples) - n_fft) / hop_length) + 1;
        if (n_frames < 1) n_frames = 1;

        std::vector<float> preemph(n_samples);
        apply_preemphasis(audio.data(), preemph.data(), n_samples, preemphasis);

        std::vector<float> fft_buf(n_fft);
        std::vector<float> power(n_fft_bins);
        Eigen::MatrixXf mel_spec(n_mels, n_frames);

        for (int t = 0; t < n_frames; ++t) {
            int start = t * hop_length;
            for (int i = 0; i < n_fft; ++i) {
                std::size_t idx = start + i;
                fft_buf[i] = (idx < n_samples ? preemph[idx] : 0.0f) * window[i];
            }
            compute_power_spectrum(fft_buf.data(), n_fft, power.data());
            for (int m = 0; m < n_mels; ++m) {
                float sum = 0.0f;
                for (int b = 0; b < n_fft_bins; ++b) {
                    sum += power[b] * mel_filter_bank[m * n_fft_bins + b];
                }
                mel_spec(m, t) = std::log(std::max(sum, 1e-10f));
            }
        }

        // Normalize per mel bin (standardization)
        for (int m = 0; m < n_mels; ++m) {
            float mean = 0.0f;
            for (int t = 0; t < n_frames; ++t) {
                mean += mel_spec(m, t);
            }
            mean /= n_frames;
            float variance = 0.0f;
            for (int t = 0; t < n_frames; ++t) {
                float d = mel_spec(m, t) - mean;
                variance += d * d;
            }
            float stddev = std::sqrt(variance / n_frames + 1e-8f);
            for (int t = 0; t < n_frames; ++t) {
                mel_spec(m, t) = (mel_spec(m, t) - mean) / stddev;
            }
        }

        // Pad or truncate to target 63 frames with edge padding
        const int target_frames = 63;
        if (n_frames < target_frames) {
            Eigen::MatrixXf padded(n_mels, target_frames);
            int copy_frames = n_frames;
            int pad_left = (target_frames - copy_frames) / 2;
            padded.block(0, pad_left, n_mels, copy_frames) = mel_spec.block(0, 0, n_mels, copy_frames);
            // Edge padding on left
            for (int t = 0; t < pad_left; ++t) {
                padded.col(t) = mel_spec.col(0);
            }
            // Edge padding on right
            for (int t = pad_left + copy_frames; t < target_frames; ++t) {
                padded.col(t) = mel_spec.col(copy_frames - 1);
            }
            return padded;
        } else if (n_frames > target_frames) {
            return mel_spec.block(0, 0, n_mels, target_frames);
        }
        return mel_spec;
    }
};

FeatureExtractor::FeatureExtractor(int sample_rate,
                                   int n_fft,
                                   int hop_length,
                                   int n_mels,
                                   float f_min,
                                   float f_max,
                                   float preemphasis)
    : impl_(std::make_unique<Impl>(sample_rate, n_fft, hop_length, n_mels, f_min, f_max, preemphasis)) {}

FeatureExtractor::~FeatureExtractor() = default;
FeatureExtractor::FeatureExtractor(FeatureExtractor&&) noexcept = default;
FeatureExtractor& FeatureExtractor::operator=(FeatureExtractor&&) noexcept = default;

Eigen::MatrixXf FeatureExtractor::MelSpectrogram(const std::vector<float>& audio_samples) {
    return impl_->ComputeMelSpec(audio_samples);
}

std::vector<Eigen::MatrixXf> FeatureExtractor::MelSpectrogramMultiChannel(
    const std::vector<float>& audio_samples,
    std::size_t num_channels) {
    if (num_channels == 0) return {};
    std::size_t total = audio_samples.size();
    std::size_t samples_per_channel = total / num_channels;
    std::vector<Eigen::MatrixXf> results;
    results.reserve(num_channels);
    for (std::size_t ch = 0; ch < num_channels; ++ch) {
        std::vector<float> ch_audio(samples_per_channel);
        for (std::size_t i = 0; i < samples_per_channel; ++i) {
            ch_audio[i] = audio_samples[i * num_channels + ch];
        }
        results.push_back(impl_->ComputeMelSpec(ch_audio));
    }
    return results;
}

int FeatureExtractor::SampleRate() const noexcept { return impl_->sample_rate; }
int FeatureExtractor::NFFT() const noexcept { return impl_->n_fft; }
int FeatureExtractor::HopLength() const noexcept { return impl_->hop_length; }
int FeatureExtractor::NMels() const noexcept { return impl_->n_mels; }

}  // namespace acoustic
