# coding: utf-8

# In[1]:

get_ipython().magic(u'matplotlib inline')
import IPython.display
from ipywidgets import interact, interactive, fixed
import numpy as np
import matplotlib.pyplot as plt
import copy
from scipy.io import wavfile
from scipy.signal import butter, lfilter
import scipy.ndimage
import socket
import os
import sys
import pandas as pd
import scipy.io as sio
import matplotlib
import scipy.signal as sg
import math
import scipy as sp
import pylab
import os
import wave
import scipy.io.wavfile
import h5py
import pickle
import tensorflow as tf
import random
import h5py
import collections
import bisect
import platform
import random

# In[2]:

#date and bird of data used
bird_id = 'z007'
num_clusters = 64
model_name = 'spectrogram prediction model/mel'

#locate source folder
hostname = socket.gethostname()

if hostname == 'txori' or hostname == 'passaro':
    data_folder = os.path.abspath('/mnt/cube/kai/data')
    results_folder = os.path.abspath('/mnt/cube/kai/results/'+model_name)
    repos_folder = os.path.abspath('/mnt/cube/kai/repositories')
    data_folder_zeke = os.path.abspath('/mnt/cube/earneodo/bci_zf/ss_data')
    bird_folder_save = os.path.join(data_folder, bird_id)
    bird_folder_data = os.path.join(data_folder_zeke, bird_id)
else:
    data_folder = os.path.abspath('/Users/Kai/Documents/UCSD/Research/Gentner group/Raw Data')
    results_folder = os.path.abspath('/Users/Kai/Documents/UCSD/Research/Gentner group/Results/'+model_name)
    repos_folder = os.path.abspath('/Users/Kai/Documents/UCSD/Research/Gentner group/Code')
    bird_folder_save = os.path.join(data_folder, bird_id)
    bird_folder_data = bird_folder_save

if not os.path.exists(bird_folder_save):
    os.makedirs(bird_folder_save)

if not os.path.exists(results_folder):
    os.makedirs(results_folder)

# In[3]:

# Most of the Spectrograms and Inversion are taken from: https://gist.github.com/kastnerkyle/179d6e9a88202ab0a2fe

def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a

def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y

