""" 

Author: Mondejar Guerra

VARPA

University of A Coruna

April 2017

Description: Train and evaluate mitdb with interpatient split (train/test)

Uses my own model clasifier with weights for imbalanced class

"""

import numpy as np

import matplotlib.pyplot as plt

import os

import csv

import pickle

import numpy as np

import matplotlib.pyplot as plt

import tensorflow as tf

import collections

from tensorflow.contrib.learn.python.learn.estimators import model_fn as model_fn_lib

tf.logging.set_verbosity(tf.logging.INFO)

def compute_accuracy(m):

  # Accuracy by column

  classes = m.shape[0]

  acc = np.zeros(classes)

  acc_global = 0

  for c in range(0, classes):

    if sum(m[:,c]) > 0:

      acc[c] = float(m[c,c]) / float(sum(m[:,c]))

    acc_global = acc_global + m[c,c]

    #print ('acc ' + str(c) + ': ' + str(acc[c]))

  acc_global = float(acc_global) / float(sum(sum(m)))

  #print ('global acc = ' + str(acc_global))

  return acc, acc_global

def load_data(output_path, window_size, compute_RR_interval_feature, compute_wavelets, binary_problem):

  extension = '_' + str(window_size)

  if compute_wavelets:

      extension = extension + '_' + 'wv'

  if compute_RR_interval_feature:

      extension = extension + '_' + 'RR'

  extension = extension + '.csv'

  # Load training and eval data

  train_data = np.loadtxt(output_path + 'train_data' + extension, delimiter=",", dtype=float)

  train_labels =  np.loadtxt(output_path + 'train_label' + extension, delimiter=",",  dtype=np.int32)

  eval_data = np.loadtxt(output_path + 'eval_data' + extension, delimiter=",", dtype=float)

  eval_labels = np.loadtxt(output_path + 'eval_label' + extension, delimiter=",",  dtype=np.int32)

  if binary_problem: 

    # Uses only two classes 

    #   [0]: N (0)

    #   [1]: SVEB, VEB, F, Q (1, 2, 3, 4)

    for i in range(0, len(train_labels), 1):

      if train_labels[i] > 0:

        train_labels[i] = 1

    for i in range(0, len(eval_labels), 1):

      if eval_labels[i] > 0:

        eval_labels[i] = 1        

  return (train_data, train_labels, eval_data, eval_labels)

# normalize data features: wave & RR intervals...

def normalize_data(train_data, eval_data):

  feature_size = len(train_data[0])

  if compute_RR_interval_feature:

    feature_size = feature_size - 4

  max_wav = np.amax(np.vstack((train_data[:, 0:feature_size], eval_data[:, 0:feature_size])))

  min_wav = np.amin(np.vstack((train_data[:, 0:feature_size], eval_data[:, 0:feature_size])))

  train_data[:, 0:feature_size] = ((train_data[:,0:feature_size] - min_wav) / (max_wav - min_wav))

  eval_data[:, 0:feature_size] = ((eval_data[:,0:feature_size] - min_wav) / (max_wav - min_wav))

  #Norm last part feature: RR interval 

  if compute_RR_interval_feature:

    max_rr = np.amax(np.vstack((train_data[:, feature_size:], eval_data[:, feature_size:])))

    min_rr = np.amin(np.vstack((train_data[:, feature_size:], eval_data[:, feature_size:])))

    train_data[:, feature_size:] = ((train_data[:, feature_size:] - min_rr) / (max_rr - min_rr))

    eval_data[:,  feature_size:] = ((eval_data[:, feature_size:] - min_rr) / (max_rr - min_rr))

  return (train_data, eval_data)

def my_model_fn(features, targets, mode, params):

  """Model function for Estimator."""

  targets_onehot = tf.one_hot(indices = targets, depth=params["num_classes"], on_value = 1)

  # Connect the first hidden layer to input layer

  # (features) with relu activation

  #first_hidden_layer = tf.contrib.layers.relu(features, 10)

  first_hidden_layer = tf.contrib.layers.fully_connected(features, params["h1"])

  # tf.nn.conv1d

  second_hidden_layer = tf.contrib.layers.fully_connected(first_hidden_layer, params["h2"])

  third_hidden_layer = tf.contrib.layers.relu(second_hidden_layer, params["h3"])

  # Connect the output layer to second hidden layer (no activation fn)

  output_layer = tf.contrib.layers.linear(third_hidden_layer, params["num_classes"])

  if mode == 'train' and params["weight_imbalanced"]:

    weights_tf = tf.constant(params["weights"])

  else:

    weights_tf = tf.ones([features.shape[0].value], tf.float32) 

  loss = tf.losses.softmax_cross_entropy(targets_onehot, output_layer, weights=weights_tf)

  train_op = tf.contrib.layers.optimize_loss(

    loss=loss,

    global_step=tf.contrib.framework.get_global_step(),

    learning_rate=params["learning_rate"],

    optimizer="SGD")

  correct_prediction = tf.equal(tf.argmax(targets_onehot, 1), tf.argmax(output_layer, 1))

  accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))   

  eval_metric_ops = {

    "accuracy": accuracy

        #tf.metrics.accuracy(targets_onehot, output_layer)

        #"rmse": tf.metrics.root_mean_squared_error(

        #    tf.cast(targets, tf.float64), predictions)

  }

  return model_fn_lib.ModelFnOps(

    mode=mode,

    predictions=output_layer,#predictions_dict,

    loss=loss,

    train_op=train_op,

    eval_metric_ops=eval_metric_ops)

