# -*- coding: utf-8 -*-
"""
Created on Sun Aug 20 15:06:12 2017

@author: Phil Withmore
"""

# -*- coding: utf-8 -*-
"""
Created on Thu Aug  3 11:14:21 2017

@author: phil
"""
import os, sys

from keras.models import Model, load_model
from keras import backend as K

from numpy.core import multiarray
import scipy
from scipy import cluster, stats
import numpy as np
import csv, math
import dwt
import matplotlib.pyplot as plt
import wfdb
import copy
from sklearn import cluster, metrics
from sklearn.externals import joblib
from sklearn.ensemble import RandomForestClassifier

import pandas as pd

pthreshold = .25
rthreshold = .5
tthreshold = .25
othreshold = .25
nthreshold = .1

window  = 300
fs_target = 300

validation = 1

def nanTo9(values):
    ind = np.where(np.isnan(values))[0]
    values[ind] = 99999
    return values

# frame to scan, long files take so long  
def find_nearest(array,value, order):
    if len(array > 0):
        d = array - value
        if order < 0:
            d[d>0] = 999
        else:
            d[d<0] = 999
            
        idx = np.abs(d).argmin()
        return np.abs(d[idx])
    
    else:
        return 0

def find_array(stack, needle, hr, order):
    result = []
    for n in needle:
        r = find_nearest(stack, n, order)
        if(r < hr):
            result.append(r)
            
    return np.array(result)

#%% Accepts a List of Peaks in a set of Data
#   Returns Probabilitys of classification

def classify(peaks, data):
	# Initialize Arrays+
    dval = copy.deepcopy(data)
    
    pWave = np.zeros(len(dval))
    rSpike = np.zeros(len(dval))
    tWave = np.zeros(len(dval))
    noise = np.zeros(len(dval))
    zeros = np.zeros(len(dval))
    other = np.zeros(len(dval))

	# Rearrange Peaks    
    peaks = np.array(peaks)    
    peaks = np.swapaxes(peaks, 0 ,1)
    
    # classify the peaks
    for x in peaks:
        if(x - window / 2 > 0) and (x + window /2 < len(dval)):
            start = int(x - window/2)
            end = int(x + window/2)
            
            dataWindow = copy.deepcopy(dval[start: end])
            dataWindow = np.float32(dataWindow)
            

            dmin = dataWindow.min()
            dmax = dataWindow.max()
    
            
            if dmax > dmin * -1:
                dataWindow *= 1/dmax
            else:
                dataWindow *= -1/dmin
            
            
            dataWindow = np.reshape(dataWindow, (1, window, 1))
            
            prediction = autoencoder.predict(dataWindow)           
    
            zeros[x] = prediction[0][0]
            pWave[x] = prediction[0][1] 
            rSpike[x] = prediction[0][2]
            tWave[x] = prediction[0][3]
            noise[x] = prediction[0][4]
            other[x] = prediction[0][5]

    	        
    return zeros, pWave, rSpike, tWave, noise, other

def correlate(values, rPos):
    lag = 3
    start = 15
    end = 45
    
    result = np.empty((len(rPos), len(rPos)))
    result[:] = 1
    
    for r in range(len(rPos)-1):
        for s in range(r+1, len(rPos)):
            correlation = []
            for l in range(-lag, lag):                
                correlation.append(np.corrcoef(values[rPos[r]-start:rPos[r]+end], values[rPos[s]-start+l:rPos[s]+end+l])[0,1])
            result[r,s] = np.max(correlation)
            result[s,r] = result[r,s]
    
    score = []
    
    
    
    
    
    for c in range(2, np.min([10, len(rPos)])):
        kmean = cluster.KMeans(n_clusters=c, random_state=170).fit_predict(result)
        score.append(metrics.silhouette_score(result, kmean, metric='euclidean'))
    
    score = np.array(score)
    
    if len(score) == 0:
        score = np.nan
        scoreMax = np.nan
    
    else:
        scoreMax = np.max(score)
        score = score[0]    
        
    if len(rPos) > 0 and len(result) > 5:
        tree = scipy.cluster.hierarchy.linkage(result, method="average", metric="euclidean")
        tree = tree[-5:, 2]
    else:
        tree = [9,9,9,9,9]
        
    return score, scoreMax, tree


