//
//  NetworkDefs.swift
//  HARKlibSample
//
//  Created by Tomoyoshi Kinoshita on 2021/04/15.
//

import Foundation
import HARKlib

class Extracting: HARKNetworkDef
{
    override func build(_ network: HARKNetwork, input: HARKDataSourceMap, output: HARKDataSinkMap)
    {
        let node_multi_fft = network.create(HARKlib.MultiFFT.self, name: "MultiFFT")
        let node_feature_remover = network.create(HARKlib.FeatureRemover.self, name: "FeatureRemover")
        let node_delta = network.create(HARKlib.Delta.self, name: "Delta")
        let node_mel_filter_bank = network.create(HARKlib.MelFilterBank.self, name: "MelFilterBank")
        let node_msls_extraction = network.create(HARKlib.MSLSExtraction.self, name: "MSLSExtraction")
        let node_pre_emphasis = network.create(HARKlib.PreEmphasis.self, name: "PreEmphasis")

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

        node_feature_remover
            .add_input("INPUT", source: node_delta.get("OUTPUT"))
            .add_input("SELECTOR", string: "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")
//        << 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
            .add_input("INPUT", source: node_pre_emphasis.get("OUTPUT"))
            .add_input("FBANK_COUNT", number: 40)
//            << 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
            .add_input("FBANK", source: node_mel_filter_bank.get("OUTPUT"))
            .add_input("SPECTRUM", source: node_pre_emphasis.get("OUTPUT"))
            .add_input("FBANK_COUNT", number: 40)
            .add_input("NORMALIZATION_MODE", string: "SPECTRAL")
            .add_input("USE_POWER", boolean: true)
//            << 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
            .add_input("INPUT", source: node_multi_fft.get("OUTPUT"))
            .add_input("INPUT_TYPE", string: "SPECTRUM")
//            << 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
            .add_input("INPUT", source: node_msls_extraction.get("OUTPUT"))
//            << hark::Input("INPUT", node_msls_extraction["OUTPUT"])
//            // << hark::Input("FBANK_COUNT", 41)

        output
            .add_input("OUTPUT", source: node_feature_remover.get("OUTPUT"))
    }
}

