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

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

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

    void
    FeatureRemover::reset()
    {
        this->_id_tobe_removed.clear();
    }

    void
    FeatureRemover::initialize(InputDataBundle &input)
    {
        const std::string &selector = input["SELECTOR"].last().cast<std::string>();
        std::istringstream ss(selector);

        std::vector<size_t>id_tobe_removed;

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

        this->_id_tobe_removed.swap(id_tobe_removed);
    }

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

        const std::map<int, data> &inputMap = input["INPUT"].last().cast<std::map<int, data>>();

        data outputMapData = data::create<std::map<int, data>>();
        auto &outputMap = outputMapData.cast<std::map<int, data>>();
        for(auto it: inputMap){
            const std::vector<float> &in_feature = it.second.cast<std::vector<float>>();

            data out_feature_data = data::create<std::vector<float>>();
            auto &out_feature = out_feature_data.cast<std::vector<float>>();

            for(int i = 0; i < in_feature.size(); i++){
                if(std::find(this->_id_tobe_removed.begin(), this->_id_tobe_removed.end(), i) == this->_id_tobe_removed.end()){
                    out_feature.push_back(in_feature[i]);
                }
            }

            outputMap[it.first] = data(out_feature);
        }

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

}
}

/* end of file */
