#import matplotlib with pdf as backend

import matplotlib 

matplotlib.use('PDF')

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

import wfdb 

import os

import numpy as np

import math

import sys

import scipy.stats as st

import glob, os

from os.path import basename

import tensorflow as tf

from keras.layers import Dense,Activation,Dropout

from keras.layers import LSTM,Bidirectional #could try TimeDistributed(Dense(...))

from keras.models import Sequential, load_model

from keras import optimizers,regularizers

from keras.layers.normalization import BatchNormalization

import keras.backend.tensorflow_backend as KTF

np.random.seed(0)

# functions

def get_ecg_data(datfile): 

    ## convert .dat/q1c to numpy arrays

    recordname=os.path.basename(datfile).split(".dat")[0]

    recordpath=os.path.dirname(datfile)

    cwd=os.getcwd()

    os.chdir(recordpath) ## somehow it only works if you chdir. 

    annotator='q1c'

    annotation = wfdb.rdann(recordname, extension=annotator, sampfrom=0,sampto = None, pbdir=None)

    Lstannot=list(zip(annotation.sample,annotation.symbol,annotation.aux_note))

    FirstLstannot=min( i[0] for i in Lstannot)

    LastLstannot=max( i[0] for i in Lstannot)-1

    print("first-last annotation:", FirstLstannot,LastLstannot)

    record = wfdb.rdsamp(recordname, sampfrom=FirstLstannot,sampto = LastLstannot) #wfdb.showanncodes()

    annotation = wfdb.rdann(recordname, annotator, sampfrom=FirstLstannot,sampto = LastLstannot) ## get annotation between first and last. 

    annotation2 = wfdb.Annotation(recordname='sel32', extension='niek', sample=(annotation.sample-FirstLstannot), symbol = annotation.symbol, aux_note=annotation.aux_note)

    Vctrecord=np.transpose(record.p_signals)

    VctAnnotationHot=np.zeros( (6,len(Vctrecord[1])), dtype=np.int)

    VctAnnotationHot[5]=1 ## inverse of the others 

    #print("ecg, 2 lead of shape" , Vctrecord.shape) 

    #print("VctAnnotationHot of shape" , VctAnnotationHot.shape) 

    #print('plotting extracted signal with annotation')

    #wfdb.plotrec(record, annotation=annotation2, title='Record 100 from MIT-BIH Arrhythmia Database', timeunits = 'seconds')

    VctAnnotations=list(zip(annotation2.sample,annotation2.symbol)) ## zip coordinates + annotations (N),(t) etc)

    #print(VctAnnotations)

    for i in range(len(VctAnnotations)):

        #print(VctAnnotations[i]) # Print to display annotations of an ecg

        try: 

            if VctAnnotations[i][1]=="p":

                if VctAnnotations[i-1][1]=="(":

                    pstart=VctAnnotations[i-1][0]

                if VctAnnotations[i+1][1]==")":

                    pend=VctAnnotations[i+1][0]

                if VctAnnotations[i+3][1]=="N":

                    rpos=VctAnnotations[i+3][0]

                    if VctAnnotations[i+2][1]=="(":

                        qpos=VctAnnotations[i+2][0]

                    if VctAnnotations[i+4][1]==")":

                        spos=VctAnnotations[i+4][0] 

                    for ii in range(0,8): ## search for t (sometimes the "(" for the t  is missing  )

                        if VctAnnotations[i+ii][1]=="t":

                            tpos=VctAnnotations[i+ii][0]

                            if VctAnnotations[i+ii+1][1]==")":

                                tendpos=VctAnnotations[i+ii+1][0]

            #               #print(ppos,qpos,rpos,spos,tendpos)

                                VctAnnotationHot[0][pstart:pend]=1 #P segment

                                VctAnnotationHot[1][pend:qpos]=1 #part "nothing" between P and Q, previously left unnanotated, but categorical probably can't deal with that

                                VctAnnotationHot[2][qpos:rpos]=1 #QR

                                VctAnnotationHot[3][rpos:spos]=1 #RS

                                VctAnnotationHot[4][spos:tendpos]=1 #ST (from end of S to end of T)

                                VctAnnotationHot[5][pstart:tendpos]=0 #tendpos:pstart becomes 1, because it is inverted above                   

        except IndexError:

            pass

    Vctrecord=np.transpose(Vctrecord) # transpose to (timesteps,feat)

    VctAnnotationHot=np.transpose(VctAnnotationHot)

    os.chdir(cwd)

    return Vctrecord, VctAnnotationHot

