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

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

namespace hark
{
namespace node
{

    PowerCalcForMatrix::PowerCalcForMatrix():
        NodeDef("PowerCalcForMatrix")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("POWER_TYPE").def("POW")
            >> Terminal::out("OUTPUT")
            >> Terminal::out("PHASE");
    }
 
    void
    PowerCalcForMatrix::reset()
    {
        this->power_type = POWER_TYPE::POW;
    }
 
    void
    PowerCalcForMatrix::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 << "PowerCalcForMatrix: Unexpected POWER_TYPE \"" << power_type_str << "\" found.\n" << "POWER_TYPE must be \"POW\" or \"MAG\".\n";
            throw hark::RuntimeError(ss.str());
        }
    }

    bool
    PowerCalcForMatrix::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<Matrix<std::complex<float> > >()){
            const auto &inputMatrix = inputData.cast<Matrix<std::complex<float> > >();

            outputData = data::create<Matrix<float> >(inputMatrix.rows(), inputMatrix.cols());
            auto &outputMatrix = outputData.cast<Matrix<float> >();

            phaseData = data::create<Matrix<std::complex<float> > >(inputMatrix.rows(), inputMatrix.cols());
            auto &phaseMatrix = phaseData.cast<Matrix<std::complex<float> > >();

            if(this->power_type == POWER_TYPE::POW){
                for(size_t i = 0; i < inputMatrix.rows(); i++){
                    for(size_t j = 0; j < inputMatrix.cols(); j++){
                        float tmp = std::abs(inputMatrix(i, j));
                        outputMatrix(i, j) = std::norm(inputMatrix(i, j));
                        phaseMatrix(i, j) = (tmp == 0.0f) ? std::complex<float>(0.0f, 0.0f) : (inputMatrix(i, j) / tmp);
                    }
                }
            }
            else if(this->power_type == POWER_TYPE::MAG){
                for(size_t i = 0; i < inputMatrix.rows(); i++){
                    for(size_t j = 0; j < inputMatrix.cols(); j++){
                        float tmp = std::abs(inputMatrix(i, j));
                        outputMatrix(i, j) = tmp;
                        phaseMatrix(i, j) = (tmp == 0.0f) ? std::complex<float>(0.0f, 0.0f) : (inputMatrix(i, j) / tmp);
                    }
                }
            }
        }
        else{
            std::stringstream ss;
            ss << "PowerCalcForMatrix: Unexpected input type. Input must be hark::data(hark::Matrix<std::complex<float> >) .";
            throw RuntimeError(ss.str());
        }

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

/* end of file */
