# -*- coding: utf-8 -*-
"""
    File name:      anomaly_detection_ecg.py
    Author:         Tetsuro Tatsuoka
    Date created:   04/12/2018
    Last modified:  04/12/2018
    Version:        1.0.0
    Python version: 3.6.3
    Description:    Anomaly detection for ECG by using autoencoder

    Revision history:
        1.0.0   04/12/2018  - Initial coding.
"""
import os
import time
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
import seaborn as sns
import pickle

from keras import backend as K
from keras.layers import Input, Dense
from keras.models import Model
from keras import regularizers

from keras.utils.vis_utils import plot_model

sns.reset_orig()

def conv_norm(x):
    """ Convert to normalized data
    """
    ms = MinMaxScaler()
    return ms.fit_transform(pd.DataFrame(x))

def conv_subseq(x, w):
    """ Convert to subsequence time series
    """
    subseq_size = len(x) - w + 1
    subseq = np.zeros((subseq_size, w))
    for i in range(subseq_size):
        start_idx = i
        end_idx = w + i
        subseq[i,:] = x[start_idx:end_idx].flatten()
    return subseq

def anomaly_detection_ecg(
        folder_path='AI',
        train_file_name='wave_data_234_148000_152000.pkl',  # Train pickle file
        test_file_name='wave_data_234_152000_156000.pkl',   # Test pickle file
        graph_save=False,           # Graph image is saved or displayed
        save_file_prefix='',        # Prefix string for save file
        half_sr_mode=False,         # Sampling rate (True:180Hz, False:360Hz)
        in_out_size=100,            # Subsequence time series size
        hidden_size=160,            # Hidden layer unit sizze
        encoder_act='relu',         # Activation function for encoder
        decoder_act='relu',         # Activation function for decoder
        reg_l1=10e-7,               # L1 regularization coefficient value
        reg_l2=0.0,                 # L2 regularization coefficient value
        opt_alg='adadelta',         # Optimizer algorithm
        loss_func='binary_crossentropy', # loss function
        batch_size=None,            # Size of mini-batch gradient descent
        epoch_num=3000):            # Number of maximum training epochs
    """ Detect anomaly ECG waveform
    """

    proc_start = time.time()
    
    # Arias
    window_size = in_out_size
    encoding_dim = hidden_size
    
    # Setup data
    FOLDER_PATH = '.' + os.sep + folder_path
    if not os.path.exists(FOLDER_PATH):
        os.mkdir(FOLDER_PATH)
    
    DATA_FOLDER_PATH = 'ecg_data'    # Saved folder of pickle data files
    TRAIN_FILE_NAME = os.path.join(DATA_FOLDER_PATH, train_file_name)
    TEST_FILE_NAME = os.path.join(DATA_FOLDER_PATH, test_file_name)
    
    with open(TRAIN_FILE_NAME, 'rb') as f:
        train_dat = (pickle.load(f)).iloc[:,0]  # Only channel 1
    
    with open(TEST_FILE_NAME, 'rb') as f:
        test_dat = (pickle.load(f)).iloc[:,0]   # Only channel 1
    
    # Decimate data in half sampling rate mode
    if half_sr_mode:
        train_dat = pd.Series(np.array(train_dat[::2]) +
                              np.array(train_dat[1::2])) / 2
        test_dat = pd.Series(np.array(test_dat[::2]) +
                             np.array(test_dat[1::2])) / 2
    
    # Normalization
    train_dat_norm = conv_norm(train_dat)
    test_dat_norm = conv_norm(test_dat)
    
    # Generate subsequence time series
    train_dat_s = conv_subseq(train_dat_norm, window_size)
    test_dat_s = conv_subseq(test_dat_norm, window_size)
    
    x_train = train_dat_s.astype('float32')
    x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
    
    # Plot training data
    plt.figure()
    plt.plot(train_dat_norm)
    plt.suptitle('Training data)')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'train_data.png'))
        plt.close()
    else:
        plt.show()
   
    # Plot test data
    plt.figure()
    plt.plot(test_dat_norm)
    plt.suptitle('Test data')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'test_data.png'))
        plt.close()
    else:
        plt.show()
    
    """ Detecting anomaly points """
    # Define model
    input_layer = Input(shape=(window_size,), name='Input')
    encoded = Dense(encoding_dim, activation=encoder_act, name='Encoder',
                    activity_regularizer=regularizers.l1(reg_l1))(input_layer)
    decoded = Dense(window_size, activation=decoder_act, name='Decoder')(encoded)
    autoencoder = Model(inputs=input_layer, outputs=decoded)
    
    # Generate model
    autoencoder.compile(optimizer=opt_alg, loss=loss_func)

    # Plot model structure    
    save_filename = os.path.join(FOLDER_PATH, save_file_prefix + 'model.png')
    plot_model(autoencoder, to_file=save_filename, show_shapes=True)
    
    # Train model
    history = autoencoder.fit(x_train, x_train,
                              epochs=epoch_num,
                              batch_size=batch_size,
                              verbose=1,
                              shuffle=True)
    
    # Print last score
    score = autoencoder.evaluate(x_train, x_train, verbose=0)
    print('Training loss:', score)

    # Plot loss curve
    plt.figure()
    plt.plot(history.history['loss'])
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.suptitle('Training loss curve')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'loss.png'))
        plt.close()
    else:
        plt.show()
    
    # Calculate score (distance)
    decoded_wave = autoencoder.predict(test_dat_s)
    dist = np.sqrt(np.sum((decoded_wave - test_dat_s)**2, axis=1))
    
    # Plot score (distance)
    plt.figure()
    plt.plot(dist)
    plt.suptitle('Score data')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'score.png'))
        plt.close()
    else:
        plt.show()
    
    idx = 260
    plt.figure()
    plt.plot(test_dat_s[idx,:])
    plt.plot(decoded_wave[idx,:])
    plt.suptitle('Waveform sample (Index: ' + str(idx) + ')')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'wave_ex.png'))
        plt.close()
    else:
        plt.show()
    
    # Max score
    max_idx = np.where(dist==max(dist))[0][0]
    plt.figure()
    plt.plot(test_dat_s[max_idx,:])
    plt.plot(decoded_wave[max_idx,:])
    plt.suptitle('Waveform sample (Index: ' + str(max_idx) + ') (Score:max)')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'wave_max.png'))
        plt.close()
    else:
        plt.show()
    
    # Detecting over threshold
    ANOMALY_TH = max(dist) * 0.8    # Threshold for score
    INHIBIT_LEN = 500               # Inhibit period
    an_idx = (np.where(dist>ANOMALY_TH))[0]
    an_idx = an_idx[np.r_[True, (an_idx[1:]-an_idx[:-1] > INHIBIT_LEN)]]
    
    # Plot anomaly waveform and socore
    PRE_LEN = 500                   # Left margin
    for idx in an_idx:
        s_idx = idx - PRE_LEN
        e_idx = idx + INHIBIT_LEN
        plt.figure()
        plt.subplot(211)
        plt.plot(test_dat_norm[s_idx:e_idx])
        plt.subplot(212)
        plt.plot(dist[s_idx:e_idx])
        plt.suptitle('Score data (Index: ' + str(s_idx) + ':' + str(e_idx) + ')')
    if graph_save:
        plt.savefig(os.path.join(FOLDER_PATH, save_file_prefix + 'score_exp.png'))
        plt.close()
    else:
        plt.show()
        
    # Save weights
    save_filename = os.path.join(FOLDER_PATH, save_file_prefix + 'model_weights.h5')
    autoencoder.save_weights(save_filename)
    
    K.clear_session()

    elapsed_per = (time.time() - proc_start)
    print('Elapsed time:', elapsed_per)


if __name__ == "__main__" :

    # Down sampling
    for i in range(5):
        anomaly_detection_ecg(folder_path='res_180sps',
                              graph_save=True,
                              save_file_prefix=str(i)+'_',
                              half_sr_mode=True,
                              in_out_size=50,
                              hidden_size=80,
                              encoder_act='relu',
                              decoder_act='relu',
                              reg_l1=10e-7,
                              reg_l2=0.0,
                              opt_alg='adadelta',
                              batch_size=None,
                              epoch_num=5000)
    # Without down sampling
    for i in range(5):
        anomaly_detection_ecg(folder_path='res_360sps',
                              graph_save=True,
                              save_file_prefix=str(i)+'_',
                              half_sr_mode=False,
                              in_out_size=100,
                              hidden_size=160,
                              encoder_act='relu',
                              decoder_act='relu',
                              reg_l1=10e-7,
                              reg_l2=0.0,
                              opt_alg='adadelta',
                              batch_size=None,
                              epoch_num=5000)