def splitseq(x,n,o):

    #split seq; should be optimized so that remove_seq_gaps is not needed. 

    upper=math.ceil( x.shape[0] / n) *n

    print("splitting on",n,"with overlap of ",o,    "total datapoints:",x.shape[0],"; upper:",upper)

    for i in range(0,upper,n):

        #print(i)

        if i==0:

            padded=np.zeros( ( o+n+o,x.shape[1])   ) ## pad with 0's on init

            padded[o:,:x.shape[1]] = x[i:i+n+o,:]

            xpart=padded

        else:

            xpart=x[i-o:i+n+o,:]

        if xpart.shape[0]<i:

            padded=np.zeros( (o+n+o,xpart.shape[1])  ) ## pad with 0's on end of seq

            padded[:xpart.shape[0],:xpart.shape[1]] = xpart

            xpart=padded

        xpart=np.expand_dims(xpart,0)## add one dimension; so that you get shape (samples,timesteps,features)

        try:

            xx=np.vstack(  (xx,xpart) )

        except UnboundLocalError: ## on init

            xx=xpart

    print("output: ",xx.shape)

    return(xx)

def remove_seq_gaps(x,y):

    #remove parts that are not annotated <- not ideal, but quickest for now.

    window=150

    c=0

    cutout=[]

    include=[]

    print("filterering.")

    print("before shape x,y",x.shape,y.shape)

    for i in range(y.shape[0]):

        c=c+1

        if c<window :

            include.append(i)

        if sum(y[i,0:5])>0:

            c=0 

        if c >= window:

            #print ('filtering')

            pass

    x,y=x[include,:],y[include,:]

    print(" after shape x,y",x.shape,y.shape)

    return(x,y)

def normalizesignal(x):

    x=st.zscore(x, ddof=0)

    return x

def normalizesignal_array(x):

    for i in range(x.shape[0]):

        x[i]=st.zscore(x[i], axis=0, ddof=0)

    return x

def plotecg(x,y,begin,end):

    #helper to plot ecg

    plt.figure(1,figsize=(11.69,8.27))

    plt.subplot(211)

    plt.plot(x[begin:end,0])

    plt.subplot(211)

    plt.plot(y[begin:end,0])

    plt.subplot(211)

    plt.plot(y[begin:end,1])

    plt.subplot(211)

    plt.plot(y[begin:end,2])

    plt.subplot(211)

    plt.plot(y[begin:end,3])

    plt.subplot(211)

    plt.plot(y[begin:end,4])

    plt.subplot(211)

    plt.plot(y[begin:end,5])

    plt.subplot(212)

    plt.plot(x[begin:end,1])

    plt.show()

def plotecg_validation(x,y_true,y_pred,begin,end):

    #helper to plot ecg

    plt.figure(1,figsize=(11.69,8.27))

    plt.subplot(211)

    plt.plot(x[begin:end,0])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,0])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,1])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,2])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,3])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,4])

    plt.subplot(211)

    plt.plot(y_pred[begin:end,5])

    plt.subplot(212)

    plt.plot(x[begin:end,1])

    plt.subplot(212)

    plt.plot(y_true[begin:end,0])

    plt.subplot(212)

    plt.plot(y_true[begin:end,1])

    plt.subplot(212)

    plt.plot(y_true[begin:end,2])

    plt.subplot(212)

    plt.plot(y_true[begin:end,3])

    plt.subplot(212)

    plt.plot(y_true[begin:end,4])

    plt.subplot(212)

    plt.plot(y_true[begin:end,5])

def LoaddDatFiles(datfiles):  

    for datfile in datfiles:

        print(datfile)

        if basename(datfile).split(".",1)[0] in exclude:

            continue

        qf=os.path.splitext(datfile)[0]+'.q1c'

        if os.path.isfile(qf):

            #print("yes",qf,datfile)

            x,y=get_ecg_data(datfile)

            x,y=remove_seq_gaps(x,y)

            x,y=splitseq(x,1000,150),splitseq(y,1000,150) ## create equal sized numpy arrays of n size and overlap of o 

            x = normalizesignal_array(x)

            ## todo; add noise, shuffle leads etc. ?

            try: ## concat

                xx=np.vstack(  (xx,x) )

                yy=np.vstack(  (yy,y) )

            except NameError: ## if xx does not exist yet (on init)

                xx = x

                yy = y

    return(xx,yy)

