#include "hark-lib/node.h"
#include "hark-lib/network.h"
#include "hark-lib/input.h"

#include "hark-lib/node/Constant.h"
#include "hark-lib/node/InputStream.h"
#include "hark-lib/node/MultiFFT.h"
#include "hark-lib/node/FeatureRemover.h"
#include "hark-lib/node/Delta.h"
#include "hark-lib/node/MelFilterBank.h"
#include "hark-lib/node/MSLSExtraction.h"
#include "hark-lib/node/PreEmphasis.h"
#include "hark-lib/node/SpeechRecognitionClient.h"
#include "hark-lib/node/SpectralMeanNormalizationIncremental.h"
#include "hark-lib/node/VADZC.h"
#include "hark-lib/node/VADPyAnnote.h"
#include "hark-lib/node/SaveFeatures.h"
#include "hark-lib/node/SaveWavePCM.h"
#include "hark-lib/node/AudioStreamFromWave.h"

#if !defined(HARKLIB_USES_LIBTORCH)
#error This sample program requires hark-lib built with libtorch.  During hark-lib build, Use -DUSES_LIBTORCH=1 for cmake option
#endif

#include <stdlib.h>
#include <string.h>
#include <exception>
#include <iostream>
#include <fstream>

using namespace hark;

class HARK_Extracting: public hark::NetworkDef
{
public:
    std::map<std::string, std::string> _args;
public:
    HARK_Extracting(const std::map<std::string, std::string> &args):
        _args(args)
    {
    }

public:
    void
    build(hark::Network &network, hark::DataSourceMap &input, hark::DataSinkMap &output)
    {
        auto &node_multi_fft = network.create<hark::Node, hark::node::MultiFFT>();
        auto &node_feature_remover = network.create<hark::Node, hark::node::FeatureRemover>();
        auto &node_delta = network.create<hark::Node, hark::node::Delta>();
        auto &node_mel_filter_bank = network.create<hark::Node, hark::node::MelFilterBank>();
        auto &node_msls_extraction = network.create<hark::Node, hark::node::MSLSExtraction>();
        auto &node_pre_emphasis = network.create<hark::Node, hark::node::PreEmphasis>();

        node_multi_fft
            << hark::Input("INPUT", input["CHANNEL_WAV"]);
            // << hark::Input("LENGTH", 512)
            // << hark::Input("WINDOW", "CONJ")
            // << hark::Input("WINDOW_LENGTH", 512);

        node_feature_remover
            << hark::Input("INPUT", node_delta["OUTPUT"])
            << hark::Input("SELECTOR", "40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81");

        node_mel_filter_bank
            << hark::Input("INPUT", node_pre_emphasis["OUTPUT"])
            // << hark::Input("LENGTH", 512)
            // << hark::Input("SAMPLING_RATE", 16000)
            // << hark::Input("CUTOFF", 8000)
            // << hark::Input("MIN_FREQUENCY", 63)
            // << hark::Input("MAX_FREQUENCY", 8000)
            << hark::Input("FBANK_COUNT", 40);

        node_msls_extraction
            << hark::Input("FBANK", node_mel_filter_bank["OUTPUT"])
            << hark::Input("SPECTRUM", node_pre_emphasis["OUTPUT"])
            << hark::Input("FBANK_COUNT", 40)
            << hark::Input("NORMALIZATION_MODE", "SPECTRAL")
            // << hark::Input("USE_LEGACY_MODE", true)
            // << hark::Input("USE_HTK_LIFTER", false)
            // << hark::Input("LIFTERING_COEF", 22)
            << hark::Input("USE_POWER", true);

        node_pre_emphasis
            << hark::Input("INPUT", node_multi_fft["OUTPUT"])
            // << hark::Input("LENGTH", 512)
            // << hark::Input("SAMPLING_RATE", 16000)
            // << hark::Input("PREEMCOEF", 0.97f)
            << hark::Input("INPUT_TYPE", "SPECTRUM");

        node_delta
            << hark::Input("INPUT", node_msls_extraction["OUTPUT"]);
            // << hark::Input("FBANK_COUNT", 41);

        output["OUTPUT"] << node_feature_remover["OUTPUT"];
    }
};

