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

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

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

namespace hark
{
namespace node
{

    PreEmphasis::PreEmphasis():
        NodeDef("PreEmphasis")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("SAMPLING_RATE").def(16000)
            << Terminal::in("PREEMCOEF").def(0.97f)
            << Terminal::in("INPUT_TYPE").def("WAV")
            >> Terminal::out("OUTPUT");
    }
 
    void
    PreEmphasis::reset()
    {
        this->length = 0;
        this->pslength = 0;
        this->sampling_rate = 0;
        this->preemcoef = 0;
        this->separ_length = 0;
        this->cutoff = 0;
        this->input_type = "WAV";
        this->preempha.clear();
    }
 
    void
    PreEmphasis::update_parameters(InputDataBundle &input)
    {
        this->length = input["LENGTH"].last().as<int>();
        this->sampling_rate = input["SAMPLING_RATE"].last().as<int>();
        this->preemcoef = input["PREEMCOEF"].last().as<float>();
        this->input_type = input["INPUT_TYPE"].last().cast<std::string>();

        if(this->input_type == "SPECTRUM"){
            this->preemcoef = -this->preemcoef;
            this->pslength = this->length / 2 + 1;
            this->cutoff = 7900;
            this->separ_length = this->length * this->cutoff / this->sampling_rate;

            rfft<float>rfftf;

            // making a spectral filter
            this->preempha.resize(this->length);
            rfftf.initialize(this->length);
            rfftf.a[0] = 1;
            rfftf.a[1] = this->preemcoef;
            for(int i = 2; i < this->length; i++){
                rfftf.a[i] = 0;
            }
            rfftf.analyze();

            for(int i = 0; i < this->pslength; i++){
                if(i < 4 ||  i > this->separ_length){
                    rfftf.A[i] = 0.0f;
                }
                this->preempha[i] = rfftf.A[i];
            }
        }
    }

    bool
    PreEmphasis::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->update_parameters(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>>();

        if(this->input_type == "WAV"){
            for(auto it: inputMap){
                const std::vector<float> &wav = it.second.cast<std::vector<float>>();

                data outData = data::create<std::vector<float>>(length);
                std::vector<float> &out = outData.cast<std::vector<float>>();

                for(int i = wav.size() - 1; i >= 1; i--){
                    out[i] = wav[i] - wav[i - 1] * this->preemcoef;
                }
                out[0] = wav[0] * (1.0f - this->preemcoef);

                outputMap[it.first] = outData;
            }

        }else if(this->input_type == "SPECTRUM"){
            for(auto it: inputMap){
                const std::vector<std::complex<float>> &spec = it.second.cast<std::vector<std::complex<float>>>();
                data outData = data::create<std::vector<std::complex<float>>>(length);
                std::vector<std::complex<float>> &out = outData.cast<std::vector<std::complex<float>>>();

                for(int i = 0; i < spec.size(); i++){
                    out[i] = spec[i] * this->preempha[i];
                }
                outputMap[it.first] = outData;
            }

        }
        else{
            std::stringstream ss;
            ss << "PreEmphasis: Unexpected input-type '" << this->input_type << "'";
            throw RuntimeError(ss.str());
        }

        output["OUTPUT"].push(outputMapData);
        return true;
    }
}
}

/* end of file */