def deep(data):
    
    zeros = np.zeros(len(data))
    pWave = np.zeros(len(data))
    rSpike = np.zeros(len(data))
    tWave = np.zeros(len(data))
    noise = np.zeros(len(data))
    other = np.zeros(len(data))
    
    
    #return peak_indexes
    peaks = scipy.signal.argrelextrema(data**2, np.greater, order=40)
    peaks = np.array(peaks)

    print("classify R Spikes")
    zeros, null, rSpike, null, null, other = classify(peaks, data)
    
    
    print("classify Rest")
    peaks = scipy.signal.argrelextrema(data, np.greater, order=20)
    peaks = np.array(peaks)
    
    zeros, pWave, null, tWave, noise, null = classify(peaks, data)
    
    p = np.where(pWave>pthreshold)[0]
    r = np.where((rSpike+other)>rthreshold)[0]    
    t = np.where(tWave>tthreshold)[0]
    o = np.where(other>othreshold)[0]
    
    
    if(debug == 1):
        plt.figure(figsize=(24, 10))
        plt.plot(data)
        plt.plot(pWave)
        #plt.plot(-zeros)
        plt.plot(rSpike)
        plt.plot(tWave)
        plt.plot(-noise)
        plt.plot(-other)
        
        #plt.xlim(7500, 10000)
        plt.show()
    
    return p, r, t, noise, o
    

