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

#include "hark-lib/node/ChannelSelector.h"
#include "hark-lib/matrix.h"

namespace hark
{
namespace node
{
    ChannelSelector::ChannelSelector():
        NodeDef("ChannelSelector")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("SELECTOR")
            >> Terminal::out("OUTPUT");
    }

    void
    ChannelSelector::reset()
    {
        this->_channel_count = 0;
        this->_length = 0;
        this->_selector.clear();
    }

    template<typename T>
    data
    ChannelSelector::select_channels(size_t count, data inputData)
    {
        const Matrix<T> &inputMatrix = inputData.cast<Matrix<T>>();

        data outputData = data::create<Matrix<T>>(this->_selector.size(), inputMatrix.cols());
        Matrix<T> &outputMatrix = outputData.cast<Matrix<T>>();

        if(count == 0){
            // check dimensionality
            this->_channel_count = inputMatrix.rows();
            this->_length        = inputMatrix.cols();

            for(size_t ch = 0; ch < this->_selector.size(); ch++){
                if(this->_channel_count <= this->_selector[ch]){
                    std::stringstream ss;
                    ss << "Selected channel " << this->_selector[ch] << " does not exist in input.  Check the input SELECTOR";
                    throw RuntimeError(ss.str());
                }
            }
        }

        if(inputMatrix.rows() != this->_channel_count || inputMatrix.cols() != this->_length){
            std::stringstream ss;
            ss << "Unexpected matrix size: (" << inputMatrix.rows() << ", " << inputMatrix.cols() << "), expected (" << this->_channel_count << ", " << this->_length << ")";
            throw RuntimeError(ss.str());
        }

        for(size_t c = 0; c < this->_selector.size(); c++){
            for(size_t j = 0; j < this->_length; j++){
                outputMatrix(c, j) = inputMatrix(this->_selector[c], j);
            }
        }

        return outputData;
    }

    bool
    ChannelSelector::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            const std::string &selector_string = input["SELECTOR"].last().cast<std::string>();
            std::istringstream ss(selector_string);

            std::vector<size_t> selector;

            size_t id;
            while(ss >> id){
                selector.push_back(id);
            }

            this->_selector.swap(selector);
            this->_channel_count = 0;
            this->_length = 0;
        }

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

        if(inputData.is_type<Matrix<float>>()){
            outputData = select_channels<float>(count, inputData);

        }else if(inputData.is_type<Matrix<std::complex<float>>>()){
            outputData = select_channels<std::complex<float>>(count, inputData);

        }else{
            throw RuntimeError("Unexpected input data type");

        }

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

/* end of file */
