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

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

#include "hark-lib/source.h"

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

namespace hark
{
namespace node
{
    SaveHTKFeatures::SaveHTKFeatures():
        NodeDef("SaveHTKFeatures")
    {
        *this
            << Terminal::in("FEATURES")
            << Terminal::in("SOURCES")
            << Terminal::in("BASENAME").def("")
            << Terminal::in("HTK_PERIOD").def(100000)
            << Terminal::in("FEATURE_TYPE").def("USER");
    }

    SaveHTKFeatures::~SaveHTKFeatures()
    {
        for(std::map<int, std::shared_ptr<HTKFileIO> >::iterator it = this->htkfileio.begin(); it != this->htkfileio.end();){
            it->second->write_header();
            it->second->closefile();
            this->htkfileio.erase(it++);
        }
        this->htkfileio.clear();
    }

    void
    SaveHTKFeatures::reset()
    {
        for(std::map<int, std::shared_ptr<HTKFileIO> >::iterator it = this->htkfileio.begin(); it != this->htkfileio.end();){
            it->second->write_header();
            it->second->closefile();
            this->htkfileio.erase(it++);
        }
        this->htkfileio.clear();
    }

    void
    SaveHTKFeatures::initialize(InputDataBundle &input)
    {
        this->basename = input["BASENAME"].last().cast<std::string>();

        int period = input["HTK_PERIOD"].last().as<int>();
        this->init_header.Period = hark::byteswap(static_cast<uint32_t>(period));

        std::string feat_type = input["FEATURE_TYPE"].last().cast<std::string>();
        if(feat_type == "USER"){
            this->init_header.FeatureType = hark::byteswap(HTKFileIO::USER);
        }
        else if(feat_type == "MFCC_E_D"){
            this->init_header.FeatureType = hark::byteswap(HTKFileIO::MFCC_E_D);
        }
        else if(feat_type == "MFCC_E_D_Z"){
            this->init_header.FeatureType = hark::byteswap(HTKFileIO::MFCC_E_D_Z);
        }
        else{
            uint16_t feat_type_le = 0;
            if(feat_type.find("MFCC") != std::string::npos){ feat_type_le = HTKFileIO::MFCC; }
            else if(feat_type.find("FBANK") != std::string::npos){ feat_type_le = HTKFileIO::FBANK; }
            else if(feat_type.find("USER") != std::string::npos){ feat_type_le = HTKFileIO::USER; }
            else if(feat_type.find("PLP") != std::string::npos){ feat_type_le = HTKFileIO::PLP; }
            else if(feat_type.find("FSTREAM") != std::string::npos){ feat_type_le = HTKFileIO::FSTREAM; }
            else{  }

            if(feat_type.find("_E") != std::string::npos){ feat_type_le |= HTKFileIO::_E; }
            if(feat_type.find("_N") != std::string::npos){ feat_type_le |= HTKFileIO::_N; }
            if(feat_type.find("_D") != std::string::npos){ feat_type_le |= HTKFileIO::_D; }
            if(feat_type.find("_A") != std::string::npos){ feat_type_le |= HTKFileIO::_A; }
            if(feat_type.find("_C") != std::string::npos){ feat_type_le |= HTKFileIO::_C; }
            if(feat_type.find("_Z") != std::string::npos){ feat_type_le |= HTKFileIO::_Z; }
            if(feat_type.find("_K") != std::string::npos){ feat_type_le |= HTKFileIO::_K; }
            if(feat_type.find("_0") != std::string::npos){ feat_type_le |= HTKFileIO::_0; }
            if(feat_type.find("_V") != std::string::npos){ feat_type_le |= HTKFileIO::_V; }
            if(feat_type.find("_T") != std::string::npos){ feat_type_le |= HTKFileIO::_T; }

            this->init_header.FeatureType = hark::byteswap(feat_type_le);
        }

        this->init_header.nSamples = 0;
        this->init_header.BytesPerSample = 0;
    }

    bool
    SaveHTKFeatures::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        const auto &input_features = input["FEATURES"].last().as<std::map<int, data> >();
        const auto &input_sources = input["SOURCES"].last().as<std::vector<data> >();

        if(count == 0){ this->initialize(input); }

        // {
        //     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(input_sources.empty()){
            for(auto it: input_features){
                source_ids.insert(it.first);
            }
        }
        else{
            for(auto &src_data: input_sources){
                const Source &src = src_data.cast<Source>();
                source_ids.insert(src.id);
            }
        }

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

        for(auto source_id: source_ids){
            auto it = this->htkfileio.find(source_id);
            if(it == this->htkfileio.end()){
                std::ostringstream filename_ss;
                filename_ss << basename << source_id << ".spec";
                std::string filename = filename_ss.str();
                this->htkfileio[source_id].reset(new HTKFileIO());
                this->htkfileio[source_id]->openfile(filename, true);
                this->htkfileio[source_id]->set_header(this->init_header);
                this->htkfileio[source_id]->write_header();
            }else{
                source_ids_vanished.erase(source_id);
            }

            auto fit = input_features.find(source_id);
            if(fit != input_features.end()){
                const std::vector<float> &feature = fit->second.cast<std::vector<float> >();
                this->htkfileio[source_id]->write_features(feature);
                this->htkfileio[source_id]->set_samples(this->htkfileio[source_id]->get_samples() + 1);
                this->htkfileio[source_id]->set_bps(sizeof(float) * feature.size());
                // Case1. Update the header every frame.
                this->htkfileio[source_id]->write_header();
            }
        }

        for(auto it: source_ids_vanished){
            // Case2. Update headers only before file closing.
            //this->htkfileio.find(it)->second->write_header();
            this->htkfileio.find(it)->second->closefile();
            this->htkfileio.erase(it);
        }

        return true;
    }
}
}

/* end of file */