def analyze(file, solution):
    def relRR(values, order):
        values = np.array([sum(values[current: current+order]) for current in range(len(values)-order)])
        if(order <= 1):
            relRRinterval = np.ediff1d(values)*2/(values[:-order]+values[order:])    
        else:
            relRRinterval = np.ediff1d(values[:-order+1])*2/(values[:-order]+values[order:])    
        return relRRinterval
    
    print(file)
    # Read .mat format
    
    signals = []
    
    file = os.path.splitext(file)[0]
    
    
    #%% Reads the signal and puts values in x and chanel in y
    signals, fields = wfdb.srdsamp(file, sampfrom=0,channels=[0])
    
    fs = fields["fs"] ## sampling rate    
    
    

    
    print("resample")
    #from scipy import signal
    #signals = signal.resample(signals, int(len(signals)*300/fs), axis=0)
    # resample and get back the data layer, discard postions. Who mad this %"§&§?
    signals = wfdb.processing.resample_sig(signals, fs, fs_target)[0] 
    

    signals = np.array(signals)
    
    signals = np.swapaxes(signals, 0, 1)[0]
    
    


    
    # denoise signal
    data = signals
        
    data = np.array(dwt.cleanUp(data))

    data = np.float32(data)
    
    #create output array


    
    p,r,t,noise,o = deep(data)
    
    mp, mr, mt, mnoise, mo = deep(-data)
    
    if(len(mp) > len(p) and len(mt) > len(t)):
        t = mt
        p = mp
        r = mr
        o = mo
        noise = mnoise
        
        flipped = 1
        
        if(debug == 1):
            print("taking inverted data")
        
    else:
        
        flipped = 0
        
        
        
        
    # calculate noisy areas
    noisewin = pd.DataFrame(noise).rolling(300, center=True, win_type="blackman").sum()
    noisewin = np.array(noisewin)[:,0]
    noisewin = np.nan_to_num(noisewin)
    
    noiseTidx = np.where(noisewin >= nthreshold)
    
    noiseT = np.zeros(len(data))
    noiseT[noiseTidx] = 1
    
    relNoiseT = len(noiseTidx[noiseTidx==True]) / len(signals)
    
    if(debug == 1):
        print("relative noise = {}".format(relNoiseT))

    
    if(len(o) > 0 and len(r) > 0):
        if(len(r) > len(o)):    
            relO = len(o)/(len(r) - len(o))
        else:
            relO = 1
    else:
        if len(o) == 0:
            relO = 0    
        else:
            relO = np.nan
        
    
    if(debug == 1):
        print("relative count of other beats: {}".format(relO))
    

    score, scoreMax, tree = correlate(data, r)
    
    # get intervals between r spikes
    rrInterval  = np.ediff1d(r)
    if(len(r)>2):
        hr = stats.trim_mean(rrInterval, 0.25)
    else:
        hr = np.nan #0.0000001
    
    if(debug == 1):
        print("passes?")
    
    
    # remove entrys longer than 150% Median
    amountDeleted = len(rrInterval[rrInterval >= hr * 1.5]) # Number of deleted intervals
    rrInterval = rrInterval[rrInterval < hr * 1.5]

    hr = stats.trim_mean(rrInterval, 0.25)
    
    if(debug == 1):
        print("hr = {}".format(hr))

    
    if(debug == 1):
        
        plt.plot(rrInterval)
        plt.show()
    

    
    hrEntropy = []
    
    baseRelRRinterval = relRR(rrInterval, 1)
    
    # get key statistics and make them a numpy array
    
    # dont devide by 0  
    
    #%% find extrasystoles
    filteredEs = np.array(dwt.checkEs(baseRelRRinterval))
    
    if(len(r)>0):
        relFilteredEs = filteredEs/len(r)
    else:
        relFilteredEs = filteredEs
    
    cleanRelRRinterval = np.array(dwt.checkEs(baseRelRRinterval, "destroy"))
    entropy0 = np.std(cleanRelRRinterval)
    
    unusualBeats = len(cleanRelRRinterval[np.abs(cleanRelRRinterval-1)>.50])/hr
    
    print(entropy0)
    
    #%% CREATE RELATIVE RR INTERVALS
    for rng in range(1,10):
        print(rng)
        
        relRRinterval = relRR(rrInterval, rng)
        entropy = np.std(relRRinterval)
        print(entropy)
        hrEntropy.append(entropy)
        
        if rng == 1:
            baseRelRRinterval = relRRinterval
            # return baseRelRRinterval
    
    hrEntropy = np.array(hrEntropy)
    
    
    #%% CALCULATE PR INTERVAL
    print("PR INTERVAL")
    if(hr > 1):
        PRinterval = find_array(p, r, int(hr), -1)

        PRinterval = PRinterval[PRinterval < hr *.5]
        relPRinterval = relRR(PRinterval, 1)
    else:
        PRinterval = [np.nan]
        relPRinterval = [np.nan]



    if(debug == 1):
        plt.plot(PRinterval)
        plt.show()
        
        plt.figure(figsize=(14, 14))
        plt.scatter(relPRinterval[:-1],relPRinterval[1:])
        plt.plot(relPRinterval[:-1],relPRinterval[1:])
        plt.xlim(-2,2)
        plt.ylim(-2,2)
        plt.show()
        
    #%% CALCULATE RT INTERVAL
    if not np.isnan(hr):
        RTinterval = find_array(t, r, int(hr), 1)
        RTinterval = RTinterval[RTinterval < hr]
        relRTinterval = relRR(RTinterval, 1)
    else:
        RTinterval = np.nan
        relRTinterval = np.nan
    # ignore measurements where RT > median
    
    


    if(debug == 1 and not np.all(np.isnan(RTinterval))):
        plt.plot(RTinterval)
        plt.show()
        
        plt.figure(figsize=(14, 14))
        plt.scatter(relRTinterval[:-1],relRTinterval[1:])
        plt.plot(relRTinterval[:-1],relRTinterval[1:])
        plt.xlim(-2,2)
        plt.ylim(-2,2)
        plt.show()
        
    # calculate the centage of qrs complexes that have a t or/and p wave
    print("len is {}".format(len(rrInterval)))
    
    #%% CALCULATE METRICS ON SLOW AND FAST PARTS
    if(len(rrInterval) > 7):
        pWavePrevalence = len(p)/len(r)
        tWavePrevalence = len(t)/len(r)
        
        numTach = len(rrInterval[rrInterval < 150])/len(r)
    
        salve3 = np.array(dwt.running_sum(rrInterval, 3))
        salve5 = np.array(dwt.running_sum(rrInterval, 5))
        
        numTach3 = len(salve3[salve3 < 150*3])/len(r)    
        numTach5 = len(salve5[salve5 < 150*5])/len(r)
        
        numBrad = len(rrInterval[rrInterval > 450])/len(r)
        
        salve3 = np.array(dwt.running_sum(rrInterval, 3))
        salve5 = np.array(dwt.running_sum(rrInterval, 5))
        
        numBrad3 = len(salve3[salve3 < 450*3])/(len(rrInterval)- 2)
        numBrad5 = len(salve5[salve5 < 450*5])/(len(rrInterval) - 4)
        
    # SET ALL 0 IF TOO FEW SAMPLES
    else:
        pWavePrevalence = np.nan
        tWavePrevalence = np.nan
        
        numTach = np.nan
        numTach3 = np.nan
        numTach5 = np.nan
        
        numBrad = np.nan
        numBrad3 = np.nan    
        numBrad5 = np.nan
    
    # normalize hrEntropy with the last entry
    if(hrEntropy[-1] > 0):
        relhrEntropy = hrEntropy / hrEntropy[-1]
    else:
        relhrEntropy = hrEntropy


    #%% GENERATE PERCENTILES         
    percentile = np.array([0.1, 0.2, .25, 0.3, 0.4, 0.5, 0.6, 0.7, .75, 0.8, 0.9])
    
    # checks if values are nan
    def describe(values):
        if(not np.all(np.isnan(values))):
            values      = np.array(pd.DataFrame(np.array(values)).describe(percentile ))[:,0]
            qdiff       = values[12]-values[6]
            return np.append(values, qdiff)
        else:
            values      = np.empty(17)
            values[:]   = np.nan
        
        return values
    
        
 

      
    rrInterval      = describe(rrInterval)
    
    relRRinterval   = describe(relRRinterval)
    PRinterval      = describe(PRinterval)
    relPRinterval   = describe(relPRinterval)
    
    RTinterval      = describe(RTinterval)
    relRTinterval   = describe(relRTinterval)
    
    if not np.all(np.isnan(rrInterval)) and not np.isnan(hr):
        medianRRinterval = describe(rrInterval/hr)
    else:
        medianRRinterval = describe(np.nan)
        
        
    #%% FUSE EVERYTHING TOGETHER
    ret = np.array([file, solution, flipped, amountDeleted, hr, pWavePrevalence, tWavePrevalence, entropy0, score, scoreMax, numTach, numTach3, numTach5, numBrad, numBrad3, numBrad5, unusualBeats, relNoiseT, relO])
    ret = np.concatenate([ret, filteredEs, relFilteredEs, tree, hrEntropy, relhrEntropy, PRinterval, relPRinterval, rrInterval, medianRRinterval, relRRinterval, RTinterval, relRTinterval])
    
    ret[2:] = nanTo9(np.float32(ret[2:]))
    
    print("finished")
    #print(ret)
    return ret

