#! /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

import plotQuickWaveformKivy
import plotQuickSpecKivy
import plotQuickMusicSpecKivy
import plotQuickSourceKivy


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

        try:
            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_localize_music = network.create(hark.node.LocalizeMUSIC)
            node_plotmusicspec_kivy = network.create(plotQuickMusicSpecKivy.plotQuickMusicSpecKivy)

            node_source_tracker = network.create(hark.node.SourceTracker)
            node_source_interval_extender = network.create(hark.node.SourceIntervalExtender)
            node_plotsource_kivy1 = network.create(plotQuickSourceKivy.plotQuickSourceKivy)
            node_plotsource_kivy2 = network.create(plotQuickSourceKivy.plotQuickSourceKivy)

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

        try:
            r = [
                ## Identitiy Matrix for the Correlation Matrix for each frequency bin.
                node_cm_identity_matrix
                    .add_input("NB_CHANNELS", 8)
                    .add_input("LENGTH", 512)
                ,
                ## Flag for whether to run the LocalizeMUSIC node.
                node_constant__for_operation_flag
                    .add_input("VALUE", True)
                ,
                ## Localize with the MUSIC algorithm.
                node_localize_music
                    .add_input("INPUT", input["SPECTRUM"])
                    .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", "SEVD")
                    #.add_input("MUSIC_ALGORITHM", "GEVD")
                    #.add_input("MUSIC_ALGORITHM", "GSVD")
                    #.add_input("WINDOW_TYPE", "PAST")
                    #.add_input("WINDOW_TYPE", "MIDDLE")
                    .add_input("LOWER_BOUND_FREQUENCY", 3000)
                    .add_input("UPPER_BOUND_FREQUENCY", 6000)
                    .add_input("WINDOW", 50)
                    .add_input("PERIOD", 1)
                    .add_input("NUM_SOURCE", 2)
                    .add_input("DEBUG", True)
                ,
                node_plotmusicspec_kivy
                    .add_input("SPECTRUM", node_localize_music["SPECTRUM"])
                ,
                ## Tracking source direction from power level.
                node_source_tracker
                    .add_input("INPUT", node_localize_music["OUTPUT"])
                    .add_input("THRESH", 22.0)
                    .add_input("PAUSE_LENGTH", 1200.0)
                    .add_input("MIN_SRC_INTERVAL", 20.0)
                    .add_input("MIN_ID", 0)
                    .add_input("DEBUG", True)
                ,
                node_plotsource_kivy1
                    .add_input("SOURCES", node_source_tracker["OUTPUT"])
                ,
                ## It is assumed that sound source existed for a certain period of time
                ## before the start time of the detected sound source. Equivalent to the
                ## operation of adding a short silent period until the start of speech.
                ## It has the effect of improving the recognition rate when the power is
                ## weak due to the starting phoneme of the utterance.
                node_source_interval_extender
                    .add_input("SOURCES", node_source_tracker["OUTPUT"])
                    #.add_input("PREROLL_LENGTH", 50)
                ,
                node_plotsource_kivy2
                    .add_input("SOURCES", node_source_interval_extender["OUTPUT"])
                ,
            ]

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

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

        return r


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

        try:
            #node_ghdss = network.create(hark.node.GHDSS)

            node_power_calc_for_map = network.create(hark.node.PowerCalcForMap)

            node_estimate_leak = network.create(hark.node.EstimateLeak)
            #node_hrle = netwrok.create(hark.node.HRLE)
            node_calc_spec_add_power = network.create(hark.node.CalcSpecAddPower)

            node_calc_spec_sub_gain = network.create(hark.node.CalcSpecSubGain)

            node_spectral_gain_filter = network.create(hark.node.SpectralGainFilter)

            node_white_noise_adder = network.create(hark.node.WhiteNoiseAdder)

            node_synthesize1 = network.create(hark.node.Synthesize)
            node_synthesize2 = network.create(hark.node.Synthesize)
            node_synthesize3 = network.create(hark.node.Synthesize)

            node_save_wave_pcm1 = network.create(hark.node.SaveWavePCM)
            node_save_wave_pcm2 = network.create(hark.node.SaveWavePCM)
            node_save_wave_pcm3 = network.create(hark.node.SaveWavePCM)

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

        try:
            r = [
                #node_ghdss
                #    .add_input("INPUT_FRAMES", input["SPECTRUM"])
                #    .add_input("INPUT_SOURCES", input["SOURCES"])
                #    #.add_input("TRANSFER_FUNCTION", node_****["TRANSFER_FUNCTION"])
                #    #.add_input("LENGTH", 512)
                #    #.add_input("ADVANCE", 160)
                #    #.add_input("SAMPLING_RATE", 16000)
                #    #.add_input("TF_INPUT_TYPE", "FILE")
                #    #.add_input("TF_INPUT_TYPE", "ONLINE")
                #    #.add_input("LOWER_BOUND_FREQUENCY", 0)
                #    #.add_input("UPPER_BOUND_FREQUENCY", 8000)
                #    .add_input("TF_CONJ_FILENAME", "tf.zip")
                #    #.add_input("INITW_FILENAME", "")
                #    #.add_input("SS_METHOD", "ADAPTIVE") ## default
                #    ##.add_input("SS_METHOD", "FIX")
                #    ##.add_input("SS_MYU", 0.001)
                #    #.add_input("SS_SCAL", 1.0)
                #    #.add_input("NOISE_FLOOR", 0.0)
                #    #.add_input("LC_CONST", "FULL") ## default
                #    #.add_input("LC_CONST", "DIAG")
                #    #.add_input("LC_METHOD", "ADAPTIVE") ## default
                #    ##.add_input("LC_METHOD", "FIX")
                #    ##.add_input("LC_MYU", 0.001)
                #    #.add_input("UPDATE_METHOD_TF_CONF", "POS") ## default
                #    #.add_input("UPDATE_METHOD_TF_CONF", "ID")
                #    #.add_input("UPDATE_METHOD_W", "ID") ## default
                #    #.add_input("UPDATE_METHOD_W", "POS")
                #    #.add_input("UPDATE_METHOD_W", "ID_POS")
                #    #.add_input("UPDATE_ACCEPT_DISTANCE", 300.0)
                #    ##.add_input("EXPORT_W", False)
                #    ##.add_input("EXPORT_W_FILENAME", "")
                #    #.add_input("UPDATE", "STEP") # default
                #    #.add_input("UPDATE", "TOTAL")
                #,
                ## Save wav file for GHDSS.
                #node_synthesize1
                #    .add_input("INPUT", node_ghdss["OUTPUT"])
                #,
                #node_save_wave_pcm1
                #    .add_input("INPUT", node_synthesize1["OUTPUT"])
                #    .add_input("ENABLE", False)
                #    .add_input("BASENAME", "pcm_ghdss_{srcid}_")
                #    # .add_input("ADVANCE", 160)
                #    # .add_input("SAMPLING_RATE", 16000)
                #    # .add_input("BITS", "int16")
                #    # .add_input("INPUT_BITS", "as_BITS")
                #,
                ## Compute the power spectrum.
                #node_power_calc_for_map
                #    .add_input("INPUT", node_ghdss["OUTPUT"])
                #    #.add_input("POWER_TYPE", "POW") ## default
                #    #.add_input("POWER_TYPE", "MAG")
                #,
                ## Estimate the leakage noise power.
                #node_estimate_leak
                #    .add_input("INPUT_POWER_SPEC", node_power_calc_for_map["OUTPUT"])
                #    #.add_input("LEAK_FACTOR", 0.25)
                #    #.add_input("OVER_CANCEL_FACTOR", 1)
                #,
                ## Estimate stationaly noise power with a histogram-based algorithm.
                #node_hrle
                #    .add_input("INPUT_SPEC", node_power_calc_for_map["OUTPUT"])
                #    #.add_input("LX", 0.85)
                #    #.add_input("TIME_CONSTANT_METHOD", "LEGACY") ## default
                #    #.add_input("TIME_CONSTANT", 0.75)
                #    #.add_input("TIME_CONSTANT_METHOD", "MILLISECOND")
                #    #.add_input("DECAY_FACTOR", 150)
                #    #.add_input("ADVANCE", 160)
                #    #.add_input("SAMPLING_RATE", 16000)
                #    #.add_input("NUM_BIN", 1000)
                #    #.add_input("MIN_LEVEL", -100)
                #    #.add_input("STEP_LEVEL", 0.2)
                #    #.add_input("DEBUG", False)
                #,
                ## Add multiple noise power spectra.
                ## Here the noise estimated by the EstimateLeak node and HRLE node are added.
                #node_calc_spec_add_power
                #    .add_input("INPUT_POWER_SPEC1", node_estimate_leak["LEAK_POWER_SPEC"])
                #    .add_input("INPUT_POWER_SPEC2", node_hrle["NOISE_SPEC"])
                #,
                ## Calculate the gain to subtract from the signal power spectrum and
                ## the estimated noise power spectrum.
                #node_calc_spec_sub_gain
                #    .add_input("INPUT_POWER_SPEC", node_power_calc_for_map["OUTPUT"])
                #    .add_input("NOISE_SPEC", node_calc_spec_add_power["OUTPUT_POWER_SPEC"])
                #    #.add_input("ALPHA", 1.0)
                #    #.add_input("BETA", 0.0)
                #    #.add_input("SS_METHOD", 2) ## default (amplitude substruction)
                #    #.add_input("SS_METHOD", 1) ## (power substruction)
                #,
                ## Apply a noise filter to the signal spectrum.
                #node_spectral_gain_filter
                #    .add_input("INPUT_SPEC", node_ghdss["OUTPUT"])
                #    .add_input("VOICE_PROB", node_calc_spec_sub_gain["VOICE_PROB"])
                #    .add_input("GAIN", node_calc_spec_sub_gain["GAIN"])
                #,
                #node_synthesize2
                #    .add_input("INPUT", node_spectral_gain_filter["OUTPUT_SPEC"])
                #,
                #node_save_wave_pcm2
                #    .add_input("INPUT", node_synthesize2["OUTPUT"])
                #    .add_input("ENABLE", False)
                #    .add_input("BASENAME", "pcm_nflt_{srcid}_")
                #    # .add_input("ADVANCE", 160)
                #    # .add_input("SAMPLING_RATE", 16000)
                #    # .add_input("BITS", "int16")
                #    # .add_input("INPUT_BITS", "as_BITS")
                #,
                ## Applying weak white noise may make the distortion less noticeable and
                ## improve the speech recognition rate. However, the effect is weak in
                ## today's deep learning models, and the learning models we provide
                ## do not use white noise addition.
                #node_white_noise_adder
                #    .add_input("INPUT", node_spectral_gain_filter["OUTPUT_SPEC"])
                #    #.add_input("LENGTH", 512)
                #    .add_input("WN_LEVEL", 0.0)
                #,
                #node_synthesize3
                #    .add_input("INPUT", node_white_noise_adder["OUTPUT"])
                #,
                #node_save_wave_pcm3
                #    .add_input("INPUT", node_synthesize3["OUTPUT"])
                #    .add_input("ENABLE", False)
                #    .add_input("BASENAME", "pcm_wn_{srcid}_")
                #    # .add_input("ADVANCE", 160)
                #    # .add_input("SAMPLING_RATE", 16000)
                #    # .add_input("BITS", "int16")
                #    # .add_input("INPUT_BITS", "as_BITS")
                #,
            ]

            #output.add_input("OUTPUT", node_ghdss["OUTPUT"])
            output.add_input("OUTPUT", node_spectral_gain_filter["OUTPUT_SPEC"])
            #output.add_input("OUTPUT", node_white_noise_adder["OUTPUT"])

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

        return r


