#! /usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import threading
import time

import numpy
import hark

class MyNode(hark.node.PythonNode):
    def __init__(self):
        super().__init__('INPUT', 'OUTPUT')

    def forward(self, **kwargs):
        # print(kwargs)
        return {'OUTPUT': kwargs['INPUT']}

class HARK_Extracting(hark.NetworkDef):
    def build(self,
              network: hark.Network,
              input:   hark.DataSourceMap,
              output:  hark.DataSinkMap):

        node_multi_fft = network.create(hark.node.MultiFFT)
        node_feature_remover = network.create(hark.node.FeatureRemover)
        node_delta = network.create(hark.node.Delta)
        node_mel_filter_bank = network.create(hark.node.MelFilterBank)
        node_msls_extraction = network.create(hark.node.MSLSExtraction)
        node_pre_emphasis = network.create(hark.node.PreEmphasis)

        r = [
            node_multi_fft.
                add_input("INPUT", input["CHANNEL_WAV"])
                # .add_input("LENGTH", 512)
                # .add_input("WINDOW", "CONJ")
                # .add_input("WINDOW_LENGTH", 512)
            ,
            node_feature_remover
                .add_input("INPUT", node_delta["OUTPUT"])
                .add_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", node_pre_emphasis["OUTPUT"])
                # .add_input("LENGTH", 512)
                # .add_input("SAMPLING_RATE", 16000)
                # .add_input("CUTOFF", 8000)
                # .add_input("MIN_FREQUENCY", 63)
                # .add_input("MAX_FREQUENCY", 8000)
                .add_input("FBANK_COUNT", 40)
            ,
            node_msls_extraction
                .add_input("FBANK", node_mel_filter_bank["OUTPUT"])
                .add_input("SPECTRUM", node_pre_emphasis["OUTPUT"])
                .add_input("FBANK_COUNT", 40)
                .add_input("NORMALIZATION_MODE", "SPECTRAL")
                # .add_input("USE_LEGACY_MODE", true)
                # .add_input("USE_HTK_LIFTER", false)
                # .add_input("LIFTERING_COEF", 22)
                .add_input("USE_POWER", True)
            ,
            node_pre_emphasis
                .add_input("INPUT", node_multi_fft["OUTPUT"])
                # .add_input("LENGTH", 512)
                # .add_input("SAMPLING_RATE", 16000)
                # .add_input("PREEMCOEF", 0.97f)
                .add_input("INPUT_TYPE", "SPECTRUM")
            ,
            node_delta
                .add_input("INPUT", node_msls_extraction["OUTPUT"])
                # .add_input("FBANK_COUNT", 41)
            ,
        ]

        output.add_input("OUTPUT", node_feature_remover["OUTPUT"])

        return r