class Recording: HARKNetworkDef
{
    override func build(_ network: HARKNetwork, input: HARKDataSourceMap, output: HARKDataSinkMap)
    {
        let use_vad_pyannote: Bool = false
        var vad_pyannote_model_path: String? = Bundle.main.path(forResource: "pyannote_jit_script", ofType: "pt")
        
        let node_extracting = network.create(Extracting.self, name: "Extracting")
        let node_speech_recognition_client = network.create(HARKlib.SpeechRecognitionClient.self, name: "SpeechRecognitionClient")
        let node_spectral_mean_normalization_incremental = network.create(HARKlib.SpectralMeanNormalizationIncremental.self, name: "SpectralMeanNormalizationIncremental")
        let node_save_features = network.create(HARKlib.SaveFeatures.self, name: "SaveFeatures")
        let node_save_wave_pcm2 = network.create(HARKlib.SaveWavePCM2.self, name: "SaveWavePCM2")
//        let node_audio_stream_from_wave = network.create(HARKlib.AudioStreamFromWave.self, dispatcher: HARKlib.TriggeredRepeatDispatcher.self)
        let node_audio_stream_from_memory = network.create(HARKlib.AudioStreamFromMemory.self, dispatcher: HARKlib.TriggeredMultiShotDispatcher.self, name: "AudioStreamFromMemory")

        let node_publisher = network.create(HARKlib.PublishData.self, dispatcher: HARKlib.RepeatDispatcher.self, name: "Publisher")
        let node_subscriber = network.create(HARKlib.SubscribeData.self, name: "Subscriber")
        let node_source_subscriber = network.create(HARKlib.SubscribeData.self, name: "Subscriber_Source")

        let node_vad = use_vad_pyannote
            ? network.create(HARKlib.VADPyAnnote.self, name: "VADPyAnnote")
            : network.create(HARKlib.VADZC.self, name: "VADZC")
        
        node_audio_stream_from_memory
            .add_input("INPUT", source: node_publisher.get("OUTPUT"))
////            << 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)

        node_vad
            .add_input("INPUT", source: node_audio_stream_from_memory.get("AUDIO"))
            .add_input("HEAD_MARGIN", number: 400.0)
            .add_input("TAIL_MARGIN", number: 600.0)

        if use_vad_pyannote {
            if let filePath = vad_pyannote_model_path {
                node_vad
                    .add_input("MODEL_PATH", string: filePath)
            }
        }else{
            node_vad
                .add_input("CHANNEL_COUNT", number: 1)
                .add_input("LEVEL_THRESHOLD", number: 60.0)
                .add_input("ZEROCROSS_THRESHOLD", number: 80)
//                .add_input("INITIAL_NORMALIZE_SAMPLE_COUNT", string: "")
//                .add_input("INITIAL_NORMALIZE_MEAN", string: "")
//                .add_input("INITIAL_NORMALIZE_STDDEV", string: "")
//            << 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_source_subscriber
            .add_input("INPUT", source: node_vad.get("SOURCES"))

        node_extracting
            .add_input("CHANNEL_WAV", source: node_vad.get("OUTPUT"))
//        // << hark::Input("LENGTH", 512)
//        // << hark::Input("SAMPLING_RATE", 16000)

        node_spectral_mean_normalization_incremental
            .add_input("INPUT", source: node_extracting.get("OUTPUT"))
            .add_input("NOT_EOF", boolean: true) // ToDo
            .add_input("SM_HISTORY", boolean: false) // ToDo
            .add_input("PERIOD", number: 1)
            .add_input("BASENAME", string: "")
            .add_input("OUTPUT_FNAME", string: "")
//            << 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", "")
//            << hark::Input("OUTPUT_FNAME", "")
//            // << hark::Input("SM_EXPORT_FILENAME", "")
//            // << hark::Input("SM_EXPORT_ALGORITHM", "LAST_SRC")

        node_speech_recognition_client
            .add_input("FEATURES", source: node_spectral_mean_normalization_incremental.get("OUTPUT"))
            .add_input("MASKS", source: node_spectral_mean_normalization_incremental.get("OUTPUT"))
            .add_input("SOURCES", source: node_vad.get("SOURCES"))
            .add_input("MFM_ENABLED", boolean: false)
//            << 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
            .add_input("FEATURES", source: node_spectral_mean_normalization_incremental.get("OUTPUT"))
            .add_input("SOURCES", source: node_vad.get("SOURCES"))
            .add_input("BASENAME", string: "feature_")
//            << hark::Input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
//            << hark::Input("SOURCES", node_vad["SOURCES"])
//            << hark::Input("BASENAME", "feature_")

        node_save_wave_pcm2
            .add_input("ENABLE", boolean: false)
            .add_input("INPUT", source: node_vad.get("OUTPUT"))
            .add_input("BASENAME", string: "pcm_{srcid}_")
//            << 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")
        
        node_subscriber
            .add_input("INPUT", source: node_spectral_mean_normalization_incremental.get("OUTPUT"))
    }
}

class Main: HARKNetworkDef
{
    override func build(_ network: HARKNetwork, input: HARKDataSourceMap, output: HARKDataSinkMap)
    {
        let loop = network.create(Recording.self, name: "Recording")
        
//        let filename =
//            network.create(HARKlib.Constant.self, dispatcher: HARKlib.OneShotDispatcher.self)
//            .add_input("VALUE", string: "input.wav")
//
//        let stream = network.create(HARKlib.InputStream.self)
//            .add_input("INPUT", source: filename.get("OUTPUT"))
//
//        let loop = network.create(Recording.self, name: "Recording")
//            .add_input("INPUT", source: stream.get("OUTPUT"))
    }

}

