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

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

namespace hark
{
namespace node
{
    CalcSpecSubGain::CalcSpecSubGain():
        NodeDef("CalcSpecSubGain")
    {
        *this
            << Terminal::in("INPUT_POWER_SPEC")
            << Terminal::in("NOISE_SPEC")
            << Terminal::in("ALPHA").def(1.0f)
            << Terminal::in("BETA").def(0.0f)
            << Terminal::in("SS_METHOD").def(2)
            >> Terminal::out("VOICE_PROB")
            >> Terminal::out("GAIN")
            >> Terminal::out("OUTPUT_POWER_SPEC");
    }

    void
    CalcSpecSubGain::reset()
    {
    }

    void
    CalcSpecSubGain::initialize(InputDataBundle &input)
    {
        this->alpha = input["ALPHA"].last().as<float>();
        this->beta = input["BETA"].last().as<float>();
        this->method = input["SS_METHOD"].last().as<int>();
        if(this->method != 1 && this->method != 2){
            std::stringstream ss;
            ss << "CalcSpecSubGain: SS_METHOD should be \"1\" or \"2\".\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }
    }

    bool
    CalcSpecSubGain::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        ////const auto &operation_flag = input["OPERATION_FLAG"].last().as<bool>();
        const auto &input_power_spectrum = input["INPUT_POWER_SPEC"].last().cast<std::map<int, data> >();
        const auto &noise_power_spectrum = input["NOISE_SPEC"].last().cast<std::map<int, data> >();

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

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

        auto out_voicep = data::create<std::map<int, data> >();
        auto &output_voice_probability = out_voicep.cast<std::map<int, data> >();

        auto out_gain = data::create<std::map<int, data> >();
        auto &output_optimum_gain = out_gain.cast<std::map<int, data> >();

        auto out = data::create<std::map<int, data> >();
        auto &output_power_spectrum = out.cast<std::map<int, data> >();

        if(input_power_spectrum.size() != noise_power_spectrum.size()){
            std::stringstream ss;
            ss << "CalcSpecSubGain: InputMap size not match.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
            throw RuntimeError(ss.str());
        }

        if(input_power_spectrum.size() > 0){
            std::map<int, data>::const_iterator it_in_spec = input_power_spectrum.begin();
            std::map<int, data>::const_iterator it_in_noise = noise_power_spectrum.begin();

            for(; it_in_spec != input_power_spectrum.end(); ++it_in_spec, ++it_in_noise){
                const std::vector<float> &input_spec = it_in_spec->second.cast<std::vector<float> >();
                const std::vector<float> &noise_spec = it_in_noise->second.cast<std::vector<float> >();
                std::vector<float> output_vp_vector;
                std::vector<float> output_gain_vector;
                std::vector<float> output_vector;
                if(it_in_spec->first != it_in_noise->first){
                    std::stringstream ss;
                    ss << "CalcSpecSubGain: InputMap keys not match.\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                    throw RuntimeError(ss.str());
                }
                output_vp_vector.resize(input_spec.size());
                output_gain_vector.resize(input_spec.size());
                output_vector.resize(input_spec.size());
                if(this->method == 1){
                    for(size_t k = 0; k < input_spec.size(); k++){
                        output_vector[k] = input_spec[k] - noise_spec[k];
                        if(output_vector[k] < 0.0f){
                            output_vector[k] = 0.0f;
                        }
                        output_gain_vector[k] = std::max(1.0f - this->alpha * (std::sqrt(noise_spec[k]) / std::max(std::sqrt(input_spec[k]), 0.0001f)), this->beta);
                        output_vp_vector[k] = 1.0f;
                    }
                }
                else if(this->method == 2){
                    for(size_t k = 0; k < input_spec.size(); k++){
                        output_vector[k] = input_spec[k] - noise_spec[k];
                        if(output_vector[k] < 0.0f){
                            output_vector[k] = 0.0f;
                        }
                        output_gain_vector[k] = 0.0f;
                        if(input_spec[k] != 0.0f){
                            output_gain_vector[k] = std::sqrt(output_vector[k] / input_spec[k]);
                        }
                        output_vp_vector[k] = 1.0f;
                    }
                }
                else{
                    std::stringstream ss;
                    ss << "CalcSpecSubGain: SS_METHOD should be \"1\" or \"2\".\n" << "FILE: " << __FILE__ << ", LINE: " << __LINE__;
                    throw RuntimeError(ss.str());
                }
                output_power_spectrum.emplace(it_in_spec->first, data(output_vector));
                output_optimum_gain.emplace(it_in_spec->first, data(output_gain_vector));
                output_voice_probability.emplace(it_in_spec->first, data(output_vp_vector));
            }
        }

        output["VOICE_PROB"].push(out_voicep);
        output["GAIN"].push(out_gain);
        output["OUTPUT_POWER_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 */
