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

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

namespace hark
{
namespace node
{

    WhiteNoiseAdder::WhiteNoiseAdder():
        NodeDef("WhiteNoiseAdder")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("LENGTH").def(512)
            << Terminal::in("WN_LEVEL").def(0)
            >> Terminal::out("OUTPUT");
    }
 
    void
    WhiteNoiseAdder::reset()
    {
        this->gen = std::mt19937(1234);
    }
 
    void
    WhiteNoiseAdder::initialize(InputDataBundle &input)
    {
        this->length = input["LENGTH"].last().as<int>();
        this->wn_level = input["WN_LEVEL"].last().as<int>();

        this->wn_factor = std::sqrt(this->length) / 2.0f;
        this->wn_coef = this->wn_factor * this->wn_level;

        this->gen = std::mt19937(1234);
        this->rand_uniform = std::uniform_real_distribution<float>(0.0f, 1.0f);
    }

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

        data inputData = input["INPUT"].last();
        data outputData;

        if(inputData.is_type<std::map<int, data> >()){
            const auto &inputMap = inputData.cast<std::map<int, data> >();

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

            for(auto it: inputMap){
                const auto &vec = it.second.cast<std::vector<std::complex<float> > >();

                data outputVecData = data::create<std::vector<std::complex<float> > >(vec.size());
                std::vector<std::complex<float> > &output_vec = outputVecData.cast<std::vector<std::complex<float> > >();

                for(size_t i = 0; i < vec.size(); i++){
                    float theta = this->rand_uniform(gen) * this->coef;
                    output_vec[i] = (vec[i] + this->wn_coef * std::complex<float>(std::cos(theta), std::sin(theta)));
                }
                outputMap[it.first] = outputVecData;
            }
        }
        else{
            std::stringstream ss;
            ss << "PowerCalcForMap: Unexpected input type. Input must be hark::data(std::map<int, hark::data(std::vector<std::complex<float> >)>) .";
            throw RuntimeError(ss.str());
        }

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

/* end of file */