#%% writes data into csv file with mode m
def write(data, file, m="a"):        
    fd = open(file, m)
    writer = csv.writer(fd, delimiter=",", lineterminator="\n")  
    writer.writerow(data)

#%% Writes CSV HEADER
def header(file):
    
    
    write(["file", "lable", "flipped", "skipped beats", "hr", "pWavePrevalence", "tWavePrevalence", "entropy 0", 
       "correlation score", "max correlation score", "number tachycardic episodes", "tachycardic 3 salves", "tachycardic 5 salves", "number bradycardic episodes", "bradycardic 3 salves", "bradycardic 5 salves", 
       "unusual beats", "percentage of noise recordingtime", "relative other beats", "simpleEs", "nSimpleEs", "nDoubleEs", "nTripleEs", "compedEs", "cDoubleEs", "cTripleEs", "es count",
       "rel simpleEs", "rel nSimpleEs", "rel nDoubleEs", "rel nTripleEs", "rel compedEs", "rel cDoubleEs", "rel cTripleEs", "rel es count",
       "tree 1","tree 2","tree 3","tree 4","tree 5",
       "hrEntropy 1","hrEntropy 2", "hrEntropy 3", "hrEntropy 4", "hrEntropy 5", "hrEntropy 6", "hrEntropy 7", "hrEntropy 8", "hrEntropy 9",
       "RhrEntropy 1","RhrEntropy 2", "RhrEntropy 3", "RhrEntropy 4", "RhrEntropy 5", "RhrEntropy 6", "RhrEntropy 7", "RhrEntropy 8", "RhrEntropy 9",
       "PR Interval", "PR mean", "PR std", "PR min", "PR 10%", "PR 20%", "PR 25%", "PR 30%", "PR 40%","PR 50%","PR 60%","PR 70%","PR 75%","PR 80%","PR 90%","PR max", "PR qdiff", 
       "rel PR Interval", "rel PR mean", "rel PR std", "rel PR min", "rel PR 10%", "rel PR 20%", "rel PR 25%", "rel PR 30%", "rel PR 40%","rel PR 50%","rel PR 60%","rel PR 70%","rel PR 75%","rel PR 80%","rel PR 90%","rel PR max","rel PR qdiff", 
       "RR Interval", "RR mean", "RR std", "RR min", "RR 10%", "RR 20%", "RR 25%", "RR 30%", "RR 40%","RR 50%","RR 60%","RR 70%", "RR 75%","RR 80%","RR 90%","RR max", "RR qdiff", 
       "median RR Interval", "median RR mean", "median RR std", "median RR min", "median RR 10%", "median RR 20%", "median RR 25%", "median RR 30%", "median RR 40%","median RR 50%","median RR 60%","median RR 70%","median RR 75%","median RR 80%","median RR 90%","median RR max","median RR qdiff", 
       "rel RR Interval", "rel RR mean", "rel RR std", "rel RR min", "rel RR 10%", "rel RR 20%", "rel RR 25%", "rel RR 30%", "rel RR 40%","rel RR 50%","rel RR 60%","rel RR 70%","rel RR 75%","rel RR 80%","rel RR 90%","rel RR max","rel RR qdiff", 
       "RT Interval", "RT mean", "RT std", "RT min", "RT 10%", "RT 20%", "RT 25%", "RT 30%", "RT 40%","RT 50%","RT 60%","RT 70%","RT 75%","RT 80%","RT 90%","RT max", "RT qdiff", 
       "rel RT Interval", "rel RT mean", "rel RT std", "rel RT min", "rel RT 10%", "rel RT 20%", "rel RT 25%","rel RT 30%", "rel RT 40%","rel RT 50%","rel RT 60%","rel RT 70%","rel RT 75%","rel RT 80%","rel RT 90%","rel RT max", "RT qdiff", 
      ], file, 'w')