class HARK_Recording(hark.NetworkDef):
    def build(self,
              network: hark.Network,
              input:   hark.DataSourceMap,
              output:  hark.DataSinkMap):

        node_extracting = network.create(HARK_Extracting, name="Extracting")

        node_speech_recognition_client = network.create(hark.node.SpeechRecognitionClient)
        node_spectral_mean_normalization_incremental = network.create(hark.node.SpectralMeanNormalizationIncremental)
        node_vad_zc = network.create(hark.node.VADZC)
        node_save_features = network.create(hark.node.SaveFeatures)
        node_save_wave_pcm2 = network.create(hark.node.SaveWavePCM)
        node_audio_stream_from_memory = network.create(hark.node.AudioStreamFromMemory, dispatch=hark.TriggeredMultiShotDispatcher, name="AudioStreamFromMemory")
        node_pynode = network.create(MyNode)

        r = [
            node_audio_stream_from_memory
                .add_input("INPUT", input["INPUT"])
            ,
            node_pynode
                .add_input("INPUT", node_audio_stream_from_memory["AUDIO"])
            ,
            node_vad_zc
                .add_input("INPUT", node_pynode["OUTPUT"])
                # .add_input("ADVANCE", 160)
                # .add_input("LENGTH", 512)
                .add_input("CHANNEL_COUNT", 1)
                # .add_input("SAMPLING_RATE", 16000)
                # .add_input("STRIP_ZERO", true)
                # .add_input("ZERO_MEAN", false)
                .add_input("LEVEL_THRESHOLD", 60.)
                .add_input("ZEROCROSS_THRESHOLD", 80)
                .add_input("HEAD_MARGIN", 400.)
                .add_input("TAIL_MARGIN", 600.)
                .add_input("SPLIT_PERIOD", 500)
            ,
            node_spectral_mean_normalization_incremental
                .add_input("INPUT", node_extracting["OUTPUT"])
                .add_input("NOT_EOF", True) # ToDo
                .add_input("SM_HISTORY", False) # ToDo
                # .add_input("FBANK_COUNT", 13)
                .add_input("PERIOD", 1)
                # .add_input("SM_ALGORITHM", "INCREMENTAL")
                # .add_input("SM_FILENAME", "")
                # .add_input("SM_HISTORY_FILENAME", "")
                # .add_input("IGNORE_FRAMES", false)
                .add_input("BASENAME", "")
                .add_input("OUTPUT_FNAME", "")
                # .add_input("SM_EXPORT_FILENAME", "")
                # .add_input("SM_EXPORT_ALGORITHM", "LAST_SRC")
            ,
            node_speech_recognition_client
                .add_input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
                .add_input("MASKS", node_spectral_mean_normalization_incremental["OUTPUT"])
                .add_input("SOURCES", node_vad_zc["SOURCES"])
                .add_input("MFM_ENABLED", False)
                # .add_input("HOST", "127.0.0.1")
                # .add_input("PORT", 5530)
                # .add_input("SOCKET_ENABLED", true)
            ,
            node_save_features
                .add_input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
                .add_input("SOURCES", node_vad_zc["SOURCES"])
                .add_input("BASENAME", "feature_")
            ,
            node_save_wave_pcm2
                .add_input("ENABLE", False)
                .add_input("INPUT", node_vad_zc["OUTPUT"])
                .add_input("BASENAME", "pcm_{srcid}_")
                # .add_input("ADVANCE", 160)
                # .add_input("SAMPLING_RATE", 16000)
                # .add_input("BITS", "int16")
                # .add_input("INPUT_BITS", "as_BITS")
            ,
        ]

        node_extracting.add_input("CHANNEL_WAV", node_vad_zc["OUTPUT"])
            # .add_input("LENGTH", 512)
            # .add_input("SAMPLING_RATE", 16000)

        output.add_input("OUTPUT", node_spectral_mean_normalization_incremental["OUTPUT"])

        return r

class HARK_Main(hark.NetworkDef):
    def __init__(self):
        hark.NetworkDef.__init__(self)

    def build(self,
              network: hark.Network,
              input:   hark.DataSourceMap,
              output:  hark.DataSinkMap):

        node_publisher = network.create(hark.node.PublishData, dispatch=hark.RepeatDispatcher, name="Publisher")
        node_subscriber = network.create(hark.node.SubscribeData, name="Subscriber")
        loop = network.create(HARK_Recording, name="HARK_Recording")

        r = [
            loop
                .add_input("INPUT", node_publisher["OUTPUT"]),
            node_subscriber.add_input("INPUT", loop["OUTPUT"])
        ]

        return r

def received(data):
    print('>>>> received: {}'.format(data))

def main(args=sys.argv[1:]):
    if len(args) > 0:
        if args[0] == '--online':
            use_online = True
        elif args[0] == '--offline':
            use_online = False
        else:
            raise BaseException("Unexpected argument {}, use --online or --offline".format(args[0]))
    else:
        use_online = True # default mode
        # raise BaseException("No args given.  Use --online or --offline")

    if use_online:
        # online
        network = hark.Network.from_networkdef(HARK_Main, name="HARK_Main1")

        publisher = network.query_nodedef("Publisher")
        subscriber = network.query_nodedef("Subscriber")
        subscriber.receive = received

        try:
            def target():
                network.execute()

            th = threading.Thread(target=network.execute)
            th.start()
        except BaseException as ex:
            print(ex)

        try:
            for t in range(1000):
                if not th.is_alive():
                    break
                if (t // 200) % 2 == 1:
                    data = (4096 * numpy.random.randn(160)).astype(numpy.int16)
                else:
                    data = numpy.zeros(160, dtype=numpy.int16)
                print('<<<< send: count={}, max={}'.format(t, numpy.max(data)))
                publisher.push(data)
                time.sleep(0.01)
        finally:
            publisher.close()
            network.stop()
            th.join()

    else:
        raise NotImplementedError("Sorry, offline mode for network-def is not implemented yet.")
        # offline
        data = numpy.zeros(160000, dtype=numpy.int16)
        data[(numpy.arange(len(data)) // 200) % 2 == 1] = (4096 * numpy.random.randn(80000)).astype(numpy.int16)

        network = hark.Network.from_networkdef(HARK_Recording)
        output = network(INPUT=data)
        print(output.OUTPUT)

if __name__ == '__main__':
    main(sys.argv[1:])

# end of file
