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

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

#include "hark-lib/source.h"

#include <set>
#include <vector>
#include <fstream>

namespace hark
{
namespace node
{
    SaveFeatures::SaveFeatures():
        NodeDef("SaveFeatures")
    {
        *this
            << Terminal::in("FEATURES")
            << Terminal::in("SOURCES")
            << Terminal::in("BASENAME");
    }

    void
    SaveFeatures::reset()
    {
        this->_streams.clear();
    }

    bool
    SaveFeatures::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        const std::map<int, data> &inputMap = input["FEATURES"].last().as<std::map<int, data>>();
        const std::vector<data> &sources = input["SOURCES"].last().as<std::vector<data>>();
        const std::string &basename = input["BASENAME"].last().as<std::string>();

        // {
        //     std::stringstream ss;
        //     ss << "sources: count=" << inputMap.size() << ": ";
        //     for(auto it: inputMap){
        //         ss << it.first << ", ";
        //     }
        //     ss << std::endl;
        //     std::cout << ss.str();
        // }
        // std::cout << "SaveFeatures: count=" << count << ", input=" << inputMap << ", sources=" << sources << std::endl;

        std::set<int>source_ids;
        if(false /*sources.isNil()*/){
            for(auto it: inputMap){
                source_ids.insert(it.first);
            }
        }else{
            for(auto &src_data: sources){
                const Source &src = src_data.cast<Source>();
                source_ids.insert(src.id);
            }
        }

        std::set<int>source_ids_vanished;
        for(auto it: this->_streams){
            source_ids_vanished.insert(it.first);
        }

        for(auto source_id: source_ids){
            auto it = this->_streams.find(source_id);
            if(it == this->_streams.end()){
                std::ostringstream filename;
                filename << basename << source_id << ".spec";
                this->_streams[source_id].reset(new std::ofstream(filename.str().c_str(), std::ios::binary));
            }else{
                source_ids_vanished.erase(source_id);
            }

            auto fit = inputMap.find(source_id);
            if(fit != inputMap.end()){
                const std::vector<float> &feature = fit->second.cast<std::vector<float>>();
                std::shared_ptr<std::ostream>stream = this->_streams[source_id];
                stream->write(reinterpret_cast<const char *>(feature.data()), sizeof(float) * feature.size());
            }
        }

        for(auto it: source_ids_vanished){
            this->_streams.find(it)->second->flush();
            this->_streams.erase(it);
        }

        return true;
    }
}
}

/* end of file */