## Loads the Tensorflow Model
autoencoder = load_model("models/model_2_30.h5")

        
#%%%    MAIN ROUTINE 
if len(sys.argv) >= 2:
    debug = 0
    #file = 'features.csv'
    dataFile = sys.argv[1]
    #header(file)
    
    testvar = np.float32(analyze(dataFile, -1)[2:])
    
    testvar = nanTo9(testvar)
    
    testvar = np.reshape(testvar, (1,175))
    
    rf = joblib.load('tree.pkl') 
    pred = rf.predict(testvar)
    write([dataFile, pred[0]], "answers.txt", m="a")
    K.clear_session()
    
else:
    #%% RUN 1 DEBUG EXAMPLE
    
    if os.environ.get('subsample') is not None:
        os.chdir("training2017")
        batch = int(os.environ['subsample'])
        print(batch)
        debug = 0
        
        file = '../2part_{}.csv'.format(batch)
        header(file)
        
        # Read the Lables and filenames
        fReader     = open('REFERENCE.csv','r')
        lablereader = csv.reader(fReader, delimiter=',', quotechar='|')
        solution    = []
        
        for row in lablereader:
            solution.append(row)         
        
        
        # Multiprocessing
        
        start   = int( math.floor( len(solution)/20 * (int(batch)-1)))
        end     = int( math.floor( len(solution)/20 * int(batch)) -1)
        
        if batch == 20:
            end += 1
    
        print("starting at {}".format(start))
        print("ending at {}".format(end))
    
        for s in solution[start:end]:
            write(analyze(s[0], s[1]), file, m="a")
    else:
        if validation == 0:
            os.chdir("validation")
            debug = 1
            file = "A04735"
            solution = -1
            testvar = np.float32(analyze(file, solution)[2:])
            
            testvar = nanTo9(testvar)
            
            testvar = np.reshape(testvar, (1,175))
            
            rf = joblib.load('../tree.pkl') 
            pred = rf.predict(testvar)
            print(pred)
        else:
            debug = 0
            fReader     = open('validation/REFERENCE.csv','r')
            lablereader = csv.reader(fReader, delimiter=',', quotechar='|')
            solution    = []
            
            for row in lablereader:
                solution.append(row)         
        
            for s in solution:
                testvar = np.float32(analyze("validation/"+s[0], -1)[2:])                
                testvar = nanTo9(testvar)                
                testvar = np.reshape(testvar, (1,175))
                
                rf = joblib.load('tree.pkl') 
                pred = rf.predict(testvar)
                write([s[0], pred[0]], "answers.txt", m="a")      