def overlap(X, window_size, window_step):
    """
    Create an overlapped version of X
    Parameters
    ----------
    X : ndarray, shape=(n_samples,)
        Input signal to window and overlap
    window_size : int
        Size of windows to take
    window_step : int
        Step size between windows
    Returns
    -------
    X_strided : shape=(n_windows, window_size)
        2D array of overlapped X
    """
    if window_size % 2 != 0:
        raise ValueError("Window size must be even!")
    # Make sure there are an even number of windows before stridetricks
    append = np.zeros((window_size - len(X) % window_size))
    X = np.hstack((X, append))

    ws = window_size
    ss = window_step
    a = X

    valid = len(a) - ws
    nw = int((valid) // ss)
    out = np.ndarray((nw,ws),dtype = a.dtype)

    for i in xrange(nw):
        # "slide" the window along the samples
        start = i * ss
        stop = start + ws
        out[i] = a[start : stop]

    return out

def stft(X, fftsize=128, step=65, mean_normalize=True, real=False,
         compute_onesided=True):
    """
    Compute STFT for 1D real valued input X
    """
    if real:
        local_fft = np.fft.rfft
        cut = -1
    else:
        local_fft = np.fft.fft
        cut = None
    if compute_onesided:
        cut = fftsize // 2
    if mean_normalize:
        X -= X.mean()

    X = overlap(X, fftsize, step)

    size = fftsize
    win = 0.54 - .46 * np.cos(2 * np.pi * np.arange(size) / (size - 1))
    X = X * win[None]
    X = local_fft(X)[:, :cut]
    return X

def pretty_spectrogram(d,log = True, thresh= 5, fft_size = 512, step_size = 64):
    """
    creates a spectrogram
    log: take the log of the spectrgram
    thresh: threshold minimum power for log spectrogram
    """
    specgram = np.abs(stft(d, fftsize=fft_size, step=step_size, real=False,
        compute_onesided=True))

    if log == True:
        specgram /= specgram.max() # volume normalize to max 1
        specgram = np.log10(specgram) # take log
        specgram[specgram < -thresh] = -thresh # set anything less than the threshold as the threshold
    else:
        specgram[specgram < thresh] = thresh # set anything less than the threshold as the threshold

    return specgram

# Also mostly modified or taken from https://gist.github.com/kastnerkyle/179d6e9a88202ab0a2fe
def invert_pretty_spectrogram(X_s, log = True, fft_size = 512, step_size = 512/4, n_iter = 10):

    if log == True:
        X_s = np.power(10, X_s)

    X_s = np.concatenate([X_s, X_s[:, ::-1]], axis=1)
    X_t = iterate_invert_spectrogram(X_s, fft_size, step_size, n_iter=n_iter)
    return X_t

def iterate_invert_spectrogram(X_s, fftsize, step, n_iter=10, verbose=False):
    """
    Under MSR-LA License
    Based on MATLAB implementation from Spectrogram Inversion Toolbox
    References
    ----------
    D. Griffin and J. Lim. Signal estimation from modified
    short-time Fourier transform. IEEE Trans. Acoust. Speech
    Signal Process., 32(2):236-243, 1984.
    Malcolm Slaney, Daniel Naar and Richard F. Lyon. Auditory
    Model Inversion for Sound Separation. Proc. IEEE-ICASSP,
    Adelaide, 1994, II.77-80.
    Xinglei Zhu, G. Beauregard, L. Wyse. Real-Time Signal
    Estimation from Modified Short-Time Fourier Transform
    Magnitude Spectra. IEEE Transactions on Audio Speech and
    Language Processing, 08/2007.
    """
    reg = np.max(X_s) / 1E8
    X_best = copy.deepcopy(X_s)
    for i in range(n_iter):
        if verbose:
            print("Runnning iter %i" % i)
        if i == 0:
            X_t = invert_spectrogram(X_best, step, calculate_offset=True,
                                     set_zero_phase=True)
        else:
            # Calculate offset was False in the MATLAB version
            # but in mine it massively improves the result
            # Possible bug in my impl?
            X_t = invert_spectrogram(X_best, step, calculate_offset=True,
                                     set_zero_phase=False)
        est = stft(X_t, fftsize=fftsize, step=step, compute_onesided=False)
        phase = est / np.maximum(reg, np.abs(est))
        X_best = X_s * phase[:len(X_s)]
    X_t = invert_spectrogram(X_best, step, calculate_offset=True,
                             set_zero_phase=False)
    return np.real(X_t)

def invert_spectrogram(X_s, step, calculate_offset=True, set_zero_phase=True):
    """
    Under MSR-LA License
    Based on MATLAB implementation from Spectrogram Inversion Toolbox
    References
    ----------
    D. Griffin and J. Lim. Signal estimation from modified
    short-time Fourier transform. IEEE Trans. Acoust. Speech
    Signal Process., 32(2):236-243, 1984.
    Malcolm Slaney, Daniel Naar and Richard F. Lyon. Auditory
    Model Inversion for Sound Separation. Proc. IEEE-ICASSP,
    Adelaide, 1994, II.77-80.
    Xinglei Zhu, G. Beauregard, L. Wyse. Real-Time Signal
    Estimation from Modified Short-Time Fourier Transform
    Magnitude Spectra. IEEE Transactions on Audio Speech and
    Language Processing, 08/2007.
    """
    size = int(X_s.shape[1] // 2)
    wave = np.zeros((X_s.shape[0] * step + size))
    # Getting overflow warnings with 32 bit...
    wave = wave.astype('float64')
    total_windowing_sum = np.zeros((X_s.shape[0] * step + size))
    win = 0.54 - .46 * np.cos(2 * np.pi * np.arange(size) / (size - 1))

    est_start = int(size // 2) - 1
    est_end = est_start + size
    for i in range(X_s.shape[0]):
        wave_start = int(step * i)
        wave_end = wave_start + size
        if set_zero_phase:
            spectral_slice = X_s[i].real + 0j
        else:
            # already complex
            spectral_slice = X_s[i]

        # Don't need fftshift due to different impl.
        wave_est = np.real(np.fft.ifft(spectral_slice))[::-1]
        if calculate_offset and i > 0:
            offset_size = size - step
            if offset_size <= 0:
                print("WARNING: Large step size >50\% detected! "
                      "This code works best with high overlap - try "
                      "with 75% or greater")
                offset_size = step
            offset = xcorr_offset(wave[wave_start:wave_start + offset_size],
                                  wave_est[est_start:est_start + offset_size])
        else:
            offset = 0
        wave[wave_start:wave_end] += win * wave_est[
            est_start - offset:est_end - offset]
        total_windowing_sum[wave_start:wave_end] += win
    wave = np.real(wave) / (total_windowing_sum + 1E-6)
    return wave

def xcorr_offset(x1, x2):
    """
    Under MSR-LA License
    Based on MATLAB implementation from Spectrogram Inversion Toolbox
    References
    ----------
    D. Griffin and J. Lim. Signal estimation from modified
    short-time Fourier transform. IEEE Trans. Acoust. Speech
    Signal Process., 32(2):236-243, 1984.
    Malcolm Slaney, Daniel Naar and Richard F. Lyon. Auditory
    Model Inversion for Sound Separation. Proc. IEEE-ICASSP,
    Adelaide, 1994, II.77-80.
    Xinglei Zhu, G. Beauregard, L. Wyse. Real-Time Signal
    Estimation from Modified Short-Time Fourier Transform
    Magnitude Spectra. IEEE Transactions on Audio Speech and
    Language Processing, 08/2007.
    """
    x1 = x1 - x1.mean()
    x2 = x2 - x2.mean()
    frame_size = len(x2)
    half = frame_size // 2
    corrs = np.convolve(x1.astype('float32'), x2[::-1].astype('float32'))
    corrs[:half] = -1E30
    corrs[-half:] = -1E30
    offset = corrs.argmax() - len(x1)
    return offset

def make_mel(spectrogram, mel_filter, shorten_factor = 1):
    mel_spec =np.transpose(mel_filter).dot(np.transpose(spectrogram))
    mel_spec = scipy.ndimage.zoom(mel_spec.astype('float32'), [1, 1./shorten_factor]).astype('float16')
    mel_spec = mel_spec[:,1:-1] # a little hacky but seemingly needed for clipping
    return mel_spec

def mel_to_spectrogram(mel_spec, mel_inversion_filter, spec_thresh, shorten_factor):
    """
    takes in an mel spectrogram and returns a normal spectrogram for inversion
    """
    mel_spec = (mel_spec+spec_thresh)
    uncompressed_spec = np.transpose(np.transpose(mel_spec).dot(mel_inversion_filter))
    uncompressed_spec = scipy.ndimage.zoom(uncompressed_spec.astype('float32'), [1,shorten_factor]).astype('float16')
    uncompressed_spec = uncompressed_spec -4
    return uncompressed_spec

# From https://github.com/jameslyons/python_speech_features

def hz2mel(hz):
    """Convert a value in Hertz to Mels
    :param hz: a value in Hz. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Mels. If an array was passed in, an identical sized array is returned.
    """
    return 2595 * np.log10(1+hz/700.)

def mel2hz(mel):
    """Convert a value in Mels to Hertz
    :param mel: a value in Mels. This can also be a numpy array, conversion proceeds element-wise.
    :returns: a value in Hertz. If an array was passed in, an identical sized array is returned.
    """
    return 700*(10**(mel/2595.0)-1)

def get_filterbanks(nfilt=20,nfft=512,samplerate=30000,lowfreq=0,highfreq=None):
    """Compute a Mel-filterbank. The filters are stored in the rows, the columns correspond
    to fft bins. The filters are returned as an array of size nfilt * (nfft/2 + 1)
    :param nfilt: the number of filters in the filterbank, default 20.

    :param samplerate: the samplerate of the signal we are working with. Affects mel spacing.
    :param lowfreq: lowest band edge of mel filters, default 0 Hz
    :param highfreq: highest band edge of mel filters, default samplerate/2
    :returns: A numpy array of size nfilt * (nfft/2 + 1) containing filterbank. Each row holds 1 filter.
    """
    highfreq= highfreq or samplerate/2
    assert highfreq <= samplerate/2, "highfreq is greater than samplerate/2"

    # compute points evenly spaced in mels
    lowmel = hz2mel(lowfreq)
    highmel = hz2mel(highfreq)
    melpoints = np.linspace(lowmel,highmel,nfilt+2)
    # our points are in Hz, but we use fft bins, so we have to convert
    #  from Hz to fft bin number
    bin = np.floor((nfft+1)*mel2hz(melpoints)/samplerate)

    fbank = np.zeros([nfilt,nfft//2])
    for j in range(0,nfilt):
        for i in range(int(bin[j]), int(bin[j+1])):
            fbank[j,i] = (i - bin[j]) / (bin[j+1]-bin[j])
        if int(bin[j+1]) == int(bin[j+2]):
            fbank[j,int(bin[j+1])] = (int(bin[j+1]) - bin[j]) / (bin[j+1]-bin[j])
        else:
            for i in range(int(bin[j+1]), int(bin[j+2])):
                fbank[j,i] = (bin[j+2]-i) / (bin[j+2]-bin[j+1])
    return fbank

def create_mel_filter(fft_size, n_freq_components = 64, start_freq = 300, end_freq = 8000, samplerate=30000):
    """
    Creates a filter to convolve with the spectrogram to get out mels

    """
    mel_inversion_filter = get_filterbanks(nfilt=n_freq_components,
                                           nfft=fft_size, samplerate=samplerate,
                                           lowfreq=start_freq, highfreq=end_freq)
    # Normalize filter
    mel_filter = mel_inversion_filter.T / mel_inversion_filter.sum(axis=1)

    return mel_filter, mel_inversion_filter

# In[4]:

def make_datasets_spec(num_lookback, flatten=True, spec_file = 'mel_spectrogram.dat'):

    spec_file = os.path.join(bird_folder, spec_name)
    spec = np.loadtxt(spec_file).tolist()

    subdirs = [combined[0] for combined in os.walk(bird_folder)][1:]
    neural_datasets = list()
    all_datasets = list()

    for directory in subdirs:
        pickle_file = os.path.join(directory, pickle_name_mel)
        '''
        neuraldata = pickle.load(open(pickle_file, "rb" ))
        '''
        try:
            neuraldata = pickle.load(open(pickle_file, "rb" ))
        except IOError:
            continue

        num_bins = len(neuraldata[0])
        #print(neuraldata[0][0])
        for i in range(len(neuraldata)):
            for _ in range(num_lookback):
                neuraldata[i].insert(0, np.zeros(len(neuraldata[i][0])))

            for j in range(num_bins):
                if flatten:
                    neuralvector = np.concatenate(neuraldata[i][j:j+num_lookback])
                else:
                    neuralvector = np.array(neuraldata[i][j:j+num_lookback])
                neural_datasets.append(neuralvector)
                #for k in range(num_lookback):

                    #neuralvector.extend(neuraldata[i][j+k])

            zipped = zip(neural_datasets, spec)
            all_datasets = all_datasets+zipped
    return all_datasets

# In[5]:

def spike_sort(neural_kwik_name, song_kwik_name, song_length, bin_size, num_clusters, num_lookbacks, start_extend_bins=0,
               end_extend_bins=0, bird_folder_save=bird_folder_save, bird_folder_data=bird_folder_data, specify_subdir=None,
              n_mel_freq_components = 64, start_freq = 200, end_freq = 15000):
    '''
    neural_kwik_name: kwik file that contains neural data
    song_kwik_name: kwik file that contains motif starting time and recordings
    song_length: length of motif
    bin_size: how many datapoints to combine into a bin
    num_clusters: how many clusters of neurons are recorded
    start_extend_bins: at the beginning of the motif, how many bins do you want to shift. Shifting the start to the left
                        is negative, and shifting to the right is positive. Default to be 0
    end_extend_bins: same as start_extend_bins, how many bins to shift at the end of motif. Shifting left is negative and
                        right is positive. Default to be 0.

    '''
    save_name = 'mel_%02d_%d_%d.p' %(num_lookbacks, start_extend_bins, end_extend_bins)

    fft_size = bin_size*16 # window size for the FFT
    step_size = bin_size # distance to slide along the window (in time)
    start_extend_bins *= int(fft_size/bin_size)
    end_extend_bins *= int(fft_size/bin_size)
    spec_thresh = 5 # threshold for spectrograms (lower filters out more noise)
    lowcut = 500 # Hz # Low cut for our butter bandpass filter
    highcut = 15000 # Hz # High cut for our butter bandpass filter
    # For mels
    # number of mel frequency channels
    shorten_factor = 1 # how much should we compress the x-axis (time)
     # Hz # What frequency to start sampling our melS from
     # Hz # What frequency to stop sampling our melS from

    mel_filter, mel_inversion_filter = create_mel_filter(fft_size = fft_size,
                                                        n_freq_components = n_mel_freq_components,
                                                        start_freq = start_freq,
                                                        end_freq = end_freq)

    num_bins = int(song_length/bin_size-2-start_extend_bins+end_extend_bins)#number of bins in each song

    if specify_subdir:
        subdirs = [os.path.join(bird_folder_data, specify_subdir)]
    else:
        subdirs = [combined[0] for combined in os.walk(bird_folder_data)][1:]

    extended_song_length = song_length+bin_size*(end_extend_bins-start_extend_bins)
    neural_song_length = extended_song_length+(num_lookbacks-2)*bin_size
    #take out two bins due to mel processing algorithm

    mel_neural_comb = list()
    tot_motifs = 0

    for directory in subdirs:
        kwik_file = os.path.join(directory, neural_kwik_name)
        song_file = os.path.join(directory, song_kwik_name)

        try:
            trial1 = h5py.File(kwik_file,'r')['/channel_groups/0/spikes/time_samples'].value
        except IOError:
            print('--->Skipping '+directory+' because it doesnt contain spiking kwik files.')
            continue

        if "sleep" in directory:
            print('--->Skipping '+directory+' because it is a sleep recording.')
            continue

        recording = directory[directory.index('day-'):]
        recording_folder_save = os.path.join(bird_folder_save, recording)
        if not os.path.exists(recording_folder_save):
            os.makedirs(recording_folder_save)

        print('Start working on '+directory)

        with h5py.File(kwik_file,'r') as f:
            time_samples = f['/channel_groups/0/spikes/time_samples'].value
            cluster_num = f['/channel_groups/0/spikes/clusters/main'].value
            recording_num = f['/channel_groups/0/spikes/recording'].value
            if max(recording_num)>20:
                error_indices = [i for i,v in enumerate(list(recording_num)) if v > 20]
                for index in error_indices:
                    recording_num[index] = recording_num[index-1]
                print('***Erroneous indices fixed.')

            #print(type(time_samples))
            rec_counter = collections.Counter(recording_num)
            rec_list = rec_counter.keys()
            spike_counts = rec_counter.values()
            time_samples_each_rec = list()
            cluster_each_rec = list()

            #separate time samples based on recording
            for rec in rec_list:
                #start_samples.append(f['/recordings/'][str(rec)].attrs['start_sample'])
                prev_num_samples = sum(spike_counts[:rec])
                time_samples_each_rec.append(time_samples[prev_num_samples:prev_num_samples+spike_counts[rec]])
                #time_samples[prev_num_samples:prev_num_samples+spike_counts[rec]] = time_samples[
                #    prev_num_samples:prev_num_samples+spike_counts[rec]]+f['/recordings/'][str(rec)].attrs['start_sample']
                cluster_each_rec.append(cluster_num[prev_num_samples:prev_num_samples+spike_counts[rec]])

            #print(np.shape(time_samples_each_rec))
            #print(time_samples_each_rec)

        with h5py.File(song_file, 'r') as f:
            motif_starts = f['event_types/singing/motiff_1/time_samples'].value  #start time of each motif
            motif_recs = f['event_types/singing/motiff_1/recording'].value  #recording number of each motif
            songtuples = zip(motif_recs, motif_starts)

            tot_motifs += len(songtuples)

            for rec, ideal_start in songtuples:

                neural_ideal_start = ideal_start+bin_size*(start_extend_bins+1-num_lookbacks)
                neural_ideal_end = neural_ideal_start+neural_song_length

                song_ideal_start = ideal_start+bin_size*start_extend_bins
                song_ideal_end = song_ideal_start+extended_song_length

                #process wav file into mel spec
                wav_name = 'experiment-rec_0%02d.mic.wav' % (rec)
                wav_file = os.path.join(directory, wav_name)

                wav_rate, wav_data = wavfile.read(wav_file)
                wav_data = wav_data[song_ideal_start:song_ideal_end]
                wav_data = butter_bandpass_filter(wav_data, lowcut, highcut, wav_rate, order=1)

                if len(wav_data)!=extended_song_length:
                    raise ValueError('less wav datapointes loaded than expected')

                #plot each filtered unmel-edwav
                fig = plt.figure()
                plt.plot(wav_data)
                plt.title(str(rec)+': '+str(song_ideal_start))
                plt.savefig(os.path.join(recording_folder_save, 'motif '+str(rec)+' '+str(ideal_start)+'.png'))
                plt.close(fig)

                #mel and compile into a list based on time bins
                wav_spectrogram = pretty_spectrogram(wav_data.astype('float64'), fft_size = fft_size,
                                   step_size = step_size, log = True, thresh = spec_thresh)
                mel_list = list(make_mel(wav_spectrogram, mel_filter, shorten_factor = shorten_factor).transpose())

                current_t_list = time_samples_each_rec[rec] #extract neural activities for current motif
                current_cluster_list = cluster_each_rec[rec]  #extract cluster info for current motif

                neural_start_index = bisect.bisect_left(current_t_list, neural_ideal_start) #locate the index of starting point of the current motif
                neural_start_val = current_t_list[neural_start_index] #the actual starting point of the current motif
                #start_empty_bins = (neural_start_val-neural_ideal_start)//bin_size

                neural_end_index = bisect.bisect_left(current_t_list, neural_ideal_end)-1 #the end index of the current motif wrt indexing within the entire recording
                #end_empty_bins = (neural_ideal_end-current_t_list[neural_end_index])//bin_size

                t_during_song = current_t_list[neural_start_index:neural_end_index+1]-neural_ideal_start#normalize time sequence
                cluster_during_song = current_cluster_list[neural_start_index:neural_end_index+1] #extract cluster sequence

                counts_list = list() #contains all counts data for this motif
                for i in range(num_bins+num_lookbacks):
                    bin_start_index = bisect.bisect_left(t_during_song, i*bin_size)
                    bin_end_index = bisect.bisect_left(t_during_song,(i+1)*bin_size)
                    counts, bins = np.histogram(cluster_during_song[bin_start_index:bin_end_index], bins=np.arange(0,num_clusters+1))
                    counts_list.append(counts)

                for i in range(num_bins):
                    mel_neural_tuple = (counts_list[i:i+10], mel_list[i])
                    mel_neural_comb.append(mel_neural_tuple)

                """for cluster in range(num_clusters):
                    index_this_cluster = [i for i, j in enumerate(cluster_during_song) if j == cluster]
                    t_this_cluster = t_during_song[index_this_cluster]
                    counts,bins = np.histogram(t_this_cluster, bins=[start_val:bin_size:end_val, sys.maxint])
                    """

                if counts_list[0].shape[0] != num_clusters:
                    raise ValueError('check cluster histogram')

                if len(counts_list)-num_lookbacks != len(mel_list):
                    raise ValueError('Neural spiking bins do not correspond to spectrogram bins.')

                    #check for counting errors
        print('Finished.')

    if len(mel_neural_comb)!= num_bins*tot_motifs:
        raise ValueError('Might be missing some motifs or bins')

    pickle.dump(mel_neural_comb, open(os.path.join(bird_folder_save, save_name),'wb'))

    return mel_neural_comb

# In[6]:

def divide_parts(num_bins, divisor):
    num_bins_each_part_min = num_bins//divisor
    num_bins_each_part_max = num_bins_each_part_min + (num_bins%divisor!=0)

    num_bins_each_part = list([num_bins_each_part_max])*int(num_bins%divisor)+list([num_bins_each_part_min])*int(divisor-num_bins%divisor)

    return num_bins_each_part

def find_start_end_index(num_list):
    start_index = list()
    end_index = list()
    for i in range(len(num_list)):
        start_index.append(sum(num_list[:i]))
        end_index.append(sum(num_list[:i+1]))
    return start_index, end_index

# In[7]:

def plotspec(spec, vmin=-4, vmax=0):
    fig, ax = plt.subplots(nrows=1,ncols=1, figsize=(20,4))
    cax = ax.matshow(spec, interpolation='nearest', aspect='auto', cmap=plt.cm.afmhot, origin='lower',vmax=vmax, vmin=vmin)
    fig.colorbar(cax)
    #plt.title('Mel Spectrogram')

# In[8]:

def spec2wav(example_spec, mel_inversion_filter, spec_thresh, shorten_factor, fft_size, step_size):
    example_inverted_spectrogram = mel_to_spectrogram(example_spec, mel_inversion_filter,
                                                    spec_thresh=5,
                                                    shorten_factor=1)
    inverted_mel_audio = invert_pretty_spectrogram(np.transpose(example_inverted_spectrogram), fft_size = 2048,
                                                step_size = 128, log = True, n_iter = 10)
    #plt.plot(inverted_mel_audio)
    return inverted_mel_audio

# In[9]:

def sort_songs(test_output_compiled, test_spec_compiled, num_songs, num_bins, divisor, break_song=False, test_indeces=[0]):
    predicted_song_compiled = list()
    original_song_compiled = list()
    rmse = list()

    if break_song:
        test_indeces = assign_test_indeces(divisor, break_song=break_song)
        num_bins_per_part = divide_parts(num_bins, divisor)
        for song_index in range(num_songs):
            predicted_song = list()
            original_song = list()
            for test_index in test_indeces:
                num_bins_this_index = num_bins_per_part[test_index]
                current_predicted_seg = test_output_compiled[test_index][song_index*num_bins_this_index:(song_index+1)*num_bins_this_index]
                current_original_seg = test_spec_compiled[test_index][song_index*num_bins_this_index:(song_index+1)*num_bins_this_index]
                predicted_song += list(current_predicted_seg)
                original_song += list(current_original_seg)
            rmse.append(np.sqrt(np.mean(np.square(np.array(predicted_song)-np.array(original_song)))))
            predicted_song_compiled.append(predicted_song)
            original_song_compiled.append(original_song)
    else:
        num_songs_each_part = divide_parts(num_songs, divisor)
        if len(test_indeces)>1:
            start_index, end_index = find_start_end_index(num_songs_each_part)
            for test_index in test_indeces:
                for song_index in range(num_songs_each_part[test_index]):
                    #offset = start_index[test_index]*num_bins

                    predicted_song = test_output_compiled[test_index][song_index*num_bins:(song_index+1)*num_bins]
                    predicted_song_compiled.append(predicted_song)
                    original_song = test_spec_compiled[test_index][song_index*num_bins:(song_index+1)*num_bins]
                    original_song_compiled.append(original_song)
                    rmse.append(np.sqrt(np.mean(np.square(np.array(predicted_song)-np.array(original_song)))))
        else:
            test_index=test_indeces
            for song_index in range(num_songs_each_part[test_index[0]]):
                predicted_song = test_output_compiled[0][song_index*num_bins:(song_index+1)*num_bins]
                predicted_song_compiled.append(predicted_song)
                original_song = test_spec_compiled[0][song_index*num_bins:(song_index+1)*num_bins]
                original_song_compiled.append(original_song)
                rmse.append(np.sqrt(np.mean(np.square(np.array(predicted_song)-np.array(original_song)))))
    print([max(rmse),min(rmse)])

    return predicted_song_compiled, original_song_compiled, rmse

def assign_test_indeces(divisor, break_song=False):
    if break_song:
        test_indeces = range(divisor)
    else:
        test_indeces = [random.randint(0, divisor-1)]
    return test_indeces

# In[10]:

def eval_performance(x, y=None, mode = 'control_w', fft_size = 2048, step_size = 128, log = True, thresh = 5, shorten_factor = 1,
              n_mel_freq_components = 64, start_freq = 200, end_freq = 15000, matric = 'corr'):

    if not y and not mode:
        raise ValueError('Either give a y or specify a mode.')
    #if mode and mode is not 'control' and mode is not 'self':
    #    raise ValueError('Mode should be control or self.')
    if y and mode:
        mode = None
        print('Mode has been ignored due to the presence of y values.')

    if matric is not 'corr' and matric is not 'rmse' and matric is not 'both':
        raise ValueError('mode should be set to either corr or rmse or both.')

    coeff = list()
    rmse = list()
    mel_filter, mel_inversion_filter = create_mel_filter(fft_size = fft_size,
                                                        n_freq_components = n_mel_freq_components,
                                                        start_freq = start_freq,
                                                        end_freq = end_freq)
    if mode == 'self':
        xnew = list()
        ynew = list()
        for i in range(len(x)):
            for j in range(i+1, len(x)):
                xnew.append(x[i])
                ynew.append(x[j])
        x = xnew
        y = ynew

    for i in range(len(x)):
        #plt.plot(original)
        xi = np.array(x[i]).flatten()
        if not y:
            if mode == 'control_w':
                original_audio = spec2wav(np.array(x[i]).transpose(), mel_inversion_filter, thresh, shorten_factor, fft_size, step_size)
                control_audio = np.random.normal(0, np.std(original_audio), len(original_audio))
                #plt.plot(control_audio)
                yi = pretty_spectrogram(control_audio.astype('float64'), fft_size = fft_size, step_size = step_size, log = log,
                                                   thresh = thresh)
                yi = make_mel(yi, mel_filter, shorten_factor = shorten_factor).flatten()

            if mode == 'control_f':
                original_audio = spec2wav(np.array(x[i]).transpose(), mel_inversion_filter, thresh, shorten_factor, fft_size, step_size)
                control_audio = np.array([random.choice([1, -1])*wavbit for wavbit in original_audio])
                yi = pretty_spectrogram(control_audio.astype('float64'), fft_size = fft_size, step_size = step_size, log = log,
                                                   thresh = thresh)
                yi = make_mel(yi, mel_filter, shorten_factor = shorten_factor).flatten()

        else:
            yi = np.array(y[i]).flatten()
        if matric == 'corr':
            coeff.append(np.corrcoef([xi, yi])[0][1])
        if matric == 'rmse':
            rmse.append(np.sqrt(np.mean(np.square(np.array(xi)-np.array(yi)))))
        if matric == 'both':
            coeff.append(np.corrcoef([xi, yi])[0][1])
            rmse.append(np.sqrt(np.mean(np.square(np.array(xi)-np.array(yi)))))

    if matric == 'corr':
        return([round(np.mean(coeff, axis=0),4), round(np.std(coeff, axis=0),4)])
    if matric == 'rmse':
        return([round(np.mean(rmse, axis=0),4), round(np.std(rmse, axis=0),4)])
    if matric =='both':
        return([round(np.mean(rmse, axis=0),4), round(np.std(rmse, axis=0),4),
               round(np.mean(coeff, axis=0),4), round(np.std(coeff, axis=0),4)])

# In[11]:

def save_wav(waveform, rate, name, amp_rate=30000):
    waveform_adjusted = np.asarray(waveform/max(waveform)*amp_rate, dtype=np.int16)
    scipy.io.wavfile.write(name, rate, waveform_adjusted)

# In[ ]:


#Import all the modules here
import pandas as pd
import decisiontree as dt
import processdata as procd
import importcsv as ic
import bayesian
import nn
import preprocessing
import svm
import arm as arm
import pickle
import DBSCAN
from processResults import processResults, generateGraphs,generateGraphsSingle
from collections import defaultdict
from sklearn.model_selection import train_test_split
#Insert main code here

def reduceDimenTest(extramodelNames = ''):
	FILL_METHODS = ["mean","median","mode","none"]

	#space to search in where num_components refer to the dimensionality of the processed dataset
	NUM_COMPONENTS = [4,8,10]
	processedResults = defaultdict(lambda: [])

	for n_components in NUM_COMPONENTS:
		for filling in FILL_METHODS:
			print("**********************************now at ",filling, n_components)

			# in above function, fill_method has 'median', 'mode', and 'mean' options to fill data with the median, mode or mean
			data = ic.separateImport()
			data = procd.fillData(data, fill_method=filling)
			X_data, Y_data = preprocessing.createFullSet(data)
			X_data, pca, ss = preprocessing.performPCA(X_data, n_components)
			m = 3* X_data.shape[0] // 10
			trainX, testX, trainY, testY = train_test_split(X_data, Y_data, test_size=m, random_state=42, stratify=Y_data)

			#print(data)

			print("training set size: ", trainX.shape[0], " test set size: ", testX.shape[0] )
			#Sandbox your code here, before transfering it into your own python file
			predictions = []
			methods = []
			#min_sup = 0 #set it to be smth
			#associateRuleMiningPredictions = arm.generate_rules(min_sup)
			#print("Associate Rule Mining Predictions", associateRuleMiningPredictions)

			nnPredictions = nn.neuralNet(testX, testY, trainX, trainY, useTrainedModel = True,modelName =  filling + str(n_components) + extramodelNames)
			#print("nnPredictions",type(nnPredictions),nnPredictions)
			predictions.append(nnPredictions)
			methods.append("nnPredictions")

			bayesPredictions = bayesian.naiveBayes(testX, testY, trainX, trainY)
			#print("bayesPredictions",type(bayesPredictions),bayesPredictions)
			predictions.append(bayesPredictions)
			methods.append("bayesPredictions")

			#gs is the grid search model that i use to find the best parameters for the svm.
			#It automatically uses k-fold cross validation to find the best parameters
			#we can call print(gs.best_params_) to determine what params were used for this model
			svmPredictions, clf = svm.svmPredict(testX, testY, trainX, trainY, filling+str(n_components)+extramodelNames, gridSearch=False)
			#print("SVMpredictions", type(svmPredictions), svmPredictions)
			predictions.append(svmPredictions)
			methods.append("SVMpredictions")

			#best hyperparams precalculated using grid search model to save time
			#res, best_params = dt.gridSearchWrapper(testX, testY, trainX, trainY)
			best_params = {'n_estimators': 10, 'max_depth': 6, 'min_samples_split': 14}
			randforestPred = dt.randomForestClassify(testX, testY, trainX, trainY, best_params)
			predictions.append(randforestPred)
			methods.append("randforest")
			# methods.append("Random forest")

			#ensemble method using a simple majority vote of all the classifiers.
			ensemblePred = []

			for result in zip(*[item.tolist() for item in predictions]):
				ensemblePred.append(max(set(result), key=result.count))

			predictions.append(ensemblePred)
			methods.append("Ensemble")
			#print("ensemblePred", ensemblePred)

			for prediction, labels in zip(predictions,methods):
				result = processResults(prediction,testY,filling,labels)
				result["n_components"] = n_components
				result["filling"] = filling
				processedResults[labels].append(result)
	generateGraphsSingle(processedResults,FILL_METHODS)
	generateGraphs(processedResults,FILL_METHODS)

def generate_association_rules(minsup=0.5, minconf=0.7):
	# print("(Itemset, Count)")
	# arm.print_frequent_itemsets(minsup)
	print("Rules generated from ARM with {} min support and {} min confidence:".format(minsup, minconf))
	arm.print_rules(minsup, minconf)

if __name__ == '__main__':
	# reduceDimenTest('tenDim')
	DBSCAN.perform_dbscan()


    # PyBrickBreaker
   # Copyright 2013 Meredith Myers

   # Licensed under the Apache License, Version 2.0 (the "License");
   # you may not use this file except in compliance with the License.
   # You may obtain a copy of the License at

   #     http://www.apache.org/licenses/LICENSE-2.0

   # Unless required by applicable law or agreed to in writing, software
   # distributed under the License is distributed on an "AS IS" BASIS,
   # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   # See the License for the specific language governing permissions and
   # limitations under the License.

from Tkinter import *

class blockData(object):
    pass
block = blockData()

class ballClass(object):
    pass
ball = ballClass()

# Handles paddle movement
def mouseMoved(event):
    global paddleID, canvas

    xV = event.x
    curBBox = canvas.bbox(paddleID)

    px = curBBox[0]
    p2x = curBBox[2]

    X = xV - px
    newPX = -X+xV

    if (xV >= 10) and (xV <= 450):
        canvas.move(paddleID, xV - px, 0)
    else:
        pass

# Begins the game
def startBall(event):
    global canvas, string2, lbl2

    string2.set("")
    canvas.after(50, moveBall)

def moveBall():
    global ballID, canvas, paddleID, blockID, leftWall, topWall, rightWall, blocks, score, string, string2

    # Get the current bounding box
    boundingBox = canvas.bbox(ballID)
    left = boundingBox[0]
    top = boundingBox[1]
    right = boundingBox[2]
    bottom = boundingBox[3]

    #check for collisions and determine what the ball has collided with
    over = canvas.find_overlapping(left,top,right,bottom)
    z = len(over)

    if z == 1:
        obj = over[0]
    else:
        for i in range(z):
            if over[i] == ballID:
                pass
            else:
                obj = over[i]

    # If the collision is with a block
    if obj in blocks:
        dx = -0.1
        dy = -0.1
        score += 1
        string.set("Score: "+(str(score)))

        # While the collision is still being detected, move back
        newx = left
        newy = top
        newx2 = right
        newy2 = bottom

        while obj in over:

            # get new x,y coordinate
            newx = newx - dx
            newy = newy - dy
            newx2 = newx2 - dx
            newy2 = newy2 - dy
            over = canvas.find_overlapping(newx,newy,newx2,newy2)

        # When the collision is no longer detected, look at the ball and block's bbox
        bboxBlock = canvas.bbox(obj)
        xBlk = bboxBlock[0]
        yBlk = bboxBlock[1]
        x2Blk = bboxBlock[2]
        y2Blk = bboxBlock[3]
        # Compute the overlaps
        y = y2Blk - newy

        if ball.dx < 0:
            x = x2Blk - newx

        else:
            x = newx2 - xBlk

        # Flip the appropriate speeds
        if x > y:
            #print("if",x,y)
            ball.dy = -ball.dy
            blocks.remove(obj)
            canvas.delete(obj)
            canvas.move(ballID, ball.dx, ball.dy)
        elif y > x:
            #print("elif",x,y)
            ball.dx = -ball.dx
            blocks.remove(obj)
            canvas.delete(obj)
            canvas.move(ballID, ball.dx, ball.dy)
        else:
            #print("else",x,y)
            ball.dy = -ball.dy
            ball.dx = -ball.dx
            blocks.remove(obj)
            canvas.delete(obj)
            canvas.move(ballID, ball.dx, ball.dy)

    # Handle paddle bounces
    elif obj == paddleID:
        paddleBox = canvas.bbox(paddleID)
        paddleLeft = paddleBox[0]
        paddleRight = paddleBox[2]
        paddleMiddle = (paddleLeft + paddleRight)/2 # detect middle of the paddle
        ball.dx += 0.6*(left - paddleMiddle)
        ball.dy = -ball.dy# Make dy negative so after hitting the paddle the ball is going up.
        canvas.move(ballID, ball.dx, ball.dy)

    # Handle wall bounces
    elif (obj == leftWall) or (obj == rightWall):
        ball.dx = -ball.dx
        canvas.move(ballID, ball.dx, ball.dy)

    elif obj == topWall:
        ball.dx = -ball.dx
        ball.dy = -ball.dy
        canvas.move(ballID, ball.dx, ball.dy)

    # Handle free motion
    else:
        canvas.move(ballID, ball.dx, ball.dy)

    #Check if game is over
    if bottom > 530:
        string2.set("GAME OVER!")
    else:
        canvas.after(50, moveBall)

def main():
    global paddleID, canvas, ballID, blockID, leftWall, topWall, rightWall, blocks, score, string, string2, lbl2

    root = Tk()
    root.grid()
    root.title("Py Brick Breaker")

    canvas = Canvas(root,width=500,height=525,bg = "black")
    canvas.grid(row=0,column=0)

    #Create walls
    leftWall = canvas.create_rectangle([-100,0,10,525], fill = "white")
    rightWall = canvas.create_rectangle([490,0,600,525], fill = "white")
    topWall = canvas.create_rectangle([10,-100,500,10], fill = "white")

    #Create paddle
    paddleID = canvas.create_rectangle([50,490,90,600],fill="green")
    canvas.bind("<B1-Motion>",mouseMoved)

    #Create ball
    ballID = canvas.create_oval([242,282,258,298],fill="green")
    bbox = canvas.bbox(ballID)
    ball.x = bbox[0]
    ball.y = bbox[1]
    ball.x2 = bbox[2]
    ball.y2 = bbox[3]
    ball.dx = 0
    ball.dy = 5

    #Create changable game start/over label
    string2 = StringVar()
    string2.set("Press Enter to Begin")
    lbl2 = Label(root,textvariable = string2, fg = "white", bg = "black")
    lbl2.grid(row = 0, column = 0)
    root.bind("<Return>",startBall)

    #Create blocks
    blocks = []
    block.xLoc = 50
    block.yLoc = 50
    block.xsize = 30
    block.ysize = 20
    for x in range(4):
        for y in range(13):

            blockID = canvas.create_rectangle([block.xLoc,block.yLoc,
                                       block.xLoc+block.xsize,
                                       block.yLoc+block.ysize], fill="green")
            canvas.grid()
            blocks.append(blockID)
            block.xLoc += 30
        block.yLoc += 20
        block.xLoc = 50

    #Create changable score label
    score = 0
    string = StringVar()
    string.set("Score: "+(str(score)))
    lbl = Label(root,textvariable = string, fg = "white", bg = "black")
    lbl.grid(row = 0, column = 0, sticky = S+E+W)

    root.mainloop()

main()


import requests
import json

from Get_auth_tokern import get_token

# Get the device list from DNAC

def main():

    #url = "https://sandboxdnac.cisco.com/dna/intent/api/v1/network-device/"

    token = get_token()
    api_path= "https://sandboxdnac.cisco.com/dna"
    #payload={}
    headers = {
    'Content-Type': 'application/json',
    'X-Auth-Token': token
    }

    response = requests.get( f"{api_path}/intent/api/v1/network-device", headers=headers)

    #print(response.text)
    #ip_list = json.dumps(response.json(), indent=2)

    #
    #print(ip_list)

    if response.ok:
        for device in response.json()["response"]:
            print(f"ID:{device['id']}   IP: {device['managementIpAddress']}"  )

    else:
        print (f"Device collection fails with code {response.status_code}")
        print (f"Failure body: {response.text}")

if __name__ == "__main__":
    main()


import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import time
import copy

class CNNEarDetector:

    def __init__(self, model, num_classes, extract_features=True):
        if model == "resnet":
            self.model = torchvision.models.resnet50(pretrained=True)
            if extract_features:
                for param in model.parameters():
                    param.requires_grad = False
            nfeatures = self.model.fc.in_features
            self.model.fc = nn.Linear(nfeatures, num_classes)
            self.input_size = 224
        elif model == "alexnet":
            self.model = torchvision.models.alexnet(pretrained=True)
            if extract_features:
                for param in model.parameters():
                    param.requires_grad = False
            nfeatures = self.model.classifier[6].in_features
            self.model.classifier[6] = nn.Linear(nfeatures, num_classes)
            self.input_size = 224

        self.extract_features = extract_features

        self.optimizer = None

    def set_otpimizer(self, learning_rate=0.001, weight_decay=0):
        # construct an optimizer
        if self.extract_features:
            params = [p for p in self.model.parameters() if p.requires_grad]
        else:
            params = self.model.parameters()

        self.optimizer = torch.optim.SGD(params, lr=learning_rate, weight_decay=weight_decay)

    def train(self, dataloader, criterion, num_epochs=50, learning_rate=0.001, weight_decay=0, since=time.time()):
        # train on the GPU or on the CPU, if a GPU is not available
        device = None
        if torch.cuda.is_available():
            print("Training model on GPUs")
            device = torch.device('cuda')
        else:
            print("Training model on CPUs")
            device = torch.device('cpu')

        self.model.to(device)

        # Initialize optimizer if it doesn't already exist
        if self.optimizer is None:
            self.set_otpimizer(learning_rate, weight_decay)

        val_acc_history = []

        best_model_wts = copy.deepcopy(self.model.state_dict())
        best_acc = 0.0

        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('-' * 10)

            # Each epoch has a training and validation phase
            for phase in ['train', 'val']:
                if phase == 'train':
                    self.model.train()  # Set model to training mode
                else:
                    self.model.eval()  # Set model to evaluate mode

                running_loss = 0.0
                running_corrects = 0

                # Iterate over data.
                for inputs, labels in dataloader[phase]:
                    inputs = inputs.to(device)
                    labels = labels.to(device)

                    #print(labels)

                    # zero the parameter gradients
                    self.optimizer.zero_grad()

                    # forward
                    # track history if only in train
                    with torch.set_grad_enabled(phase == 'train'):
                        # Get model outputs and calculate loss
                        outputs = self.model(inputs)
                        loss = criterion(outputs, labels)

                        _, preds = torch.max(outputs, 1)

                        # backward + optimize only if in training phase
                        if phase == 'train':
                            loss.backward()
                            self.optimizer.step()

                    # statistics
                    running_loss += loss.item() * inputs.size(0)
                    running_corrects += torch.sum(preds == labels.data)

                epoch_loss = running_loss / len(dataloader[phase].dataset)
                epoch_acc = running_corrects.double() / len(dataloader[phase].dataset)

                print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

                # deep copy the model
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model_wts = copy.deepcopy(self.model.state_dict())
                if phase == 'val':
                    val_acc_history.append(epoch_acc)

            print()

        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print('Best val Acc: {:4f}'.format(best_acc))

        # load best model weights
        self.model.load_state_dict(best_model_wts)

        return val_acc_history

    def save(self, root, filename):
        torch.save({"model_state_dict": self.model.state_dict(),
                    "optimizer_state_dict": self.optimizer.state_dict()},
                   root+"/feature_extractors/your_super_extractor/"+filename+".pt")

    def load(self, root, filename):
        if torch.cuda.is_available():
            checkpoint = torch.load(root+"/feature_extractors/your_super_extractor/"+filename+".pt")
        else:
            checkpoint = torch.load(root +"/feature_extractors/your_super_extractor/" + filename + ".pt", map_location=torch.device('cpu'))

        self.model.load_state_dict(checkpoint['model_state_dict'])

        params = [p for p in self.model.parameters() if p.requires_grad]
        self.optimizer = torch.optim.SGD(params, lr=0.001)
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

        if torch.cuda.is_available():
            self.model.to('cuda')

    def __call__(self, *args):
        return self.model(*args)

    def set_train(self):
        self.model.train()

    def set_eval(self):
        self.model.eval()


from pathlib import Path

from setuptools import setup, find_packages

long_description = (Path(__file__).parent / "README.md").read_text('utf-8').split('# Installation')[0]

setup(
    name="mokuro",
    version='0.1.7',
    description="Browser reader for manga with selectable text",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/kha-white/mokuro",
    author="Maciej Budyś",
    author_email="kha-white@mail.com",
    license="GPLv3",
    classifiers=[
        "Programming Language :: Python :: 3",
    ],
    packages=find_packages(),
    include_package_data=True,
    install_requires=[
        "fire",
        "loguru",
        "manga-ocr>=0.1.7",
        "natsort",
        "numpy",
        "opencv-python>=4.1.2",
        "Pillow>=7.1.2",
        "pyclipper",
        "requests",
        "scipy",
        "shapely",
        "torch>=1.7.0",
        "torchsummary",
        "torchvision>=0.8.1",
        "transformers>=4.25.0",
        "tqdm>=4.41.0",
        "yattag",
    ],
    entry_points={
        "console_scripts": [
            "mokuro=mokuro.__main__:main",
        ]
    },
)


import json
import os
import shlex
import re
import sys
import codecs
from shutil import copyfile

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

from node.app import settings
from task.scmmgr import SCMMgr
from util.pathlib import PathMgr
from util.pathfilter import FilterPathUtil
from task.codelintmodel import CodeLintModel
from util.configlib import ConfigReader
from util.subprocc import SubProcController
from util.logutil import LogPrinter
from util.textutil import StringMgr
from util.tooldisplay import ToolDisplay
from util.exceptions import ConfigError

class Eslint(CodeLintModel):
    def __init__(self, params):
        CodeLintModel.__init__(self, params)
        self.sensitive_word_maps = {"eslint": "Tool", "Eslint": "Tool", "node": "tool", "Node": "tool"}

    def print_log(self, message):
        """
        将日志中的敏感词替换掉，再打印日志
        :param message:
        :return:
        """
        if "rules Loading rule" not in message and "Configuration was determined" not in message:
            ToolDisplay.print_log(self.sensitive, self.sensitive_word_maps, message)

    def analyze(self, params):
        """
        扫描分析源码
        :param params:
        :return:
        """
        envs = os.environ
        work_dir = os.getcwd()

        # 表示Eslint扫描Js以及React框架
        eslint_type = "JAVASCRIPT"
        eslint_ext = ".js,.jsx"
        if "ESLINT_%s_EXT" % eslint_type in envs:
            eslint_ext = envs.get("ESLINT_%s_EXT" % eslint_type)
        # 扫描html内嵌js
        if envs.get("ESLINT_EXT_HTML"):
            eslint_ext = eslint_ext + ",.html"

        error_output = os.path.join(work_dir, "%s_eslint_result.xml" % eslint_type.lower())

        config_path, rule_filte_flag = self.config(params, eslint_type, True)

        self.scan(params, config_path, error_output, eslint_ext)

        return self.data_handle(params, error_output, rule_filte_flag)

    def config(self, params, eslint_type, use_custom_prority=False):
        """
        设置eslint的配置文件
        :param params:
        :param eslint_type:
        :param use_custom_prority: REACT默认使用代码库下的配置文件，VUE和TYPESCRIPT默认使用CodeDog的配置文件
        :return:
        """
        source_dir = params.source_dir
        work_dir = params.work_dir
        envs = os.environ

        # 不使用CodeDog维护的配置文件时候，不对扫描结果规则过滤
        rule_filte_flag = False

        # 设置配置文件
        config_path = None
        # 兼容之前的ESLINT_CONFIG和ESLINT_CONFIG_TYPE
        eslint_config = envs.get("ESLINT_CONFIG", None)
        eslint_type_config = envs.get("ESLINT_%s_CONFIG" % eslint_type, None)
        eslint_config_type = envs.get("ESLINT_CONFIG_TYPE", None)
        eslint_type_config_type = envs.get("ESLINT_%s_CONFIG_TYPE" % eslint_type, None)

        if eslint_type == "JAVASCRIPT" and eslint_config:
            eslint_type_config = eslint_config

        if eslint_type_config:
            LogPrinter.info("复用项目中指定的配置文件进行扫描")
            config_path = eslint_type_config
            return config_path, rule_filte_flag

        if eslint_type == "JAVASCRIPT" and eslint_config_type:
            eslint_type_config_type = eslint_config_type

        # default custom google
        if eslint_type_config_type not in ("default", "custom", None):
            LogPrinter.info("使用" + eslint_type_config_type + "风格配置文件进行扫描")
            config_file = "%s_%s_eslintrc.json" % (eslint_type_config_type, eslint_type.lower())
            config_path = os.path.join(work_dir, config_file)
            copyfile(os.path.join(envs.get("NODE_HOME"), config_file), config_path)

            # 支持设置配置文件
            self._set_eslint_config(params, config_path, eslint_type, set_rule_flag=False)
            return config_path, rule_filte_flag

        if eslint_type_config_type == "custom":
            self.print_log("发现项目中含有eslint配置文件，复用该配置文件进行扫描")
            return config_path, rule_filte_flag

        if eslint_type_config_type == "default":
            pass
        elif use_custom_prority and self._find_eslintrc(source_dir):
            self.print_log("发现项目中含有eslint配置文件，复用该配置文件进行扫描")
            return config_path, rule_filte_flag

        LogPrinter.info("使用默认的AlloyTeam风格配置文件进行扫描")
        config_path = os.path.join(work_dir, os.path.join(work_dir, "%s_eslintrc.json" % eslint_type.lower()))
        copyfile(os.path.join(envs.get("NODE_HOME"), "%s_eslintrc.json" % eslint_type.lower()), config_path)
        # 设置配置
        self._set_eslint_config(params, config_path, eslint_type)
        rule_filte_flag = True

        return config_path, rule_filte_flag

    def _set_eslint_config(self, params, config_file, eslint_type, set_rule_flag=True):
        """
        给配置文件设置规则，将之置为warn级别
        注意：
        1. react扫描，即babel-eslint解析时候，no-unused-vars规则可能会解析失败，导致eslint执行失败
        :param params:
        :param config_file:
        :param eslint_type:
        :param set_rule_flag: True时候，设置配置文件的规则，对应not rule_filte_flag
        :return:
        """
        source_dir = params.source_dir
        work_dir = params.work_dir
        rule_list = params["rule_list"]
        envs = os.environ

        # 判断是否为ts,如果是配置tsconfig参数
        parserOptions = {}
        if "typescript" == eslint_type.lower():
            tsconfig = self._find_tsconfig(source_dir)
            LogPrinter.info("读取tsconfig文件配置: {}".format(tsconfig))
            parserOptions["project"] = tsconfig if tsconfig else self._set_tsconfig_file(work_dir)
            # 不存在配置文件 指定绝对路径空tsconfig.json文件

        # 设置eslint的global配置，以分号分隔，比如：ESLINT_GLOBALS=_:readonly;jtest:readonly
        # 对于每个全局变量键，将对应的值设置为 "writable" 以允许重写变量，或 "readonly" 不允许重写变量, "off" 禁用全局变量
        # 由于历史原因，布尔值 false 和字符串值 "readable" 等价于 "readonly"。类似地，布尔值 true 和字符串值 "writeable" 等价于 "writable"
        eslint_globals = envs.get("ESLINT_%s_GLOBALS" % eslint_type, "")

        # 设置eslint的环境，一个环境定义了一组预定义的全局变量。以分号分隔
        eslint_env = envs.get("ESLINT_%s_ENV" % eslint_type, "")

        # Eslint配置文件，添加在现有维护的配置文件中
        eslint_options = os.path.join(source_dir, envs.get("ESLINT_%s_OPTIONS" % eslint_type, ""))

        with open(config_file, "r") as f:
            configContent = f.read()
            configContentJson = json.loads(configContent)
            confi_env = configContentJson["env"]
            confi_globals = configContentJson["globals"]
            config_rules = configContentJson["rules"]
            config_plugins = configContentJson["plugins"]

        # enable rules
        if set_rule_flag and rule_list:
            self._set_rule_config(config_rules, rule_list, config_plugins)

        # 设置env
        for e in eslint_env.split(";"):
            if not e:
                continue
            confi_env[e] = True

        # 设置globals配置
        for g in eslint_globals.split(";"):
            tem = g.split(":")
            if len(tem) != 2 or tem[1] == "":
                continue
            confi_globals[tem[0]] = tem[1]

        # 设置HTML插件
        if envs.get("ESLINT_EXT_HTML"):
            if "plugins" in configContentJson.keys():
                configContentJson["plugins"].append("html")
            else:
                configContentJson["plugins"] = ["html"]

        # 使用文件设置
        if eslint_options and os.path.isfile(eslint_options):
            self.print_log("eslint_options: %s" % eslint_options)
            with open(eslint_options, "r") as f:
                add = json.load(f)
            self._recursive_merge(configContentJson, add)
        # 配置 tsconfig 配置
        if parserOptions:
            if "parserOptions" in configContentJson.keys():
                configContentJson["parserOptions"]["project"] = parserOptions["project"]
            else:
                configContentJson["parserOptions"] = parserOptions
        with codecs.open(config_file, "w", encoding="utf-8") as f:
            json.dump(configContentJson, f, ensure_ascii=False, indent=2)

    def _set_rule_config(self, config_rules, rule_list, config_plugins):
        """
        规则的配置
        """
        # 设置eslint插件，根据所选规则添加对应插件
        jsx_a11y_plugin = False
        vuejs_accessibility_plugin = False

        for rule in rule_list:
            rule_name = rule["name"]
            # 根据规则名判断插件
            if rule_name == "parse-error":
                continue
            if rule_name.startswith("jsx-a11y/"):
                jsx_a11y_plugin = True
            if rule_name.startswith("vuejs-accessibility/"):
                vuejs_accessibility_plugin = True
            # 兼容旧版typescript的规则
            if rule_name.startswith("typescript/"):
                rule_name = rule_name.replace("typescript/", "@typescript-eslint/")
            config_rules[rule_name] = "warn"
            # 支持规则参数设置
            if not rule["params"]:
                continue
            param = rule["params"]
            # 1. 直接拷贝在配置文件Json中的设置
            try:
                # demo: [\"error\", \"tab\", { \"SwitchCase\": 60 }]
                config_rules[rule_name] = json.loads(param)
            except json.JSONDecodeError as e:
                # 2. 使用key-value方式
                # 数组的非字典元素，使用options指定
                # 数组的字典元素，就使用key/value的方式
                # 若十分复杂，就直接传参json字符串
                # demo: options=tab\nSwitchCase=1
                if "[block]" in param:
                    rule_params = param
                else:
                    rule_params = "[block]\r\n" + param
                rule_params_dict = ConfigReader(cfg_string=rule_params).read("block")
                rule_params = ["warn"]
                if rule_params_dict.get("options"):
                    rule_params.extend(
                        [StringMgr.trans_type(option) for option in rule_params_dict["options"].split(",")]
                    )
                    rule_params_dict.pop("options")
                if len(rule_params_dict.keys()) > 0:
                    for key in rule_params_dict:
                        rule_params_dict[key] = StringMgr.trans_type(rule_params_dict[key])
                    rule_params.append(rule_params_dict)
                config_rules[rule_name] = rule_params

        # 设置plugin
        if jsx_a11y_plugin:
            config_plugins.append("jsx-a11y")
        if vuejs_accessibility_plugin:
            config_plugins.append("vuejs-accessibility")

    def _find_eslintrc(self, root_dir):
        """
        查找项目目录下所有的.eslintrc.*
        这里只检查根目录下的文件是否包含eslint配置文件
        package.json文件也可以设置eslint配置，但这里不对package.josn进行检查，因为每个Js项目都可能有该文件
        :param root_dir:
        :return:
        """
        for fileOrDir in os.listdir(root_dir):
            path = os.path.join(root_dir, fileOrDir)
            if os.path.isdir(path):
                continue
            if re.match("(\.eslintrc.*)", fileOrDir, re.I):
                return True
        return False

    def _find_tsconfig(self, root_dir):
        """
        查找项目目录下所有的tsconfig.json文件
        某些规则需要配置tsconfig.json路径
        这里检查项目中已存在的tsconfig文件
        :param root_dir: 项目地址
        :return List: 所有tsconfig路径
        """
        tsconfigFiles = []
        dirList = os.walk(root_dir)
        for root, dirs, files in dirList:
            for f in files:
                if "tsconfig.json" == f:
                    tsconfigFiles.append(os.path.join(root, f))
        return tsconfigFiles

    def _set_tsconfig_file(self, work_dir):
        """写入默认tsconfig文件, 返回project配置
        :param work_dir:
        :return tsconfigFile: 默认tsconfig配置路径
        """
        tsconfigFile = "{}/tsconfig.json".format(work_dir)
        with open(tsconfigFile, "w") as f:
            json.dump(
                """{
                        "extends": "",
                        "compilerOptions": {},
                        "include": [],
                        "exclude": []
                        }""",
                f,
                indent=2,
            )
        return tsconfigFile

    def scan(self, params, config_path, error_output, eslint_ext):
        """

        :param params:
        :param config_path:
        :param error_output:
        :return:
        """
        source_dir = params.source_dir
        path_exclude = params.path_filters.get("wildcard_exclusion", [])
        path_include = params.path_filters.get("wildcard_inclusion", [])
        incr_scan = params["incr_scan"]
        envs = os.environ

        if "ESLINT_MAX_OLD_SPACE_SIZE" in envs:
            # 或者设置NODE_OPTIONS="--max-old-space-size=4096"
            scan_cmd = ["node", "--max-old-space-size=" + envs["ESLINT_MAX_OLD_SPACE_SIZE"]]
            if sys.platform == "win32":
                scan_cmd.append(os.path.join(envs.get("NODE_HOME"), "node_modules", "eslint", "bin", "eslint.js"))
            else:
                scan_cmd.append(
                    os.path.join(envs.get("NODE_HOME"), "lib", "node_modules", "eslint", "bin", "eslint.js")
                )
        else:
            scan_cmd = ["eslint"]

        if config_path:
            scan_cmd.extend(["--config", config_path, "--no-eslintrc"])

        scan_cmd.extend(["-f", "checkstyle", "-o", error_output, "--debug"])

        if "ESLINT_OPTION_PARAMS" in envs:
            scan_cmd.extend(shlex.split(envs["ESLINT_OPTION_PARAMS"]))

        scan_cmd.extend(["--ext", eslint_ext])

        # 设置扫描路径过滤 exclusion and inclusion
        scan_cmd.extend(self._get_path_filte(path_exclude, path_include))
        # 默认不使用代码库下eslintignore文件中的过滤路径, 除非设置环境变量ESLINT_IGNORE指定
        eslint_ignore = envs.get("ESLINT_IGNORE", None)
        ignore_file = eslint_ignore if eslint_ignore else self.get_ignore_file(params)
        scan_cmd.extend(["--ignore-path", ignore_file])

        # 指定扫描目录或者文件，支持增量
        relpos = len(source_dir) + 1
        if incr_scan:
            diffs = SCMMgr(params).get_scm_diff()
            toscans = [
                os.path.join(source_dir, diff.path)
                for diff in diffs
                if diff.path.endswith(tuple(eslint_ext.split(","))) and diff.state != "del"
            ]
            toscans = FilterPathUtil(params).get_include_files(toscans, relpos)
            toscans = [path[relpos:].replace(os.sep, "/") for path in toscans]

            # windows下增量文件路径总长度超长，改为扫描整个soucedir，否则会执行异常
            # 因为 windows 下执行subprocess.Popen方法且shell=False的情况下，命令长度限制为32768字节
            # 这里文件路径长度限制为32500，因为前面还要预留cppcheck命令的长度
            if sys.platform == "win32" and len(" ".join(toscans)) > 32500:
                toscans = ["."]
        else:
            toscans = ["."]

        if not toscans:
            LogPrinter.debug("To-be-scanned files is empty ")
            return
        scan_cmd.extend(toscans)

        # 由于eslint会按照js的严格模式来解析js代码，若解析错误便会产生一个error，导致返回值非0
        self.print_log("cmd: %s" % " ".join(scan_cmd))
        sub_proc = SubProcController(
            scan_cmd, cwd=source_dir, stdout_line_callback=self.print_log, stderr_line_callback=self.print_log
        )
        sub_proc.wait()
        # 适配代码库下没有js相关文件的情况
        if sub_proc.returncode != 0:
            LogPrinter.error("命令执行返回码为: %d" % sub_proc.returncode)
            err_log = self.__get_stderr_log(sub_proc)
            if err_log.find("No files matching the pattern") != -1:
                LogPrinter.error("指定扫描的代码目录下面无可扫描文件，请检查配置")
                return
            # 查找"Cannot find module" 可能会出现误报
            elif err_log.find("ESLint couldn't find the") != -1:
                raise ConfigError("复用客户代码库中的Eslint配置文件，但没有安装相关插件，请在机器上安装，或者评估使用默认Alloy配置")
            # 设置没有配置文件错误提示
            elif err_log.find("ESLint couldn't find a configuration file.") != -1:
                raise ConfigError("代码库根目录中没有发现Eslint配置文件")

            # 捕获JS堆溢出异常
            elif (
                err_log.find(
                    "FATAL ERROR: Ineffective mark-compacts near heap limit Allocation failed \
                    - JavaScript heap out of memory"
                )
                != -1
            ):
                raise ConfigError("JS堆溢出错误，麻烦设置过滤库文件或者设置增加JS堆")

    def data_handle(self, params, error_output, rule_filte_flag):
        """
        处理结果文件
        :param params:
        :param error_output:
        :param rule_filte_flag:
        :return:
        """
        source_dir = params.source_dir
        filter_mgr = FilterPathUtil(params)
        pos = len(source_dir) + 1
        rules = params.get("rules", [])

        if not os.path.exists(error_output) or os.stat(error_output).st_size == 0:
            LogPrinter.info("result is empty ")
            return []

        # 处理unicode字符，否则会导致解析失败, 并非所有有效的Unicode字符在XML中都有效
        text = open(error_output, "r").read()
        # 替换所有无效Unicode字符
        text = re.sub(
            r"&#([0-8]);|&#1([1-24-9]);|&#2([0-9]);|&#3([0-1]);|&#x0?([0-8bcBCefEF]);|&#x1([0-9a-fA-F])",
            "",
            text,
            count=0,
            flags=re.I,
        )
        raw_warning = ET.fromstring(text)
        issues = []
        for file in raw_warning.findall("file"):
            # tag, attrib
            path = file.attrib.get("name")[pos:]
            if filter_mgr.should_filter_path(path):
                continue
            for error in file:
                # 若解析文件失败，便可能会line为undefined，该结果无效
                if error.attrib.get("line") == "undefined":
                    continue
                line = int(error.attrib.get("line"))
                column = error.attrib.get("column")
                if column in ("NaN",):
                    column = 0
                # need rule real name
                tmp = error.attrib.get("source").split(".")
                rule = tmp[-1]
                msg = error.attrib.get("message")
                # 记录解析错误带来的问题
                if rule == "" and msg.startswith("Parsing error:"):
                    rule = "parse-error"
                    # parse-error规则，移除msg中的代码块
                    if msg.find(".\n\n") != -1:
                        msg = "".join(msg.split(".\n\n")[:-1])
                # 兼容旧版typescript的规则
                # if alloy_version == "2.0.5" and rule.startswith("@typescript-eslint/"):
                if rule.startswith("@typescript-eslint/"):
                    rule = rule.replace("@typescript-eslint/", "typescript/")
                # 若使用客户的配置文件，不进行规则过滤
                if rule_filte_flag and rule not in rules:
                    continue
                if msg.startswith("Definition for rule"):
                    continue
                issues.append({"path": path, "rule": rule, "msg": msg, "line": line, "column": column})
        return issues

    def __get_stderr_log(self, sub_proc):
        """
        获取执行stderr的log
        :param sub_proc:
        :return: stderr中的log
        """
        if not os.path.exists(sub_proc.stderr_filepath):
            return ""
        spc_err = open(sub_proc.stderr_filepath, "r", encoding="utf-8")
        log = spc_err.read()
        spc_err.close()
        return log

    def _recursive_merge(self, src, add):
        """
        递归合并两json内容
        :param src: 原始json
        :param add: 需要添加的内容
        :return:
        """
        if isinstance(src, dict):
            for a in add:
                if a in src:
                    # 如果在src中，且是复合类型，则递归
                    if isinstance(src[a], dict) or isinstance(src[a], list):
                        self._recursive_merge(src[a], add[a])
                    else:
                        # 若是String或者bool等基本类型，则直接覆盖
                        src[a] = add[a]
                else:
                    # 如果不在src中，则添加
                    src[a] = add[a]
        elif isinstance(src, list):
            # list只能添加，不能继续递归修改
            for a in add:
                if a not in src:
                    src.append(a)

    def _get_path_filte(self, path_exclude, path_include):
        """
        给scan_cmd设置路径过滤扫描的命令行配置
        其中：
        **匹配多个目录，当test/**时候，匹配test目录下所有目录和文件
        *匹配除了/以外的任意多个字符
        ?匹配出了/以外的一个字符
        :param path_exclude:
        :param path_include:
        :return:
        """
        path_filter = []

        # 使用!来做--ignore-pattern的否定判断
        # 默认会过滤掉/**/node_modules/*  但是--ignore-pattern优先级更高
        # 但是这个不能不扫描include以外的文件，只是表示要扫描哪些文件
        if path_include:
            for tmp_path in path_include:
                path_filter.append("--ignore-pattern")
                tmp_path = re.sub("^\*/", "**/", tmp_path)
                tmp_path = re.sub("/\*$", "/**", tmp_path)
                tmp_path = re.sub("/\*/", "/**/", tmp_path)
                path_filter.append(f"'!{tmp_path}'")

        if path_exclude:
            for tmp_path in path_exclude:
                path_filter.append("--ignore-pattern")
                tmp_path = re.sub("^\*/", "**/", tmp_path)
                tmp_path = re.sub("/\*$", "/**", tmp_path)
                tmp_path = re.sub("/\*/", "/**/", tmp_path)
                path_filter.append(f"'{tmp_path}'")

        return path_filter

    def get_ignore_file(self, params, eslint_ext=None):
        """
        1. 忽略文件需要在代码库根目录下，否则base目录有误，会取代码库根目录和忽略文件路径的交叉处作为baseDir
        2. 会增加内存。目前8GB堆内存，可以支撑54300条，耗时从原本的5min增加到16min，若没有**/前缀，则是13min
        :param params:
        :param eslint_ext:
        :return:
        """
        source_dir = params.source_dir
        relpos = len(source_dir) + 1
        work_dir = params.work_dir
        ignore_file = os.path.join(work_dir, "codedog_eslintignore")

        if eslint_ext:
            want_suffix = eslint_ext.split(",")
            files = PathMgr().get_dir_files(source_dir, tuple(want_suffix))
            files = [path.replace(os.sep, "/") for path in files]
            include_files = FilterPathUtil(params).get_include_files(files, relpos)
            exclude_files = list(set(files) - set(include_files))
            ignore_files = ["**/%s" % path[relpos:] for path in exclude_files]
        else:
            ignore_files = list()
        with open(ignore_file, "w", encoding="UTF-8") as f:
            f.write("\n".join(ignore_files))
        return ignore_file

    def check_tool_usable(self, tool_params):
        """
        这里判断机器是否支持运行eslint
        1. 支持的话，便在客户机器上扫描
        2. 不支持的话，就发布任务到公线机器扫描
        :return:
        """
        if SubProcController(["node", "--version"], stderr_line_callback=self.print_log).wait() != 0:
            return []
        if SubProcController(["eslint", "--version"], stderr_line_callback=self.print_log).wait() != 0:
            return []
        return ["analyze"]

tool = Eslint

if __name__ == "__main__":
    pass


'''
时间：2021年03月30日
日志：复现KEFE方法中对于数据与处理的部分
步骤：去除噪声值，包括
    · 非英文数据
    · emoji
    · 特殊符号
    · 特殊数字
'''
import re
import json

app_dictionary = {
    "yahoo-mail": 100001, "gmail": 100002,
    "instagram": 100003, "snapchat": 100004,
    "com.chase.sig.android": 100005, "com.bpi.ng.mobilebanking": 100006,
}

def remove_serial_number(text):
    text = re.sub(r'[^a-zA-Z0-9,.\!?]+', ' ', text)
    text = re.sub(r'[●,=√#￥%&,[]【】◆ ★◎☆⊙]', ' ', text)
    return text

def remove_exce_punc(text):
    duels = [x + x for x in list('。，,=！!-—#')]
    # print(duels)
    for d in duels:
      while d in text:
        text = text.replace(d, d[0])
    return text

def pre_process(file_path):
    print("[INFO] Start to manage: " + file_path)
    lines = open(file_path, encoding="utf-8").readlines()
    print("[FILE] The file has " + str(len(lines)) + " lines.")
    # write_file = open("train/preprocess_result/" + file_path.strip().split("/")[8]+ "_clean_" + file_path.strip().split("/")[9], "w", encoding="utf-8")
    write_file = open("test/preprocess_result/" + file_path.strip().split("/")[8]+ "_clean_" + file_path.strip().split("/")[9], "w", encoding="utf-8")
    id_list = []
    for line in lines:
        line = json.loads(line)
        id = line["id"]
        id_list.append(id)
        text = line["text"].strip()
        # 除去文本之中的噪声项
        text = remove_serial_number(str(text))
        text_remove_noisy = remove_exce_punc(text)
        write_file.write(str(id) + "******" + text_remove_noisy + "\n")
    write_file.close()
    print("[END] Finish writing file " + "preprocess_result/" + file_path.strip().split("/")[8]+ "_clean_data.json")

def read_file(trainData_or_testData):
    file_path_list = []
    for i in range(0, 10): file_path_list.append("data/fold_"+ str(i) + "/" + str(trainData_or_testData) + ".json")
    return file_path_list

if __name__ == "__main__":
    # trainData_or_testData = "train" # 预处理训练集数据中的「特征候选短语」
    trainData_or_testData = "dev" # 预处理测试集数据中的「特征候选短语」
    file_path_list = read_file(trainData_or_testData)
    for one_file_path in file_path_list: pre_process(one_file_path)


import smartpy as sp

class PaymentCapturer(sp.Contract):
    def __init__(self, admin_address, treasury_address, refund_manager_address):
        self.init_type(sp.TRecord(admin_address = sp.TAddress, refund_manager_address = sp.TAddress, treasury_address = sp.TAddress))
        self.init(admin_address=admin_address,
                  treasury_address=treasury_address,
                  refund_manager_address=refund_manager_address
                  )

    @sp.entry_point
    def pay(self, metadata):
        sp.set_type(metadata, sp.TString)
        prepared_metadata = (metadata, sp.amount, sp.source)
        sp.emit(prepared_metadata, with_type=True, tag="PAY")

    @sp.entry_point
    def refund(self, metadata, amount, buyer):
        sp.set_type(metadata, sp.TString)
        sp.verify_equal(self.data.refund_manager_address, sp.source, "Refund not authorised.")
        sp.verify(sp.balance >= amount, "Contract does not have enough balance.")
        sp.send(buyer, amount, "Failed to refund.")
        sp.emit(metadata, with_type=True, tag="REFUND")

    @sp.entry_point
    def withdraw(self):
        sp.verify_equal(self.data.admin_address, sp.source, "Withdraw not authorised.")
        sp.send(self.data.treasury_address, sp.balance, "Could not withdraw.")

    @sp.entry_point
    def withdrawAmount(self, amount):
        sp.verify_equal(self.data.admin_address, sp.source, "Withdraw not authorised.")
        sp.send(self.data.treasury_address, amount, "Could not withdraw.")

    @sp.add_test(name="PaymentsCapturer")
    def test():
        admin = sp.test_account("admin")
        treasury = sp.test_account("treasury")
        refund_manager = sp.test_account("refund_manager")
        user = sp.test_account("user")

        r = PaymentCapturer(admin.address, treasury.address, refund_manager.address)
        scenario = sp.test_scenario()
        scenario.h1("Payments Capturer")
        scenario += r

        scenario.h2("Regular flow")
        scenario.h3("Pay")
        pay_amount = sp.tez(100)
        pay_metadata = "ok"
        scenario += r.pay(pay_metadata).run(source=user.address, amount=pay_amount, valid=True)
        scenario.verify_equal(r.balance, pay_amount)

        scenario.h3("Unauthorised Refund")
        refund_amount = sp.tez(40)
        refund_metadata = "refund"
        user = sp.test_account("user")
        scenario += r.refund(metadata=refund_metadata, amount=refund_amount, buyer=user.address).run(source=user.address, valid=False)
        scenario.verify_equal(r.balance, pay_amount)

        scenario.h3("Authorised Refund")
        refund_amount = sp.tez(40)
        refund_metadata = "refund"
        user = sp.test_account("user")
        scenario += r.refund(metadata=refund_metadata, amount=refund_amount, buyer=user.address).run(source=refund_manager.address, valid=True)
        expected_contract_balance = pay_amount - refund_amount
        scenario.verify_equal(r.balance, expected_contract_balance)

        scenario.h3("Unauthorised WithdrawAmount")
        amount = sp.tez(10)
        scenario += r.withdrawAmount(amount).run(source=user.address, valid=False)
        scenario.verify_equal(r.balance, expected_contract_balance)

        scenario.h3("Authorised WithdrawAmount")
        scenario += r.withdrawAmount(amount).run(source=admin.address, valid=True)
        scenario.verify_equal(r.balance, expected_contract_balance - amount)

        scenario.h3("Unauthorised Withdraw")
        scenario += r.withdraw().run(source=user.address, valid=False)
        scenario.verify_equal(r.balance, expected_contract_balance - amount)

        scenario.h3("Authorised Withdraw")
        scenario += r.withdraw().run(source=admin.address, valid=True)
        scenario.verify_equal(r.balance, sp.tez(0))

sp.add_compilation_target("my_contract_compiled", PaymentCapturer(admin_address=sp.address("tz1NHjoiSpRj8gzJbAxQSyTfiPwfxcRnnVfn"), treasury_address=sp.address("tz1NHjoiSpRj8gzJbAxQSyTfiPwfxcRnnVfn"), refund_manager_address=sp.address("tz1NHjoiSpRj8gzJbAxQSyTfiPwfxcRnnVfn")))