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

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

#include <fstream>
#include <sstream>

namespace hark
{
namespace node
{
    SpectralMeanNormalizationIncremental::SpectralMeanNormalizationIncremental():
        NodeDef("SpectralMeanNormalizationIncremental")
    {
        *this
            << Terminal::in("INPUT")
            << Terminal::in("SM_HISTORY")
            << Terminal::in("NOT_EOF")
            << Terminal::in("FBANK_COUNT").def(13)
            << Terminal::in("PERIOD").def(20)
            << Terminal::in("SM_ALGORITHM").def("INCREMENTAL")
            << Terminal::in("SM_FILENAME").def("")
            << Terminal::in("SM_HISTORY_FILENAME").def("")
            << Terminal::in("IGNORE_FRAMES").def(false)
            << Terminal::in("BASENAME").def("smn_")
            << Terminal::in("OUTPUT_FNAME").def("out.txt")
            << Terminal::in("SM_EXPORT_FILENAME").def("")
            << Terminal::in("SM_EXPORT_ALGORITHM").def("LAST_SRC")
            >> Terminal::out("OUTPUT");
    }

    void
    SpectralMeanNormalizationIncremental::reset()
    {
        this->fbank_count = 0;
        this->period = 0;
        this->sm_algorithm.clear();
        this->infileSM.clear();
        this->b_sm_history_filename = false;
        this->historySM.clear();
        this->ignore_frames = false;
        this->prevSM.clear();
        this->curSM.clear();
        this->mapSM.clear();
        this->basename.clear();
        this->b_basename = false;
        this->streams.clear();
        this->output_fname.clear();
        this->b_output_fname = false;
        this->fout2.clear();
        this->b_sm_export = false;
        this->sm_export_filename.clear();
        this->sm_export_algorithm.clear();
        this->sm_export_nframe.clear();
        this->sm_export_sm.clear();
    }

    void
    SpectralMeanNormalizationIncremental::initialize(InputDataBundle &input)
    {
        fbank_count = input["FBANK_COUNT"].last().as<int>();
        period = input["PERIOD"].last().as<int>();
        sm_algorithm = input["SM_ALGORITHM"].last().cast<std::string>();
        if(sm_algorithm == "FILE"){
            std::string sm_filename = input["SM_FILENAME"].last().cast<std::string>();
            std::ifstream fin(sm_filename.c_str());
            if(fin.fail()){
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: Can't open '" << sm_filename << "'";
                throw RuntimeError(ss.str());
            }
            if(fin.eof()){
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: No_coef file defined in '" << sm_filename << "'";
                throw RuntimeError(ss.str());
            }

            std::string str;
            std::getline(fin, str);

            std::string token;
            std::istringstream stream(str);
            infileSM.resize(0);
            while(std::getline(stream, token, ',')){
                float val;
                std::stringstream ss(token);
                ss >> val;
                infileSM.push_back(val);
            }
            if(fbank_count != infileSM.size()){
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: Vector size of SM_FILENAME mismatches to FBANK_COUNT";
                throw RuntimeError(ss.str());
            }
        }

        b_sm_history_filename = false;
        if(false /*input.has("SM_HISTORY_FILENAME")*/){ // ToDo
            b_sm_history_filename = true;
            std::string sm_history_filename = input["SM_HISTORY_FILENAME"].last().cast<std::string>();
            std::ifstream fin(sm_history_filename.c_str());
            if(fin.fail()){
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: Can't open '" << sm_history_filename << "'";
                throw RuntimeError(ss.str());
            }
            if(fin.eof()){
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: No_coef file defined in '" << sm_history_filename << "'";
                throw RuntimeError(ss.str());
            }

            std::string str;
            historySM.resize(0);
            while(std::getline(fin, str)){
                std::string token;
                std::istringstream stream(str);
                std::vector<float>tmpV;
                tmpV.resize(0);
                while(std::getline(stream, token, ',')){
                    float val;
                    std::stringstream ss(token);
                    ss >> val;
                    tmpV.push_back(val);
                }
                historySM.push_back(tmpV);
            }
            for(int i = 0; i < historySM.size(); i++){
                if(fbank_count != historySM[i].size()){
                    std::stringstream ss;
                    ss << "SpectralMeanNormalizationIncremental: Vector size of SM_HISTORY_FILENAME mismatches to FBANK_COUNT";
                    throw RuntimeError(ss.str());
                }
            }
        }

        ignore_frames = input["IGNORE_FRAMES"].last().as<bool>();
        if(input.has("BASENAME")){
            basename = input["BASENAME"].last().cast<std::string>();
            b_basename = true;
        }else{
            b_basename = false;
        }

        if(false /*input.has("OUTPUT_FNAME")*/){ // ToDo
            output_fname = input["OUTPUT_FNAME"].last().cast<std::string>();
            fout2.open(output_fname.c_str());
            if(fout2.fail())
            {
                std::stringstream ss;
                ss << "SpectralMeanNormalizationIncremental: Can't open '" << output_fname << "'";
                throw RuntimeError(ss.str());
            }
            b_output_fname = true;
        }else{
            b_output_fname = false;
        }

        prevSM.resize(fbank_count);
        for(int i = 0; i < fbank_count; i++){
            prevSM[i] = 0.0f;
        }
        curSM.resize(fbank_count);

        b_sm_export = false;
        if(input.has("SM_EXPORT_FILENAME")){
            b_sm_export = true;
            sm_export_filename = input["SM_EXPORT_FILENAME"].last().cast<std::string>();
            sm_export_algorithm = input["SM_EXPORT_ALGORITHM"].last().cast<std::string>();
            sm_export_nframe.resize(0);
            sm_export_sm.resize(0);
        }

        input["INPUT"].set_look_window(0, 1);
    }

