#!/usr/bin/python
#challenge entry
import matplotlib
matplotlib.use('Agg')
from keras import backend as K
K.set_image_dim_ordering('tf')
import os
import sys
import scipy.io
import math as m
import numpy as np
from scipy.signal import resample,filtfilt
import matplotlib.pyplot as plt
import pylab as pylab
from pylab import plot,show,specgram
import keras
from keras.models import Sequential, load_model
from keras.models import Model, model_from_json
from keras.layers.core import Dense, Dropout, Flatten, Activation
from keras.layers.convolutional import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.optimizers import SGD
from keras.callbacks import EarlyStopping
from keras.utils import np_utils
from sklearn.metrics import log_loss
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
from sklearn.model_selection import train_test_split
import warnings
from keras.layers import Input
from keras import layers
from keras.layers import Activation
from keras.layers import Flatten

b1 = [-7.757327341237223e-5,  -2.357742589814283e-4, -6.689305101192819e-4, -0.001770119249103, -0.004364327211358, -0.010013251577232, -0.021344241245400, -0.042182820580118, -0.077080889653194,     -0.129740392318591, -0.200064921294891, -0.280328573340852, -0.352139052257134, -0.386867664739069,-0.351974030208595, -0.223363323458050, 0, 0.286427448595213, 0.574058766243311, 0.788100265785590, 0.867325070584078, 0.788100265785590, 0.574058766243311, 0.286427448595213, 0, -0.223363323458050, -0.351974030208595, -0.386867664739069, -0.352139052257134, -0.280328573340852, -0.200064921294891, -0.129740392318591, -0.077080889653194, -0.042182820580118, -0.021344241245400, -0.010013251577232, -0.004364327211358, -0.001770119249103, -6.689305101192819e-4,-2.357742589814283e-4, -7.757327341237223e-5]
b1=resample(b1,50)
input_C = 2
input_H = 17 
input_W = 513 
nType=4
labs = ['N', 'A', 'O', '~']

def ecg_padding(ecgwave):
    ecg_duration=ecgwave.size
    if ecg_duration<18300:
	    ecg_padded=np.concatenate((ecgwave,np.zeros(18300-ecg_duration)),axis=0)
    return ecg_padded

def bpfilterecg(ecgwave):
    filterecg=filtfilt(b1,1,ecgwave)
    return filterecg

def ecg_rep(ecgwave):
    ecg_duration=ecgwave.size
    reps=0;remainder=0; outecg=[];
    # max duration of ecg= 61s; sampling rate=300Hz, total number of samples=18300
    # following code repeats ecg waveform until total duration is 18300
    if (ecg_duration<18300): 
        reps=m.floor(18300/ecg_duration);
        remainder=m.fmod(18300,ecg_duration);
	#import pdb;pdb.set_trace()
        outecg_A=ecgwave
	for i in range(int(reps)-1):
            outecg_A=np.concatenate((outecg_A,ecgwave),0);
        if (int(remainder)>0):
            outecg=np.concatenate((outecg_A,ecgwave[0:int(remainder)]),axis=0);
            #print(outecg.size)
	    
        else:
	    outecg=outecg_A
    else:
        outecg=ecgwave
    return outecg



