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

import sys
import threading
import time

import numpy
import hark

import struct
#from scipy.io.wavfile import read
import wave

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)

        try:
            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"])

        except BaseException as ex:
            print('error: {}'.format(ex))

        return r

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

        try:
            node_extracting = network.create(HARK_Extracting, name="Extracting")
        except BaseException as ex:
            print(ex)

        try:
            #node_print = network.create(hark.node.Print)
            node_print = network.create(hark.node.Print, dispatch=hark.TriggeredMultiShotDispatcher)
            node_localize_music = network.create(hark.node.LocalizeMUSIC)
            node_cm_identity_matrix = network.create(hark.node.CMIdentityMatrix, dispatch=hark.RepeatDispatcher)
            node_constant__for_operation_flag = network.create(hark.node.Constant, dispatch=hark.RepeatDispatcher)
            node_multi_fft = network.create(hark.node.MultiFFT)

            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")
        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_audio_stream_from_memory
                    .add_input("INPUT", input["INPUT"])
                    .add_input("CHANNEL_COUNT", 8)
                ,
                node_cm_identity_matrix
                    .add_input("NB_CHANNELS", 8)
                    .add_input("LENGTH", 512)
                ,
                node_constant__for_operation_flag
                    .add_input("VALUE", True)
                ,
                node_multi_fft
                    .add_input("INPUT", node_audio_stream_from_memory["AUDIO"])
                    # .add_input("LENGTH", 512)
                    # .add_input("WINDOW", "CONJ")
                    # .add_input("WINDOW_LENGTH", 512)
                ,
                node_localize_music
                    .add_input("INPUT", node_multi_fft["OUTPUT"])
                    .add_input("NOISECM", node_cm_identity_matrix["OUTPUT"])
                    .add_input("OPERATION_FLAG", node_constant__for_operation_flag["OUTPUT"])
                    .add_input("A_MATRIX", "tf.zip")
                    #.add_input("MUSIC_ALGORITHM", "GEVD")
                    .add_input("MUSIC_ALGORITHM", "GSVD")
                    #.add_input("WINDOW_TYPE", "PAST")
                    #.add_input("WINDOW_TYPE", "MIDDLE")
                    .add_input("PERIOD", 50)
                    .add_input("DEBUG", True)
                ,
                node_print
                    #.add_input("INPUT", node_constant__for_operation_flag["OUTPUT"])
                    #.add_input("INPUT", node_cm_identity_matrix["OUTPUT"])
                    #.add_input("INPUT", node_audio_stream_from_memory["AUDIO"])
                    #.add_input("INPUT", node_multi_fft["OUTPUT"])
                    #.add_input("INPUT", node_localize_music["OUTPUT"])
                    #.add_input("INPUT", node_localize_music["SPECTRUM"])
                    #.add_input("INPUT", node_localize_music["SPECTRUM_BIN"])
                    #.add_input("INPUT", node_localize_music["TF_INFO"])
                    #.add_input("INPUT", node_localize_music["RXXS"])
                    #.add_input("INPUT", node_localize_music["RXXN"])
                ,
                node_vad_zc
                    .add_input("INPUT", node_audio_stream_from_memory["AUDIO"])
                    # .add_input("ADVANCE", 160)
                    # .add_input("LENGTH", 512)
                    .add_input("CHANNEL_COUNT", 8)
                    # .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")
                ,
            ]
        except BaseException as ex:
            print(ex)

        try:
            node_extracting.add_input("CHANNEL_WAV", node_vad_zc["OUTPUT"])
                # .add_input("LENGTH", 512)
                # .add_input("SAMPLING_RATE", 16000)
        except BaseException as ex:
            print(ex)

        try:
            #output.add_input("OUTPUT", node_spectral_mean_normalization_incremental["OUTPUT"])
            output.add_input("OUTPUT", node_localize_music["SPECTRUM"])
        except BaseException as ex:
            print(ex)

        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):

        try:
            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")
        except BaseException as ex:
            print(ex)

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

        return r

def received(data):
    pass
    #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:
            wf = wave.open('input.wav', mode='rb')
            rate = wf.getframerate()
            bytes = wf.getsampwidth()
            channels = wf.getnchannels()
            samples = wf.getnframes()
            wf.rewind()
            buf = wf.readframes(-1)
            wf.close()
            print("DEBUG:\n  rate={}, bytes={}, channels={}, samples={}, len(buf)={}\n".format(rate, bytes, channels, samples, len(buf)))

            # Read audio buffer to numpy array: convert to dtype='int32'
            enable_norm = False
            data = numpy.array([])
            if bytes == 1:
                data = numpy.frombuffer(buf, dtype=numpy.uint8).astype(numpy.int32)
                data -= 128
                if enable_norm:
                    data *= 2**24
            elif bytes == 2:
                data = numpy.frombuffer(buf, dtype=numpy.int16).astype(numpy.int32)
                if enable_norm:
                    data *= 2**16
            elif bytes == 3:
                #data = numpy.frombuffer(buf, dtype='int24')
                data = numpy.array([struct.unpack("<i", bytearray([0]) + buf[bytes * k:bytes * (k+1)])[0] for k in range(samples * channels)], dtype='int32')
                if enable_norm:
                    data *= 2**8
            elif bytes == 4:
                data = numpy.frombuffer(buf, dtype=numpy.int32)
            else:
                pass
            audio = numpy.reshape(data, (channels, samples))
            print("DEBUG:\n  audio.shape={}\n".format(audio.shape))

            # Generate audio frames with window shift: ADVANCE=160, LENGTH=512
            frames = numpy.lib.stride_tricks.as_strided(audio, strides=audio.strides+tuple(audio.strides[ax] for ax in tuple(range(audio.ndim))), shape=tuple(list(audio.shape))+(channels, 512), subok=False, writeable=False)[0, ::160]
            print("DEBUG:\n  frames.shape={}\n".format(frames.shape))

            for t in range(len(frames)):
                if not th.is_alive():
                    break
                print('<<<< send: count={}, max={}'.format(t, numpy.max(frames[t])))
                #print(frames[t].shape)
                #publisher.push(numpy.reshape(frames[t], frames[t].shape[0] * frames[t].shape[1]).astype(numpy.int16))
                publisher.push(frames[t].astype(numpy.int16))
                time.sleep(0.01)

            #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