    bool
    SpectralMeanNormalizationIncremental::operator () (size_t count, InputDataBundle &input, OutputDataBundle &output)
    {
        if(count == 0){
            this->initialize(input);
        }
        if(!input["INPUT"].is_source_ready()){
            // std::stringstream ss;
            // ss << "SpectralMeanNormalizationIncremental: Input not ready" << std::endl;
            // std::cout << ss.str();
            return false;
        }

        data input1Data = input["INPUT"].last();
        data input2Data = input["INPUT"].relative(1);

        const std::map<int, data> &in1 = input1Data.cast<std::map<int, data>>();
        const std::map<int, data> &in2 = input2Data.cast<std::map<int, data>>();
        // RCPtr<Map<int, ObjectRef> > input1 = getInput(inputID, count);
        // const Map<int, ObjectRef>& in1 = *input1;
        //
        // RCPtr<Map<int, ObjectRef> > input2 = getInput(inputID, count + 1);
        // const Map<int, ObjectRef>& in2 = *input2;

        // {
        //     std::stringstream ss;
        //     ss << "SpectralMeanNormalizationIncremental: in1: count=" << in1.size() << ": ";
        //     for(auto it: in1){
        //         ss << it.first << ", ";
        //     }
        //     ss << std::endl;
        //     std::cout << ss.str();
        // }
        // {
        //     std::stringstream ss;
        //     ss << "SpectralMeanNormalizationIncremental: in2: count=" << in2.size() << ": ";
        //     for(auto it: in2){
        //         ss << it.first << ", ";
        //     }
        //     ss << std::endl;
        //     std::cout << ss.str();
        // }

        bool eofn = true;
        if(input.has("NOT_EOF")){
            eofn = input["NOT_EOF"].last().as<bool>();
        }
        // if(noteofID != -1) {
        //     ObjectRef peofn = getInput(noteofID, count+1);
        //     eofn = dereference_cast<bool>(peofn);
        // }

        std::map<int, data>smhistory_map;
        if(false /*input.has("SM_HISTORY")*/){ // ToDo
            smhistory_map = input["SM_HISTORY"].last().cast<std::map<int, data>>();
        }
        // RCPtr<Map<int, ObjectRef> > smhistory_map(NULL);
        // if(smhistoryID != -1) {
        //     smhistory_map = getInput(smhistoryID, count);
        // }

        data outputMapData = data::create<std::map<int, data>>();
        std::map<int, data> &outputMap = outputMapData.cast<std::map<int, data>>();
        // Map<int,ObjectRef>& output = *new Map<int,ObjectRef>;
        // out[count] = &output;

        for(auto it: in1){
            // for (Map<int, ObjectRef>::const_iterator it = in1.begin(); it != in1.end(); ++it) {

            std::vector<float> &fbank = it.second.cast<std::vector<float>>();
            // RCPtr<Vector<float> > fbank = it->second;
            if(fbank.size() < fbank_count){
                throw RuntimeError("Size of input feature vector is less than FBANK_COUNT.");
                // throw new NodeException(this, "Size of input feature vector is less than FBANK_COUNT.", __FILE__, __LINE__);
            }

            // Detect beggining of sources
            if(mapSM.find(it.first) == mapSM.end()){
                SpectralMeanMod tmpSM;
                tmpSM.Initialize(fbank_count);
                if(b_sm_history_filename){
                    for(int i = 0; i < historySM.size(); i++){
                        tmpSM.AddSpectrum(historySM[i]);
                    }
                }

                if(input.has("SM_HISTORY")){
                    // if(smhistoryID != -1) {
                    auto smhistory_it = smhistory_map.find(it.first);
                    if(smhistory_it != smhistory_map.end()){
                        std::vector<data> &tmpVV = smhistory_it->second.cast<std::vector<data>>();
                        for(auto v: tmpVV){
                            tmpSM.AddSpectrum(v.cast<std::vector<float>>());
                        }
                        // Map<int, ObjectRef>::const_iterator smhistory_it = smhistory_map->find(it->first);
                        // Vector<ObjectRef>& tmpVV = object_cast<Vector<ObjectRef> >(smhistory_it->second);
                        // for (int i = 0; i < tmpVV.size(); i++) {
                        //     RCPtr<Vector<float> > tmpV = tmpVV[i];
                        //     tmpSM.AddSpectrum(*tmpV);
                        // }
                    }
                }
                mapSM[it.first] = tmpSM;
            }

            // Load spectrum
            map<int, SpectralMeanMod>::iterator it_mapSM = mapSM.find(it.first);
            int source_id = it_mapSM->first;
            SpectralMeanMod& mod = it_mapSM->second;
            mod.AddDuration();

            // Add Spectrum
            if(mod.GetDuration() > ignore_frames){
                mod.AddSpectrum(fbank);
            }

            // Calculate SM
            if(period != 0 && mod.GetSpecCount() % period == 0 && mod.GetSpecCount() != 0){
                mod.CalculateMean();
            }

            // Decide SM
            // std::cout << "duration=" << mod.GetDuration() << std::endl;
            if(mod.GetDuration() < ignore_frames + period || period == 0){
                if(sm_algorithm == "INCREMENTAL"){
                    if(mod.GetDuration() <= ignore_frames){
                        for(int i = 0; i < fbank_count; i++){
                            curSM[i] = 0.0f;
                        }
                    }else{
                        mod.CalculateMean();
                        for(int i = 0; i < fbank_count; i++){
                            curSM[i] = mod.GetSpectralMean()[i];
                        }
                    }
                }else if(sm_algorithm == "PREV_SM"){
                    for(int i = 0; i < fbank_count; i++){
                        curSM[i] = prevSM[i];
                    }
                }else if(sm_algorithm == "ZERO"){
                    for(int i = 0; i < fbank_count; i++){
                        curSM[i] = 0.0f;
                    }
                }else if(sm_algorithm == "FILE"){
                    for(int i = 0; i < fbank_count; i++){
                        curSM[i] = infileSM[i];
                    }
                }
            }else{
                for(int i = 0; i < fbank_count; i++){
                    curSM[i] = mod.GetSpectralMean()[i];
                }
            }

            // Logging
            if(b_basename){
                shared_ptr<ofstream>fout;
                if(streams.find(it.first) == streams.end()){
                    std::ostringstream name_str;
                    name_str << basename << it.first << ".csv";
                    fout.reset(new ofstream(name_str.str().c_str()));
                    streams[it.first] = fout;
                }else{
                    fout = streams[it.first];
                }
                mod.CalculateMeanTmp();
                *fout << count << ", " << it.first << ", " << mod.GetDuration() << ", " << mod.GetSpecCount();
                for(int i = 0; i < fbank_count; i++){
                    *fout << ", " << mod.GetSpectralMeanTmp()[i];
                }
                for(int i = 0; i < fbank_count; i++){
                    *fout << ", " << curSM[i];
                }
                *fout << endl;
            }

            // Subtract SM
            data fbankoutData = data::create<std::vector<float>>(fbank.size());
            std::vector<float> &fbankout = fbankoutData.cast<std::vector<float>>();
            // RCPtr<Vector<float> > fbankout(Vector<float>::alloc(fbank.size()));
            for(int i = 0; i < fbank.size(); i++){
                if(i < fbank_count)
                    fbankout[i] = fbank[i] - curSM[i];
                else
                    fbankout[i] = fbank[i];
            }
            outputMap[it.first] = fbankoutData;
            // std::cout << "SpectralMeanNormalizationIncremental: source=" << it.first << std::endl;

            // Detect end of sources -> normalization and saving prevSM
            if(in2.find(it.first) == in2.end() || !eofn){
                mod.CalculateMean();
                if(b_output_fname){
                    fout2 << count << ", " << it.first << ", " << mod.GetDuration() << ", " << mod.GetSpecCount();
                    for(int i = 0; i < fbank_count; i++){
                        fout2 << ", " << mod.GetSpectralMean()[i];
                    }
                    fout2 << endl;
                }
                for(int i = 0; i < fbank_count; i++){
                    prevSM[i] = mod.GetSpectralMean()[i];
                }
                if(b_sm_export){
                    sm_export_nframe.push_back(mod.GetSpecCount());
                    sm_export_sm.push_back(prevSM);
                }
                mapSM.erase(it_mapSM);
                if(b_basename){
                    streams.erase(streams.find(source_id));
                }
            }

        }

        output["OUTPUT"].push(outputMapData);
        // std::cout << "SpectralMeanNormalizationIncremental: count=" << count << ", output=" << outputMap << std::endl;
        return true;
    }
}
}

/* end of file */