def spectrogram_gen(ecgwave,flbase=''):
    ecg_duration=ecgwave.size
    reps=0;remainder=0; outecg=[];
    #outecg=ecg_rep(ecgwave)  
    #outecg=ecg_padding(ecgwave)
    outecg=bpfilterecg(ecgwave)
    
    positiveecg,negativeecg=[],[]
    for i in outecg:
        if i<0:
            positiveecg.append(0)
        else:
            positiveecg.append(i)
    for i in outecg:
        if i<0:
            negativeecg.append(i)
        else:
            negativeecg.append(0)
     
    [ptspectrum,ptfreqs,pt_t,pt_im]=pylab.specgram(positiveecg,NFFT=1024,  noverlap=0)
    [negspectrum,negfreqs,negt,neg_im]=pylab.specgram(negativeecg,NFFT=1024,  noverlap=0)
    [spectrum,freqs,t,im]=pylab.specgram(outecg,NFFT=1024,noverlap=0)
    
    negspectrum=np.swapaxes(negspectrum,0,1)
    negspectrum= negspectrum.astype('float32')
    ptspectrum=np.swapaxes( ptspectrum,0,1)
    ptspectrum = ptspectrum.astype('float32')
    spectrum=np.swapaxes(spectrum,0,1)
    spectrum= spectrum.astype('float32')
    
    for item in range(pt_t.size):
        if np.std(ptspectrum[item,:])!=0:
           ptspectrum[item,:]=(ptspectrum[item,:]-np.mean(ptspectrum[item,:]))/np.std(ptspectrum[item,:])
        else:
            ptspectrum[item,:]=(ptspectrum[item,:]-np.mean(ptspectrum[item,:]))
    for item in range(t.size):
        if np.std(spectrum[item,:])!=0:
           spectrum[item,:]=(spectrum[item,:]-np.mean(spectrum[item,:]))/np.std(spectrum[item,:])
        else:
           spectrum[item,:]=(spectrum[item,:]-np.mean(spectrum[item,:]))
    for item in range(negt.size):
        if np.std( negspectrum[item,:])!=0:
            negspectrum[item ,:]=( negspectrum[item,:]-np.mean( negspectrum[item,:]))/np.std( negspectrum[item,:])
        else:
             negspectrum[item,:]=( negspectrum[item,:]-np.mean( negspectrum[item,:]))
    
    zeropad=np.zeros((1,513))
    if t.size<17:
        pad_reps=17-t.size
        for i in range(pad_reps):
            spectrum=np.vstack((spectrum,zeropad)) 
    if pt_t.size<17:
        pad_reps=17-pt_t.size
        for i in range(pad_reps):
            ptspectrum=np.vstack((ptspectrum,zeropad))      
    if negt.size<17:
        pad_reps=17-negt.size
        for i in range(pad_reps):
            negspectrum=np.vstack((negspectrum,zeropad))
    ptspectrum=np.expand_dims(ptspectrum,axis=2)
    negspectrum=np.expand_dims(negspectrum,axis=2)
    spectra=np.concatenate((ptspectrum,negspectrum),axis=2)
    return spectra


def build_model(nType):
    inputs = Input((input_H, input_W,input_C))
    
    conv1 = Conv2D(32, (3,3),  padding="same", activation="relu")(inputs)
    pool1 = MaxPooling2D(pool_size=(2, 2),padding="same")(conv1)
    
    conv2 = Conv2D(64, (3,3), padding="same", activation="relu")(pool1)
    pool2 = MaxPooling2D(pool_size=(2, 2),padding="same")(conv2)
    
    conv3 = Conv2D(128, (3,3),  padding="same", activation="relu")(pool2)
    pool3 = MaxPooling2D(pool_size=(2, 2),padding="same")(conv3)
    
    conv4 = Conv2D(256, (3,3),  padding="same", activation="relu")(pool3)
    pool4 = MaxPooling2D(pool_size=(2, 2),padding="same")(conv4)
    
    conv5 = Conv2D(512, (3,3),  padding="same", activation="relu")(pool4)
    pool5 = MaxPooling2D(pool_size=(2, 2),  padding="same")(conv5)
   
    flat1=Flatten()(pool5)

    dropout1=Dropout(.25)(flat1)

    dense1 =Dense(nType, activation='softmax')(dropout1)

    sgd = SGD(lr=1e-3, decay=1e-6, momentum=0.9, nesterov=True)
    model = Model(input=inputs, output=dense1)
    model.compile(optimizer=sgd, loss='categorical_crossentropy')
    return model

if __name__ == '__main__':

    record = sys.argv[1]
    # Read waveform samples (input is in WFDB-MAT format)
    mat_data = scipy.io.loadmat(record + ".mat")
    samples = mat_data['val'][0]
    X_test=spectrogram_gen(samples,record)
    X_test=np.expand_dims(X_test,axis=0)            
    ensemble_store=np.zeros((1,4))
    for i in range(1,5):
        best_model_file = './weights_'+str(i)+'.h5'
	model = build_model(nType)
	model.load_weights(best_model_file)
        Y_out=model.predict(X_test)  
        ensemble_store+=Y_out   
    Y_out_label=np.argmax(ensemble_store,axis=1)
    answer=labs[Y_out_label[0]]
    answers_file = open("answers.txt", "a")
    answers_file.write("%s,%s\n" % (record, answer))
    answers_file.close()
        