class HARK_Extracting_Subnet(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 = [
                ## Executes high frequency enhancement (pre-emphasis).
                #node_pre_emphasis
                #    .add_input("INPUT", input["SEP_SPEC"])
                #    # .add_input("LENGTH", 512)
                #    # .add_input("SAMPLING_RATE", 16000)
                #    # .add_input("PREEMCOEF", 0.97)
                #    .add_input("INPUT_TYPE", "SPECTRUM")
                #,
                ## Performs mel fileter bank processing on the input spectrum to obtain
                ## the output energy for each filter channel.
                #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)
                #,
                ## Calculate Mel-Scale Log-Spectrum (MSLS) features and logarithmic power.
                #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)
                #,
                ## Calculate Delta features.
                #node_delta
                #    .add_input("INPUT", node_msls_extraction["OUTPUT"])
                #    # .add_input("FBANK_COUNT", 41)
                #,
                ## Remove features.
                #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_spectral_mean_normalization_incremental
                #    .add_input("INPUT", node_feature_remover["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_")
                #,
            ]

            output.add_input("OUTPUT", node_spectral_mean_normalization_incremental["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_localize_subnet = network.create(HARK_Localize_Subnet, name="Localize")
            node_separation_subnet = network.create(HARK_Separation_Subnet, name="Separation")
            node_extracting_subnet = network.create(HARK_Extracting_Subnet, name="Extracting")

        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_audio_stream_from_memory
                    .add_input("INPUT", input["INPUT"])
                    .add_input("CHANNEL_COUNT", 8)
                ,
                node_plotwaveform_kivy
                    .add_input("WAVE", node_audio_stream_from_memory["AUDIO"])
                ,
                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_plotspec_kivy
                    .add_input("FFT", node_multi_fft["OUTPUT"])
                ,
                node_localize_subnet
                    .add_input("FFT", node_multi_fft["OUTPUT"])
                ,
                node_separation_subnet
                    .add_input("SPECTRUM", node_multi_fft["OUTPUT"])
                    .add_input("SOURCES", node_localize_subnet["OUTPUT"])
                ,
                node_extracting_subnet
                    .add_input("SEP_SPEC", node_separation_subnet["OUTPUT"])
                    # .add_input("LENGTH", 512)
                    # .add_input("SAMPLING_RATE", 16000)
                ,
            ]

        except BaseException as ex:
            print(ex)

        try:
            #output.add_input("OUTPUT", node_multi_fft["OUTPUT"])
            #output.add_input("OUTPUT", node_localize_subnet["OUTPUT"])
            #output.add_input("OUTPUT", node_separation_subnet["OUTPUT"])
            output.add_input("OUTPUT", node_extracting_subnet["OUTPUT"])
        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_16bit.wav', mode='rb')
            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
