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

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

#include "hark-lib/extlib/fft.hpp"

#include <map>
#include <math.h>

namespace hark
{
namespace node
{
    MSLSExtraction::MSLSExtraction():
        NodeDef("MSLSExtraction")
    {
        *this
            << Terminal::in("FBANK")
            << Terminal::in("SPECTRUM")
            << Terminal::in("FBANK_COUNT").def(13)
            << Terminal::in("NORMALIZATION_MODE").def("CEPSTRAL")
            << Terminal::in("USE_LEGACY_MODE").def(true)
            << Terminal::in("USE_HTK_LIFTER").def(false)
            << Terminal::in("LIFTERING_COEF").def(22)
            << Terminal::in("USE_POWER").def(false)
            >> Terminal::out("OUTPUT");
    }

    void
    MSLSExtraction::reset()
    {
        this->fbank_count = 0;
        this->duration = 0;
        this->use_power = false;
        this->dim_static = 0;
        this->lifter.clear();
        this->cepstrum.clear();
    }

    void
    MSLSExtraction::initialize(InputDataBundle &input)
    {
        fbank_count = input["FBANK_COUNT"].last().as<int>();
        if(fbank_count > 0){
        }
        else{
            std::stringstream ss;
            ss << "MSLSExtraction: FBANK_COUNT must be positive value";
            throw RuntimeError(ss.str());
        }

        std::string mode = input["NORMALIZATION_MODE"].last().cast<std::string>();
        bool use_lecacy_mode = input["USE_LEGACY_MODE"].last().as<bool>();

        if(mode == "CEPSTRAL"){
            normalization_mode = use_lecacy_mode? LEGACY_MODE: CEPSTRAL_MODE;
        }
        else if(mode == "SPECTRAL"){
            normalization_mode = SPECTRAL_MODE;
        }
        else{
            std::stringstream ss;
            ss << "MSLSExtraction: NORMALIZATION_MODE is not CEPSTRAL nor SPECTRAL";
            throw RuntimeError(ss.str());
        }

        bool use_htk_lifter = input["USE_HTK_LIFTER"].last().as<bool>();
        int coef = input["LIFTERING_COEF"].last().as<int>();
        use_power = input["USE_POWER"].last().as<bool>();
        dim_static = fbank_count + (use_power? 1: 0);

        lifter.resize(fbank_count);
        if(use_htk_lifter){
            if(coef <= 0){coef = 1;}
            for(int i = 0; i < lifter.size(); i++){
                lifter[i] = 1.0f + (coef / 2.0f) * std::sin((i + 1) * (M_PI / coef));
            }
        }
        else{
            for(int i = 0; i < lifter.size(); i++){
                lifter[i] = 1.0f;
            }
        }
        lifter[0] = 0.0f;
        cepstrum.resize(fbank_count);
    }

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

        const std::map<int, data> &fbanks = input["FBANK"].last().cast<std::map<int, data>>();
        // {
        //     std::stringstream ss;
        //     ss << "MSLSExtraction: fbanks: count=" << fbanks.size() << ": ";
        //     for(auto it: fbanks){
        //         ss << it.first << ", ";
        //     }
        //     ss << std::endl;
        //     std::cout << ss.str();
        // }

        std::map<int, data>spectra;
        if(use_power){
            spectra = input["SPECTRUM"].last().cast<std::map<int, data>>();
        }

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

        for(auto it: fbanks){
            // std::cout << "fbank: " << it.first << std::endl;
            const std::vector<float> &in = it.second.cast<std::vector<float>>();
            data fbank_data = data::create<std::vector<float>>(dim_static * 2);
            std::vector<float> &fbank = fbank_data.cast<std::vector<float>>();

            int i = 0;
            while(i < fbank_count){
                fbank[i] = std::log(in[i] < 1.0f? 1.0f: in[i]);
                i++;
            }
            while(i < fbank.size()){
                fbank[i] = 0.0f;
                i++;
            }

            switch(normalization_mode){
            case LEGACY_MODE:
                LifterOld(fbank);
                break;
            case CEPSTRAL_MODE:
                Lifter(fbank);
                break;
            case SPECTRAL_MODE:
                SubtractMean(fbank);
                EmphasizeSpectralPeak(fbank);
                break;
            }

            if(use_power){
                auto itspec = spectra.find(it.first);

                float power = 0.0f;
                if(itspec != spectra.end()){
                    std::vector<std::complex<float>> &spec = itspec->second.cast<std::vector<std::complex<float>>>();

                    int k = 0;
                    power += std::norm(spec[k++]);
                    while(k < spec.size() - 1){
                        power += std::norm(spec[k++]) * 2;
                    }
                    power += std::norm(spec[k]);
                    power /= (spec.size() - 1) * 2;

                    //// This code is fast, but there are slight differences in results.
                    // int k = spec.size() - 1;
                    // power += std::norm(spec[k--]);
                    // while(k > 0){
                    //  power += std::norm(spec[k--]) * 2;
                    // }
                    // power += std::norm(spec[k]);
                    // power /= (spec.size() - 1) * 2;

                }
                fbank[dim_static - 1] = (power >= 1e-40? std::log(power): 0.0f);
            }

            output_map[it.first] = fbank_data;
            // std::cout << "size(output_map): " << output_map.size() << std::endl;
        }

        output["OUTPUT"].push(output_map_data);
        // std::cout << "MSLSExtraction: count=" << count << ", output=" << output_map << std::endl;
        return true;
    }

    void
    MSLSExtraction::Lifter(std::vector<float> &fbank)
    {
        dct(&fbank[0], &cepstrum[0], fbank_count);
        for(int i = 0; i < fbank_count; i++){
            cepstrum[i] *= lifter[i];
        }
        idct(&cepstrum[0], &fbank[0], fbank_count);
    }

    void
    MSLSExtraction::LifterOld(std::vector<float> &fbank)
    {
        dct(&fbank[0], &cepstrum[0], fbank_count);
        cepstrum[0] = 0.0f;
        for(int i = 13; i < fbank_count; i++){
            cepstrum[i] = 0.0f;
        }
        idct(&cepstrum[0], &fbank[0], fbank_count);
    }

    void
    MSLSExtraction::SubtractMean(std::vector<float> &fbank) const
    {
        float mean = 0.0f;

        for(int i = 0; i < fbank_count; i++){
            mean += fbank[i];
        }
        mean /= fbank_count;

        for(int i = 0; i < fbank_count; i++){
            fbank[i] -= mean;
        }

    }

    void MSLSExtraction::dct(const float *x, float *y, int n) const
    {
        for(int i = 0; i < n; i++){
            y[i] = 0;
            for(int j = 0; j < n; j++){
                y[i] += x[j] * std::cos(M_PI * (j + 0.5f) * i / n);
            }
            y[i] *= 2.0f;
        }
    }

    void MSLSExtraction::idct(const float *x, float *y, int n) const
    {
        const float a = 1.0 / n;
        for(int i = 0; i < n; i++){
            y[i] = x[0] * 0.5f;
            for(int j = 1; j < n; j++){
                y[i] += x[j] * std::cos(M_PI * j * (i + 0.5f) / n);
            }
            y[i] *= a;
        }
    }

    void
    MSLSExtraction::EmphasizeSpectralPeak(std::vector<float> &vec) const
    {
        const float p = 0.9f;
        for(int i = vec.size() - 1; i > 0; i--){
            vec[i] -= p * vec[i - 1];
        }
    }

}
}

/* end of file */
