# -*- coding: utf-8 -*-
"""
    File name:      load_data_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:    Load specified ECG data from MIT-BIH Arrhythmia Database
                    Remove baseline drift artifact
    Reference:
        MIT-BIH Arrhythmia Database
        https://physionet.org/physiobank/database/mitdb/
        https://www.physionet.org/physiobank/database/html/mitdbdir/mitdbdir.htm#toc
        
        WFDB(Waveform database) Application
        https://github.com/MIT-LCP/wfdb-python
        https://github.com/MIT-LCP/wfdb-python/blob/master/demo.ipynb

    Revision history:
        1.0.0   04/12/2018  - Initial coding.
"""
import os
import numpy as np
import pandas as pd
import pickle
import wfdb
import collections
import time

RECORD_NAME = '234'

DB_DIR = 'mitdb/'
SAVE_FOLDER_PATH = '.' + os.sep + 'ecg_data'

def load_data_ecg(rec_name=RECORD_NAME, begin_idx=0, end_idx=0, N=200):
    """
        Load specified ECG data from MIT-BIH Arrhythmia Database
        and remove baseline drift artifact.
    
        Parameters:
            rec_name    Record name 
            begin_idx   Begin index
            end_idx     End index
            N           Range of moving averaging for removing drift baseline
        Returns:
            rec         wfdb.io.record.Record object (Original waveform)
            ann         wfdb.io.annotation.Annotation object
            rec2        wfdb.io.record.Record object (After removing drift)
                        (Data length is N * 200 fewer than original)
        Usage:
            1)  load_data_ecg('234')  # for entire data length
            2)  load_data_ecg('234', begin_idx=148000, end_idx=152000)
    """
    # Data file hostd on Physiobank
    if end_idx == 0:
        rec = wfdb.rdrecord(rec_name, pb_dir = DB_DIR)
    else:
        rec = wfdb.rdrecord(rec_name, pb_dir = DB_DIR, sampfrom=begin_idx - N, sampto=end_idx + N)
    print('\n\nRecord data:')
    print(vars(rec))
    
    # Draw waveforms
    wfdb.plot_wfdb(record=rec, title='Record '+rec_name) 
    
    # Read header file
    hdr = wfdb.rdheader(rec_name, pb_dir=DB_DIR)
    print('\n\nHeader data:')
    print(vars(hdr))

    # Annotation file
    if end_idx == 0:
        ann = wfdb.rdann(rec_name, 'atr', pb_dir=DB_DIR)
    else:
        ann = wfdb.rdann(rec_name, 'atr', pb_dir=DB_DIR, sampfrom=begin_idx - N, sampto=end_idx + N)
        ann.sample = ann.sample - (begin_idx - N)

    print('\n\nAnnotation data:')
    print(ann.ann_len)
    print('\n\nIndex of anomaly data:')
    print(collections.Counter(ann.symbol))
    ano = ann.sample[np.where(np.array(ann.symbol) != 'N')]
    astr = np.array(ann.symbol)[np.where(np.array(ann.symbol) != 'N')] 
    print("Annotation: " + str(len(ano)) + ' beats')
    for (i, a) in zip(ano, astr):
        print(str(i) + "\t" + a)        

    # Draw waveforms    
    wfdb.plot_wfdb(record=rec, annotation=ann,
                   title='Record ' + rec_name + ' with annotation',
                   time_units='seconds')
    
    # Removed baseline drift artifact
    M = 2 * N + 1
    dat_len = rec.p_signal.shape[0] - 2 * N
    ch_num = rec.p_signal.shape[1]

    b = np.ones(M) / float(M)
    start_time = time.time()
    drift_wave = np.zeros((dat_len, ch_num), dtype=np.int16)
    if end_idx == 0:
        rec2 = wfdb.rdrecord(rec_name, pb_dir = DB_DIR, sampfrom=N, sampto=N+dat_len)
    else:
        rec2 = wfdb.rdrecord(rec_name, pb_dir = DB_DIR, sampfrom=begin_idx, sampto=end_idx)

    for i in range(ch_num):
         drift_wave = np.convolve(rec.p_signal[:,i], b, 'valid')
         rec2.p_signal[:,i] = rec2.p_signal[:,i] - drift_wave
    elapsed_time = time.time() - start_time
    print ("elapsed_time:{0}".format(elapsed_time) + "[sec]")
    
    # Draw waveforms
    wfdb.plot_wfdb(record=rec2, title='Record '+rec_name + ' removed baseline drift') 
    
    # Save waveforms
    if not os.path.exists(SAVE_FOLDER_PATH):
        os.mkdir(SAVE_FOLDER_PATH)
    if end_idx == 0:
        file_name = 'wave_data_' + rec_name
    else:
        file_name = 'wave_data_' + rec_name + "_" + str(begin_idx) + "_" + str(begin_idx+dat_len)
    wave_dat_path = os.path.join(SAVE_FOLDER_PATH, file_name + '.pkl')
    wave_csv_path = os.path.join(SAVE_FOLDER_PATH, file_name + '.csv')
    wave_data = pd.DataFrame(rec2.p_signal)
    with open(wave_dat_path,'wb') as f:
        pickle.dump(wave_data, f)
    wave_data.to_csv(wave_csv_path, index=False)
    
    return rec, ann, rec2

if __name__ == "__main__" :
    load_data_ecg('234', begin_idx=148000, end_idx=152000)
    load_data_ecg('234', begin_idx=152000, end_idx=156000)
    