def main():

  window_size = 160

  compute_RR_interval_feature = True

  compute_wavelets = True

  dataset = '/home/mondejar/dataset/ECG/mitdb/'

  output_path = dataset + 'm_learning/'

  binary_problem = False

  weight_imbalanced = True

  # 0 Load Data

  train_data, train_labels, eval_data, eval_labels = load_data(output_path, window_size, compute_RR_interval_feature, compute_wavelets, binary_problem)

  # 1 TODO Preprocess data? norm? if RR interval, last 4 features are pre, post, local and global RR

  # Apply some norm? convolution? another approach?

  normalize = False

  if normalize:

    train_data, eval_data =  normalize_data(train_data, eval_data)

  # 2 Create my own model

  # Imbalanced class: weights

  # https://www.tensorflow.org/api_guides/python/contrib.losses

  # Learning rate for the model

  LEARNING_RATE = 0.001

  if binary_problem:

    num_classes = 2

  else:

    num_classes = 5

  # Set model params

  count = collections.Counter(train_labels)

  total = 0

  max_class = 0

  for c in range(0,num_classes):

    total = count[c] + total

    if count[c] > max_class:

      max_class = count[c]

  class_weight = np.zeros(num_classes)

  for c in range(0,num_classes):

    if count[c] > 0:

      #class_weight[c] = 1- float(count[c]) / float(total)

      class_weight[c] = float(max_class) / float(count[c]) # the class with more instance will have weight = 1, and the others x times ...

  # TODO give more weigth to anomaly classes? We want to detect always these bad anomalies

  weights = np.zeros((len(train_labels)), dtype='float')

  for i in range(0,len(train_labels)):

    weights[i] = class_weight[train_labels[i]]

  hn_1 = [128, 64, 32]

  hn_2 = [64, 32, 16]

  hn_3 = [32, 16, 8]

  steps = [500, 1000, 2000, 8000]

  for h1 in hn_1:

    for h2 in hn_2:

      for h3 in hn_3:

        for s in steps:

          model_params = {

            "learning_rate": LEARNING_RATE, 

            "num_classes": num_classes, 

            "weights": weights, 

            "weight_imbalanced": weight_imbalanced, 

            "h1": h1,

            "h2": h2, 

            "h3": h3}

          nn = tf.contrib.learn.Estimator(model_fn=my_model_fn, params=model_params)

          def get_train_inputs():

            x = tf.constant(train_data)

            y = tf.constant(train_labels)

            return x, y

          # Fit

          nn.fit(input_fn=get_train_inputs, steps=s)

          # Score accuracy

          def get_test_inputs():

            x = tf.constant(eval_data)

            y = tf.constant(eval_labels)

            return x, y

          ev = nn.evaluate(input_fn=get_test_inputs, steps=1)["accuracy"]

          # Compute the matrix confussion

          predictions = list(nn.predict(input_fn=get_test_inputs))

          confusion_matrix = np.zeros((num_classes,num_classes), dtype='int')

          for p in range(0, len(predictions), 1):

            ind_p = np.argmax(predictions[p])

            confusion_matrix[ind_p][eval_labels[p]] = confusion_matrix[ind_p][eval_labels[p]] + 1

          acc, acc_g = compute_accuracy(confusion_matrix)

          np.savetxt('nn_' + str(h1) + '_' + str(h2) + '_' + str(h3) + '_' + str(s) + '_cm.txt', confusion_matrix, fmt='%-7.0f')    

          np.savetxt('nn_' + str(h1) + '_' + str(h2) + '_' + str(h3) + '_' + str(s) + '_acc.txt', acc, fmt='%-7.2f')    

if __name__ == "__main__":

  main()