class HARK_Recording: public hark::NetworkDef
{
public:
    std::map<std::string, std::string> _args;
    bool use_vad_pyannote;
    std::string output_path;
    float onset_threshold, offset_threshold;
    std::string norm_mean, norm_sigma, norm_count;
public:
    HARK_Recording(const std::map<std::string, std::string> &args):
        _args(args),
        use_vad_pyannote(false),
        output_path("."),
        onset_threshold(0.1),
        offset_threshold(-0.2),
        norm_mean("0 0 0 0"),
        norm_sigma("0 0 0 0"),
        norm_count("0 0 0 0")
    {
        std::map<std::string, std::string>::const_iterator it;
        if((it = args.find("use-vad-pyannote")) != args.end()){
            std::stringstream ss(it->second);
            ss >> this->use_vad_pyannote;
            std::cout << "use-vad-pyannote: " << this->use_vad_pyannote << std::endl;
        }
        if((it = args.find("output-path")) != args.end()){
            this->output_path = it->second;
        }
        if((it = args.find("onset-threshold")) != args.end()){
            std::stringstream(it->second) >> this->onset_threshold;
        }
        if((it = args.find("offset-threshold")) != args.end()){
            std::stringstream(it->second) >> this->offset_threshold;
        }
        if((it = args.find("norm-mean")) != args.end()){
            this->norm_mean = it->second;
            std::cout << "norm-mean: " << this->norm_mean << std::endl;
        }
        if((it = args.find("norm-sigma")) != args.end()){
            this->norm_sigma = it->second;
            std::cout << "norm-sigma: " << this->norm_sigma << std::endl;
        }
        if((it = args.find("norm-count")) != args.end()){
            this->norm_count = it->second;
            std::cout << "norm-count: " << this->norm_count << std::endl;
        }
    }

public:
    void
    build(hark::Network &network, hark::DataSourceMap &input, hark::DataSinkMap &output)
    {
        auto &node_extracting = network.create<hark::Network, HARK_Extracting>("Extracting", this->_args);
        auto &node_speech_recognition_client = network.create<hark::Node, hark::node::SpeechRecognitionClient>();
        auto &node_spectral_mean_normalization_incremental = network.create<hark::Node, hark::node::SpectralMeanNormalizationIncremental>();
        auto &node_vad = this->use_vad_pyannote
            ? network.create<hark::Node, hark::node::VADPyAnnote>()
            : network.create<hark::Node, hark::node::VADZC>();
        auto &node_save_features = network.create<hark::Node, hark::node::SaveFeatures>();
        auto &node_save_wave_pcm = network.create<hark::Node, hark::node::SaveWavePCM>();
        auto &node_audio_stream_from_wave = network.create<hark::Node, hark::node::AudioStreamFromWave, hark::dispatch::TriggeredRepeat>();

        node_audio_stream_from_wave
            << hark::LoopCondition(node_audio_stream_from_wave["NOT_EOF"])
            << hark::Input("INPUT", input["INPUT"]);
            // << hark::Input("LENGTH", 512)
            // << hark::Input("ADVANCE", 160)
            // << hark::Input("USE_WAIT", true);

        if(this->use_vad_pyannote){
            node_vad
                << hark::Input("INPUT", node_audio_stream_from_wave["AUDIO"])
                // << hark::Input("MODEL_PATH", "/home/kino/tmp/ipad/models/pyannote_jit_script.pt")
                << hark::Input("MODEL_PATH", "../../models/pyannote_jit_script.pt")
                << hark::Input("INITIAL_NORMALIZE_SAMPLE_COUNT", this->norm_count)
                << hark::Input("INITIAL_NORMALIZE_MEAN", this->norm_mean)
                << hark::Input("INITIAL_NORMALIZE_STDDEV", this->norm_sigma)
                << hark::Input("ONSET_SCORE_THRESOLD", this->onset_threshold)
                << hark::Input("OFFSET_SCORE_THRESOLD", this->offset_threshold)
                // << hark::Input("ADVANCE", 160)
                // << hark::Input("LENGTH", 512)
                << hark::Input("CHANNEL_COUNT", 2)
                // << hark::Input("SAMPLING_RATE", 16000)
                // << hark::Input("STRIP_ZERO", true)
                // << hark::Input("ZERO_MEAN", false)
                // << hark::Input("LEVEL_THRESHOLD", 60.f)
                // << hark::Input("ZEROCROSS_THRESHOLD", 80)
                << hark::Input("HEAD_MARGIN", 400.f)
                << hark::Input("TAIL_MARGIN", 600.f);
        }else{
            node_vad
                << hark::Input("INPUT", node_audio_stream_from_wave["AUDIO"])
                // << hark::Input("ADVANCE", 160)
                // << hark::Input("LENGTH", 512)
                << hark::Input("CHANNEL_COUNT", 2)
                    // << hark::Input("SAMPLING_RATE", 16000)
                // << hark::Input("STRIP_ZERO", true)
                // << hark::Input("ZERO_MEAN", false)
                << hark::Input("LEVEL_THRESHOLD", 60.f)
                << hark::Input("ZEROCROSS_THRESHOLD", 80)
                << hark::Input("HEAD_MARGIN", 400.f)
                << hark::Input("TAIL_MARGIN", 600.f);
        }

        node_extracting
            << hark::Input("CHANNEL_WAV", node_vad["OUTPUT"]);
        // << hark::Input("LENGTH", 512)
        // << hark::Input("SAMPLING_RATE", 16000);

        node_spectral_mean_normalization_incremental
            << hark::Input("INPUT", node_extracting["OUTPUT"])
            << hark::Input("NOT_EOF", true) // ToDo
            << hark::Input("SM_HISTORY", false) // ToDo
            // << hark::Input("FBANK_COUNT", 13)
            << hark::Input("PERIOD", 1)
            // << hark::Input("SM_ALGORITHM", "INCREMENTAL")
            // << hark::Input("SM_FILENAME", "")
            // << hark::Input("SM_HISTORY_FILENAME", "")
            // << hark::Input("IGNORE_FRAMES", false)
            << hark::Input("BASENAME", this->output_path + "/")
            << hark::Input("OUTPUT_FNAME", "");
            // << hark::Input("SM_EXPORT_FILENAME", "")
            // << hark::Input("SM_EXPORT_ALGORITHM", "LAST_SRC");

        node_speech_recognition_client
            << hark::Input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
            << hark::Input("MASKS", node_spectral_mean_normalization_incremental["OUTPUT"])
            << hark::Input("SOURCES", node_vad["SOURCES"])
            << hark::Input("MFM_ENABLED", false);
            // << hark::Input("HOST", "127.0.0.1")
            // << hark::Input("PORT", 5530)
            // << hark::Input("SOCKET_ENABLED", true);

        node_save_features
            << hark::Input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
            << hark::Input("SOURCES", node_vad["SOURCES"]);

        node_save_wave_pcm
            << hark::Input("ENABLE", false)
            << hark::Input("INPUT", node_vad["OUTPUT"])
            << hark::Input("BASENAME", "pcm_{srcid}_");
            // << hark::Input("ADVANCE", 160)
            // << hark::Input("SAMPLING_RATE", 16000)
            // << hark::Input("BITS", "int16")
            // << hark::Input("INPUT_BITS", "as_BITS");
    }
};

