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

import sys
import threading
import time

import numpy
import hark

import struct
import wave
import soundfile

import plotQuickWaveformKivy
import plotQuickSpecKivy
import plotQuickMusicSpecKivy
import plotQuickSourceKivy


def check_version_requirements(package, version):
    for t in [x for x in zip(package.__version__.split('.'), version.split('.'))]:
        if t[0] > t[1]:
            return True
        elif t[0] < t[1]:
            return False
    return True


class HARK_sub_localization(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_source_tracker = network.create(hark.node.SourceTracker)
            node_source_interval_extender = network.create(hark.node.SourceIntervalExtender)
            node_plotmusicspec_kivy = network.create(plotQuickMusicSpecKivy.plotQuickMusicSpecKivy)
            node_plotsource_kivy = network.create(plotQuickSourceKivy.plotQuickSourceKivy)

        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_cm_identity_matrix
                    .add_input("NB_CHANNELS", 8)
                    .add_input("LENGTH", 512)
                ,
                node_constant__for_operation_flag
                    .add_input("VALUE", True)
                ,
                node_localize_music
                    .add_input("INPUT", input["SPEC"])
                    .add_input("NOISECM", node_cm_identity_matrix["OUTPUT"])
                    .add_input("OPERATION_FLAG", node_constant__for_operation_flag["OUTPUT"])
                    .add_input("MUSIC_ALGORITHM", "SEVD")
                    #.add_input("TF_CHANNEL_SELECTION", [])
                    #.add_input("LENGTH", 512)
                    #.add_input("SAMPLING_RATE", 16000)
                    #.add_input("TF_INPUT_TYPE", "FILE")
                    #.add_input("A_MATRIX", "")
                    .add_input("A_MATRIX", "tamago_rectf.zip")
                    #.add_input("WINDOW", 50)
                    #.add_input("WINDOW_TYPE", "FUTURE")
                    .add_input("WINDOW_TYPE", "MIDDLE")
                    #.add_input("PERIOD", 50)
                    #.add_input("NUM_SOURCE", 2)
                    #.add_input("MIN_DEG", -180)
                    #.add_input("MAX_DEG", 180)
                    #.add_input("LOWER_BOUND_FREQUENCY", 500)
                    .add_input("LOWER_BOUND_FREQUENCY", 3000)
                    #.add_input("UPPER_BOUND_FREQUENCY", 2800)
                    .add_input("UPPER_BOUND_FREQUENCY", 6000)
                    #.add_input("SPECTRUM_WEIGHT_TYPE", "Uniform")
                    .add_input("SPECTRUM_WEIGHT_TYPE", "A_Characteristic")
                    #.add_input("A_CHAR_SCALING", 1.0)
                    #.add_input("ENABLE_EIGENVALUE_WEIGHT", True)
                    .add_input("ENABLE_EIGENVALUE_WEIGHT", False)
                    #.add_input("MAXIMUM_OUT_PEAKS", -1)
                    #.add_input("ENABLE_OUTPUT_SPECTRUM", False)
                    .add_input("ENABLE_OUTPUT_SPECTRUM", True)
                    #.add_input("ENABLE_OUTPUT_SPECTRUM_BIN", False)
                    #.add_input("ENABLE_OUTPUT_TF_INFO", False)
                    #.add_input("ENABLE_OUTPUT_RXXS", False)
                    #.add_input("ENABLE_OUTPUT_RXXN", False)
                    #.add_input("DEBUG", False)
                    .add_input("DEBUG", True)
                ,
                node_plotmusicspec_kivy
                    .add_input("SPECTRUM", node_localize_music["SPECTRUM"])
                ,
                node_source_tracker
                    .add_input("INPUT", node_localize_music["OUTPUT"])
                    #.add_input("THRESH", 20.0)
                    .add_input("THRESH", 25.0)
                    #.add_input("PAUSE_LENGTH", 800.0)
                    .add_input("PAUSE_LENGTH", 1200.0)
                    #.add_input("MIN_SRC_INTERVAL", 20.0)
                    #.add_input("MIN_ID", 0)
                    #.add_input("DEBUG", False)
                ,
                node_source_interval_extender
                    .add_input("SOURCES", node_source_tracker["OUTPUT"])
                    #.add_input("PREROLL_LENGTH", 50)
                    .add_input("PREROLL_LENGTH", 80)
                ,
                node_plotsource_kivy
                    .add_input("SOURCES", node_source_interval_extender["OUTPUT"])
                ,
            ]

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

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

        return r


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

        try:
            node_ghdss = network.create(hark.node.GHDSS)
            node_synthesize = network.create(hark.node.Synthesize)
            node_save_wave_pcm = network.create(hark.node.SaveWavePCM)
            node_power_calc_for_map = network.create(hark.node.PowerCalcForMap)
            #node_hrle = netwrok.create(hark.node.HRLE)
            node_estimate_leak = network.create(hark.node.EstimateLeak)
            #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)

        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_ghdss
                    .add_input("INPUT_FRAMES", input["SPEC"])
                    .add_input("INPUT_SOURCES", input["SOURCE"])
                    #.add_input("TRANSFER_FUNCTION", node_****["TRANSFER_FUNCTION"])
                    #.add_input("LENGTH", 512)
                    #.add_input("ADVANCE", 160)
                    #.add_input("SAMPLING_RATE", 16000)
                    #.add_input("LOWER_BOUND_FREQUENCY", 0)
                    #.add_input("UPPER_BOUND_FREQUENCY", 8000)
                    .add_input("TF_INPUT_TYPE", "FILE")
                    #.add_input("TF_CONJ_FILENAME", "")
                    .add_input("TF_CONJ_FILENAME", "tamago_rectf.zip")
                    #.add_input("INITW_FILENAME", "")
                    #.add_input("SS_METHOD", "ADAPTIVE")
                    ##.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")
                    .add_input("LC_CONST", "DIAG")
                    #.add_input("LC_METHOD", "ADAPTIVE")
                    ##.add_input("LC_METHOD", "FIX")
                    ##.add_input("LC_MYU", 0.001)
                    #.add_input("UPDATE_METHOD_TF_CONF", "POS")
                    #.add_input("UPDATE_METHOD_W", "ID_POS")
                    .add_input("UPDATE_METHOD_W", "ID")
                    #.add_input("UPDATE_ACCEPT_DISTANCE", 300.0)
                    ##.add_input("EXPORT_W", False)
                    ##.add_input("EXPORT_W_FILENAME", "")
                    #.add_input("UPDATE", "STEP")
                ,
                node_synthesize
                    .add_input("INPUT", node_ghdss["OUTPUT"])
                    #.add_input("LENGTH", 512)
                    #.add_input("ADVANCE", 160)
                    #.add_input("SAMPLING_RATE", 16000)
                    #.add_input("MIN_FREQUENCY", 125)
                    #.add_input("MAX_FREQUENCY", 7900)
                    #.add_input("WINDOW", "HAMMING")
                    #.add_input("OUTPUT_GAIN", 1.0)
                    .add_input("OUTPUT_GAIN", 0.0003125)
                ,
                node_save_wave_pcm
                    .add_input("INPUT", node_synthesize["OUTPUT"])
                    #.add_input("ENABLE", True)
                    #.add_input("BASENAME", "sep_")
                    .add_input("BASENAME", "ghdss_{srcid}_")
                    #.add_input("ADVANCE", 160)
                    #.add_input("SAMPLING_RATE", 16000)
                    #.add_input("BITS", "int16")
                    #.add_input("INPUT_BITS", "as_BITS")
                ,
                node_power_calc_for_map
                    .add_input("INPUT", node_ghdss["OUTPUT"])
                    #.add_input("POWER_TYPE", "POW")
                ,
                #node_hrle
                #    .add_input("INPUT_SPEC", node_power_calc_for_map["OUTPUT"])
                #    #.add_input("LX", 0.85)
                #    .add_input("LX", 0.10)
                #    #.add_input("TIME_CONSTANT_METHOD", "LEGACY")
                #    #.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("NUM_BIN", 2000)
                #    #.add_input("MIN_LEVEL", -100)
                #    .add_input("MIN_LEVEL", -200)
                #    #.add_input("STEP_LEVEL", 0.2)
                #    #.add_input("DEBUG", False)
                #,
                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)
                ,
                #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"])
                #,
                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("NOISE_SPEC", node_estimate_leak["LEAK_POWER_SPEC"])
                    #.add_input("ALPHA", 1.0)
                    #.add_input("BETA", 0.0)
                    #.add_input("SS_METHOD", 2)
                ,
                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"])
                ,
            ]

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

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

        return r


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

        try:
            node_pre_emphasis = network.create(hark.node.PreEmphasis)
            node_mel_filter_bank = network.create(hark.node.MelFilterBank)
            node_msls_extraction = network.create(hark.node.MSLSExtraction)
            node_delta = network.create(hark.node.Delta)
            node_feature_remover = network.create(hark.node.FeatureRemover)
            node_spectral_mean_normalization_incremental = network.create(hark.node.SpectralMeanNormalizationIncremental)
            node_speech_recognition_client = network.create(hark.node.SpeechRecognitionClient)
            node_save_features = network.create(hark.node.SaveHTKFeatures)

        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_pre_emphasis
                    .add_input("INPUT", input["SPEC"])
                    #.add_input("INPUT", node_spectral_gain_filter["OUTPUT_SPEC"])
                    #.add_input("LENGTH", 512)
                    #.add_input("SAMPLING_RATE", 16000)
                    #.add_input("PREEMCOEF", 0.97)
                    #.add_input("INPUT_TYPE", "WAV")
                    .add_input("INPUT_TYPE", "SPECTRUM")
                ,
                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", 13)
                    .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", 13)
                    .add_input("FBANK_COUNT", 40)
                    #.add_input("NORMALIZATION_MODE", "CEPSTRAL")
                    .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", False)
                    .add_input("USE_POWER", True)
                ,
                node_delta
                    .add_input("INPUT", node_msls_extraction["OUTPUT"])
                    #.add_input("FBANK_COUNT", 41)
                ,
                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("SM_HISTORY", False) # ToDo
                    .add_input("NOT_EOF", True) # ToDo
                    #.add_input("FBANK_COUNT", 13)
                    .add_input("FBANK_COUNT", 40)
                    #.add_input("PERIOD", 20)
                    .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", "smn_")
                    .add_input("BASENAME", "")
                    #.add_input("OUTPUT_FNAME", "out.txt")
                    .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", input["SOURCE"])
                    #.add_input("MFM_ENABLED", False)
                    #.add_input("HOST", "127.0.0.1")
                    .add_input("HOST", "localhost")
                    #.add_input("PORT", 5530)
                    .add_input("SOCKET_ENABLED", True)
                    #.add_input("SOCKET_ENABLED", False)
                    #.add_input("PROTOCOL_TYPE", 1)
                ,
                node_save_features
                    .add_input("FEATURES", node_spectral_mean_normalization_incremental["OUTPUT"])
                    .add_input("SOURCES", input["SOURCE"])
                    #.add_input("BASENAME", "")
                    .add_input("BASENAME", "feature_")
                    #.add_input("HTK_PERIOD", 100000)
                    #.add_input("FEATURE_TYPE", "USER")
                ,
            ]

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

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

        return r


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

        try:
            node_constant_input_filename = network.create(hark.node.Constant, dispatch=hark.OneShotDispatcher, name="Filename")
            node_input_stream = network.create(hark.node.InputStream)
            node_audio_stream_from_wave = network.create(hark.node.AudioStreamFromWave, dispatch=hark.TriggeredRepeatDispatcher, name="AudioStreamFromWave")
            node_multi_gain = network.create(hark.node.MultiGain)
            node_multi_fft = network.create(hark.node.MultiFFT)
            node_sub_localization = network.create(HARK_sub_localization, name="Localization")
            node_sub_separation = network.create(HARK_sub_separation, name="Separation")
            node_sub_recognition = network.create(HARK_sub_recognition, name="Recognition")

        except BaseException as ex:
            print(ex)

        try:
            r = [
                node_constant_input_filename
                    .add_input("VALUE", "2SPK-ja.wav")
                ,
                node_input_stream
                    .add_input("INPUT", node_constant_input_filename["OUTPUT"])
                ,
                node_audio_stream_from_wave
                    #.add_input("INPUT", input["INPUT"])
                    .add_input("INPUT", node_input_stream["OUTPUT"])
                    # LoopCondition is not implements yet.
                    #.loop_condition(node_audio_stream_from_wave["NOT_EOF"])
                ,
                #node_plotwaveform_kivy
                #    .add_input("WAVE", node_audio_stream_from_memory["AUDIO"])
                #,
                node_multi_gain
                    #.add_input("INPUT", node_audio_stream_from_memory["AUDIO"])
                    .add_input("INPUT", node_audio_stream_from_wave["AUDIO"])
                    #.add_input("GAIN", 0.0009765625 * 2**8)
                    #.add_input("GAIN", 0.0009765625)
                    .add_input("GAIN", 1.0)
                ,
                node_multi_fft
                    .add_input("INPUT", node_multi_gain["OUTPUT"])
                    # .add_input("LENGTH", 512)
                    # .add_input("WINDOW", "CONJ")
                    # .add_input("WINDOW_LENGTH", 512)
                ,
                node_sub_localization
                    .add_input("SPEC", node_multi_fft["OUTPUT"])
                ,
                node_sub_separation
                    .add_input("SPEC", node_multi_fft["OUTPUT"])
                    .add_input("SOURCE", node_sub_localization["OUTPUT"])
                ,
                node_sub_recognition
                    .add_input("SPEC", node_sub_separation["OUTPUT"])
                    .add_input("SOURCE", node_sub_localization["OUTPUT"])
                ,
            ]

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

        except BaseException as ex:
            print('error: {}'.format(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_subscriber = network.create(hark.node.SubscribeData, name="Subscriber")
            loop = network.create(HARK_MAIN_LOOP, name="HARK_Main_Loop")

        except BaseException as ex:
            print(ex)

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

        except BaseException as ex:
            print(ex)

        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

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

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

        try:
            th = threading.Thread(target=network.execute)
            th.start()
        except BaseException as ex:
            print(ex)
        except:
            network.stop()
        finally:
            if th.ident is not None:
                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
