
#include <cmath>
#include <complex>
#include <algorithm>
#include <iostream>

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

namespace hark
{
namespace node
{
    SpectralGainFilter::SpectralGainFilter():
        NodeDef("SpectralGainFilter")
    {
        *this
            << Terminal::in("INPUT_SPEC")
            << Terminal::in("VOICE_PROB")
            << Terminal::in("GAIN")
            >> Terminal::out("OUTPUT_SPEC");
    }

    void
    SpectralGainFilter::reset()
    {
    }

    void
    SpectralGainFilter::initialize(InputDataBundle &input)
    {
    }

    bool
    SpectralGainFilter::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        const auto &input_spectrum = input["INPUT_SPEC"].last().cast<std::map<int, std::vector<std::complex<float> > > >();
        const auto &input_voice_probability = input["VOICE_PROB"].last().cast<std::map<int, std::vector<float> > >();
        const auto &input_optimum_gain = input["GAIN"].last().cast<std::map<int, std::vector<float> > >();

#ifdef DEBUG
        ptime start_time;
        ptime end_time;
        start_time = microsec_clock::local_time();
#endif

        //if(count == 0){ this->initialize(input); }

        auto out = data::create<std::map<int, std::vector<std::complex<float> > > >();
        auto &output_spectrum = out.cast<std::map<int, std::vector<std::complex<float> > > >();

        if(input_spectrum.size() > 0){
            std::map<int, std::vector<std::complex<float> > >::const_iterator it_in_spec = input_spectrum.begin();

            for(; it_in_spec != input_spectrum.end(); ++it_in_spec){
                if(input_voice_probability.find(it_in_spec->first) == input_voice_probability.end()){
                    std::stringstream ss;
                    ss << "SpectralGainFilter: Failed to find ID in VOICE_PROB.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                    throw RuntimeError(ss.str());
                }
                if(input_optimum_gain.find(it_in_spec->first) == input_optimum_gain.end()){
                    std::stringstream ss;
                    ss << "SpectralGainFilter: Failed to find ID in GAIN.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                    throw RuntimeError(ss.str());
                }
                std::vector<std::complex<float> > output_vector;
                output_vector.resize(it_in_spec->second.size());
                std::map<int, std::vector<float> >::const_iterator vp_vector = input_voice_probability.find(it_in_spec->first);
                std::map<int, std::vector<float> >::const_iterator gain_vector = input_optimum_gain.find(it_in_spec->first);
                for(size_t k = 0; k < it_in_spec->second.size(); k++){
                    if(k < vp_vector->second.size()){
                        output_vector[k] = std::complex<float>(it_in_spec->second[k].real() * vp_vector->second[k] * gain_vector->second[k],
                                                               it_in_spec->second[k].imag() * vp_vector->second[k] * gain_vector->second[k]);
                    }
                    else{
                        output_vector[k] = it_in_spec->second[k];
                    }
                    output_spectrum.insert(std::pair<int, std::vector<std::complex<float> > >(it_in_spec->first, output_vector));
                }
            }
        }

        output["OUTPUT_SPEC"].push(out);

#ifdef DEBUG
        end_time = microsec_clock::local_time();
        std::cout << end_time - start_time << std::endl;
#endif

        return true;
    }
}
}

/* end of file */