class HARK_Main: public hark::NetworkDef
{
public:
    std::map<std::string, std::string> _args;
public:
    HARK_Main(const std::map<std::string, std::string> &args): _args(args)
    {
    }

public:
    void
    build(hark::Network &network, hark::DataSourceMap &input, hark::DataSinkMap &output){

        auto &filename = network.create<hark::Node, hark::node::Constant, hark::dispatch::OneShot>()
                         << hark::Input("VALUE", "../input.wav");

        auto &stream = network.create<hark::Node, hark::node::InputStream>()
                       << hark::Input("INPUT", filename["OUTPUT"]);

        auto &loop = network.create<hark::Network, HARK_Recording, hark::dispatch::TriggeredRepeat>("HARK_Recording", this->_args)
                     << hark::Input("INPUT", stream["OUTPUT"]);
    }

};

int
process(int argc, char *argv[])
{
    std::map<std::string, std::string> args;
    for(int i = 1; i < argc; i++){
        // printf("%s\n", argv[i]);
        char *del = strchr(argv[i], '=');
        *del = '\0';
        std::string key = argv[i];
        std::string value = del + 1;
        // std::stringstream ss(argv[i]);
        // std::string key, value;
        // std::getline(ss, key, '=');
        // value = argv[i] + (key.size() + 1);
        // // std::getline(ss, value, '\0');
        // fprintf(stdout, "k=%s, v=%s\n", key.c_str(), value.c_str());
        args[key] = value;
    }
    // hark:data data(123.45);
    // std::cout << data.as<int>() << std::endl;
    // std::cout << data.as<float>() << std::endl;
    //
    std::shared_ptr<Network> network(hark::Network::create<HARK_Main>("HARK_Main1", args));
    network->execute();

    return 0;
}

int
main(int argc, char *argv[])
{
    try{
        process(argc, argv);
    }catch(std::exception &ex){
        std::cerr << "Error in process: " << ex.what() << std::endl;
        return EXIT_FAILURE;
    }catch(...){
        std::cerr << "Error in process: " << "not std::exception type" << std::endl;
        throw;
    }
    return EXIT_SUCCESS;
}
