/*
   HARK Library
   Copyright (C) 2021 Honda Research Institute Japan Inc.
 */

#include "hark-lib/node/MelFilterBank.h"

#include <complex>

namespace hark
{
namespace node
{

    MelFilterBank::MelFilterBank():
        NodeDef("MelFilterBank")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("CUTOFF").def(8000)
            << Terminal::in("MIN_FREQUENCY").def(63)
            << Terminal::in("MAX_FREQUENCY").def(8000)
            << Terminal::in("FBANK_COUNT").def(13)
            >> Terminal::out("OUTPUT");
    }

    void
    MelFilterBank::reset()
    {
        this->_length = 0;
        this->_sampling_rate = 0;
        this->_cutoff = 0;
        this->_fbank_count = 0;
        this->_lo = 0;
        this->_hi = 0;
        this->_spectrum.clear();
        this->_lo_index.clear();
        this->_lo_weight.clear();
        this->_mel_center.clear();
    }

    void
    MelFilterBank::initialize(InputDataBundle &input)
    {
        //// parse values
        int length = input["LENGTH"].last().as<int>();
        int sampling_rate = input["SAMPLING_RATE"].last().as<int>();
        int cutoff = input["CUTOFF"].last().as<int>();
        int fbank_count = input["FBANK_COUNT"].last().as<int>();
        int min_freq = input["MIN_FREQUENCY"].last().as<int>();
        int max_freq = input["MAX_FREQUENCY"].last().as<int>();

        //// validate
        if(min_freq < 0 || min_freq > sampling_rate / 2){
            std::stringstream ss;
            ss << "MelFilterBank: Minimum frequency (" << min_freq << ") is not between 0 and half of sampling rate (" << sampling_rate << ")";
            throw RuntimeError(ss.str());
        }
        if(max_freq < 0 || max_freq > sampling_rate / 2){
            std::stringstream ss;
            ss << "MelFilterBank: Maximum frequency (" << max_freq << ") is not between 0 and half of sampling rate (" << sampling_rate << ")";
            throw RuntimeError(ss.str());
        }
        if(min_freq > max_freq){
            std::stringstream ss;
            ss << "MelFilterBank: Maximum frequency (" << max_freq << ") is less than minimum frequency (" << min_freq << ")";
            throw RuntimeError(ss.str());
        }

        //// calculate workset
        int lo = min_freq * length / sampling_rate;
        int hi = max_freq * length / sampling_rate;

        //// store values
        this->_length = length;
        this->_sampling_rate = sampling_rate;
        this->_cutoff = cutoff;
        this->_fbank_count = fbank_count;
        this->_lo = lo;
        this->_hi = hi;
        this->_spectrum.resize(length / 2 + 1);

        this->initialize_filters();
    }

    bool
    MelFilterBank::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->initialize(input);
        }

        const std::map<int, data> &inputMap = input["INPUT"].last().cast<std::map<int, data>>();

        data outputMapData = data::create<std::map<int, data>>();
        std::map<int, data> &outputMap = outputMapData.cast<std::map<int, data>>();

        for(auto &it: inputMap){
            int id = it.first;

            data fbankData = data::create<std::vector<float>>(this->_fbank_count * 2, 0);
            std::vector<float> &fbank = fbankData.cast<std::vector<float>>();

            if(it.second.is_type<std::vector<float>>()){
                this->filter(it.second.cast<std::vector<float>>(), fbank);

            }else if(it.second.is_type<std::vector<std::complex<float>>>()){
                const std::vector<std::complex<float>> &spec = it.second.cast<std::vector<std::complex<float>>>();
                std::vector<float>rspec(spec.size());
                std::transform(spec.begin(), spec.end(), rspec.begin(), [](const std::complex<float> &z){
                        return abs(z);
                    });

                this->filter(rspec, fbank);

            }else{
                std::stringstream ss;
                ss << "MelFilterBank: Unexpected data type.  std::vector<float> or std::vector<std::complex<float>> is expected";
                throw RuntimeError(ss.str());
            }

            outputMap[id] = fbankData;
        }
        output["OUTPUT"].push(outputMapData);
        // std::cout << "MelFilterBank: count=" << count << ", output=" << outputMap << std::endl;

        return true;
    }

    inline static
    float Mel(int k, int fs, int n){
        /* k  : frequency bin ID         */
        /* fs : sampling frequency [Hz]  */
        /* n  : analysis frame length    */
        return 1127 * log(1 + (k - 1) * fs * (1 / 700.0f) / n);
    }

    inline static
    float Mel(int f){
        /* f : frequency [Hz] */
        return 1127 * log(1 + f * (1 / 700.0f));
    }

    void
    MelFilterBank::initialize_filters()
    {
        int slength = this->_length / 2 + 1;

        this->_mel_center.resize(this->_fbank_count + 2);

        // float mel_hi = Mel(slength, fs, length);
        float mel_hi = Mel(this->_cutoff);
        float mel_lo = 0;
        float mel_s = mel_hi - mel_lo;
        for(int c = 1; c < this->_mel_center.size(); c++){
            this->_mel_center[c] = ((float)c / (this->_fbank_count + 1)) * mel_s + mel_lo;
        }

        this->_lo_index.resize(slength);
        for(int k = 1, c = 1; k < slength; k++){
            float mel_k = Mel(k, this->_sampling_rate, this->_length);
            while(this->_mel_center[c] < mel_k && c < this->_mel_center.size()){
                c++;
            }
            this->_lo_index[k] = c - 1;
        }

        this->_lo_weight.resize(slength);
        for(int k = 1; k <= this->_hi; k++){
            int c = this->_lo_index[k];
            if(c > 0){
                this->_lo_weight[k] = (this->_mel_center[c + 1] - Mel(k, this->_sampling_rate, this->_length)) /
                                      (this->_mel_center[c + 1] - this->_mel_center[c]);
            }
            else{
                this->_lo_weight[k] = (this->_mel_center[1] - Mel(k, this->_sampling_rate, this->_length)) /
                                      (this->_mel_center[1] - mel_lo);
            }
        }
    }

    void
    MelFilterBank::filter(const std::vector<float> &spec, std::vector<float> &fbank)
    {
        for(int i = 0; i < this->_fbank_count; i++){
            fbank[i] = 0.0f;
        }

        for(int k = this->_lo; k <= this->_hi; k++){
            float a = spec[k];
            int bin = this->_lo_index[k];
            float weighted_a = this->_lo_weight[k] * a;

            if(bin > 0 && bin <= this->_fbank_count){
                fbank[bin - 1] += weighted_a;
            }
            if(bin < this->_fbank_count){
                fbank[bin] += a - weighted_a;
            }
        }
    }
}
}

/* end of file */
