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

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

namespace hark
{
namespace node
{

    PowerCalcForMap::PowerCalcForMap():
        NodeDef("PowerCalcForMap")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("POWER_TYPE").def("POW")
            >> Terminal::out("OUTPUT")
            >> Terminal::out("PHASE");
    }
 
    void
    PowerCalcForMap::reset()
    {
        this->power_type = POWER_TYPE::POW;
    }
 
    void
    PowerCalcForMap::initialize(InputDataBundle &input)
    {
        const auto &power_type_str = input["POWER_TYPE"].last().cast<std::string>();
        if(power_type_str == "POW"){ this->power_type = POWER_TYPE::POW; }
        else if(power_type_str == "MAG"){ this->power_type = POWER_TYPE::MAG; }
        else{
            std::stringstream ss;
            ss << "PowerCalcForMap: Unexpected POWER_TYPE \"" << power_type_str << "\" found.\n" << "POWER_TYPE must be \"POW\" or \"MAG\".\n";
            throw hark::RuntimeError(ss.str());
        }
    }

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

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

        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> >();

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

            if(this->power_type == POWER_TYPE::POW){
                for(auto it: inputMap){
                    const auto &vec = it.second.cast<std::vector<std::complex<float> > >();

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

                    for(size_t i = 0; i < vec.size(); i++){
                        float tmp = std::abs(vec[i]);
                        output_vec[i] = std::norm(vec[i]);
                        phase_vec[i] = (tmp == 0.0f) ? std::complex<float>(0.0f, 0.0f) : (vec[i] / tmp);
                    }

                    outputMap[it.first] = outputVecData;
                    phaseMap[it.first] = phaseVecData;
                }
            }
            else if(this->power_type == POWER_TYPE::MAG){
                for(auto it: inputMap){
                    const auto &vec = it.second.cast<std::vector<std::complex<float> > >();

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

                    for(size_t i = 0; i < vec.size(); i++){
                        float tmp = std::abs(vec[i]);
                        output_vec[i] = tmp;
                        phase_vec[i] = (tmp == 0.0f) ? std::complex<float>(0.0f, 0.0f) : (vec[i] / tmp);
                    }

                    outputMap[it.first] = outputVecData;
                    phaseMap[it.first] = phaseVecData;
                }
            }
        }
        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);
        output["PHASE"].push(phaseData);
        return true;
    }
}
}

/* end of file */