def unison_shuffled_copies(a, b):

    assert len(a) == len(b)

    p = np.random.permutation(len(a))

    return a[p], b[p]

def get_session(gpu_fraction=0.8):

    #allocate % of gpu memory.

    num_threads = os.environ.get('OMP_NUM_THREADS')

    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)

    if num_threads:

        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, intra_op_parallelism_threads=num_threads))

    else:

        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

def getmodel():

    model = Sequential()

    model.add(Dense(32,W_regularizer=regularizers.l2(l=0.01), input_shape=(seqlength, features)))

    model.add(Bidirectional(LSTM(32, return_sequences=True)))#, input_shape=(seqlength, features)) ) ### bidirectional ---><---

    model.add(Dropout(0.2))

    model.add(BatchNormalization())

    model.add(Dense(64, activation='relu',W_regularizer=regularizers.l2(l=0.01)))

    model.add(Dropout(0.2))

    model.add(BatchNormalization())

    model.add(Dense(dimout, activation='softmax'))

    adam = optimizers.adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

    model.compile(loss='categorical_crossentropy', optimizer=adam, metrics=['accuracy']) 

    print(model.summary())

    return(model)

##################################################################

##################################################################

qtdbpath=sys.argv[1] ## first argument = qtdb database from physionet. 

perct=0.81 #percentage training

percv=0.19 #percentage validation

exclude = set()

exclude.update(["sel35","sel36","sel37","sel50","sel102","sel104","sel221","sel232", "sel310"])# no P annotated:

##################################################################

# datfile=qtdbpath+"sel49.dat"  ## single ECG to test if loading works.  

# x,y=get_ecg_data(datfile)

# print(x.shape,y.shape)

# # for i in range(y.shape[0]): #Invert QT-label to actually represent QT. Does give overlapping labels

# #     y[i][0] = 1 - y[i][0]

# plotecg(x,y,0,y.shape[0]) ## plot all

# x,y=remove_seq_gaps(x,y) ## remove 'annotation gaps'

# plotecg(x,y,0,y.shape[0]) ## plot all

# x,y=splitseq(x,750,150),splitseq(y,750,150) ## create equal sized numpy arrays of n size and overlap of o 

# exit()

##################################################################

# load data

datfiles=glob.glob(qtdbpath+"*.dat")

xxt,yyt=LoaddDatFiles(datfiles[ :round(len(datfiles)*perct) ]) # training data. 

xxt,yyt=unison_shuffled_copies(xxt,yyt) ### shuffle

xxv,yyv=LoaddDatFiles(datfiles[ -round(len(datfiles)*percv): ] ) ## validation data.

seqlength=xxt.shape[1]

features=xxt.shape[2]

dimout=yyt.shape[2]

print("xxv/validation shape: {}, Seqlength: {}, Features: {}".format(xxv.shape[0],seqlength,features))

# #plot validation ecgs 

# with PdfPages('ecgs_xxv.pdf') as pdf:

#   for i in range( xxv.shape[0] ): 

#       print (i)

#       plotecg(xxv[i,:,:],yyv[i,:,:],0,yyv.shape[1])

#       pdf.savefig()

#       plt.close()

# call keras/tensorflow and build lstm model 

KTF.set_session(get_session())

with tf.device('/cpu:0'): #switch to /cpu:0 to use cpu 

    if not os.path.isfile('model.h5'):

        model = getmodel() # build model

        model.fit(xxt, yyt, batch_size=32, epochs=100, verbose=1) # train the model

        model.save('model.h5')

    model = load_model('model.h5')

    score, acc = model.evaluate(xxv, yyv, batch_size=4, verbose=1)

    print('Test score: {} , Test accuracy: {}'.format(score, acc))

    # predict

    yy_predicted = model.predict(xxv) 

    # maximize probabilities of prediction.

    for i in range(yyv.shape[0]): 

        b = np.zeros_like(yy_predicted[i,:,:])

        b[np.arange(len(yy_predicted[i,:,:])), yy_predicted[i,:,:].argmax(1)] = 1

        yy_predicted[i,:,:] = b

    # plot: 

    with PdfPages('ecg.pdf') as pdf:

        for i in range( xxv.shape[0] ): 

            print (i)

            plotecg_validation(xxv[i,:,:],yy_predicted[i,:,:],yyv[i,:,:],0,yy_predicted.shape[1])  # top = predicted, bottom=true

            pdf.savefig()

            plt.close()

    #plotecg(xv[1,:,:],yv[1,:,:],0,yv.shape[1]) ## plot first seq