import wfdb

import matplotlib.pyplot as plt

import numpy as np

from hrv.filters import quotient, moving_median

from scipy import interpolate

from tqdm import tqdm

import pickle

import os

FS = 100.0

# From https://github.com/rhenanbartels/hrv/blob/develop/hrv/classical.py

def create_time_info(rri):

    rri_time = np.cumsum(rri) / 1000.0  # make it seconds

    return rri_time - rri_time[0]   # force it to start at zero

def create_interp_time(rri, fs):

    time_rri = create_time_info(rri)

    return np.arange(0, time_rri[-1], 1 / float(fs))

def interp_cubic_spline(rri, fs):

    time_rri = create_time_info(rri)

    time_rri_interp = create_interp_time(rri, fs)

    tck = interpolate.splrep(time_rri, rri, s=0)

    rri_interp = interpolate.splev(time_rri_interp, tck, der=0)

    return time_rri_interp, rri_interp

def interp_cubic_spline_qrs(qrs_index, qrs_amp, fs):

    time_qrs = qrs_index / float(FS)

    time_qrs = time_qrs - time_qrs[0]

    time_qrs_interp = np.arange(0, time_qrs[-1], 1/float(fs))

    tck = interpolate.splrep(time_qrs, qrs_amp, s=0)

    qrs_interp = interpolate.splev(time_qrs_interp, tck, der=0)

    return time_qrs_interp, qrs_interp

data_path = './data/'

train_data_name = ['a02', 'a03', 'a04', 'a05',

             'a06', 'a07', 'a08', 'a09', 'a10',

             'a11', 'a12', 'a13', 'a14', 'a15',

             'a16', 'a17', 'a18', 'a19',

             'b02', 'b03', 'b04',

             'c02', 'c03', 'c04', 'c05',

             'c06', 'c07', 'c08', 'c09',

             ]

val_data_name = ['a01', 'b01', 'c01']

test_data_name = ['a20','b05','c10']

age = [51, 38, 54, 52, 58,

       63, 44, 51, 52, 58,

       58, 52, 51, 51, 60,

       44, 40, 52, 55, 58,

       44, 53, 53, 42, 52,

       31, 37, 39, 41, 28,

       28, 30, 42, 37, 27]

sex = [1, 1, 1, 1, 1,

       1, 1, 1, 1, 1,

       1, 1, 1, 1, 1,

       1, 1, 1, 1, 1,

       0, 1, 1, 1, 1,

       1, 1, 1, 0, 0,

       0, 0, 1, 1, 1]

def get_qrs_amp(ecg, qrs):

    interval = int(FS * 0.250)

    qrs_amp = []

    for index in range(len(qrs)):

        curr_qrs = qrs[index]

        amp = np.max(ecg[curr_qrs-interval:curr_qrs+interval])

        qrs_amp.append(amp)

    return qrs_amp

MARGIN = 10

FS_INTP = 4

MAX_HR = 300.0

MIN_HR = 20.0

MIN_RRI = 1.0 / (MAX_HR / 60.0) * 1000

MAX_RRI = 1.0 / (MIN_HR / 60.0) * 1000

train_input_array = []

train_label_array = []

for data_index in range(len(train_data_name)):

    print (train_data_name[data_index])

    win_num = len(wfdb.rdann(os.path.join(data_path,train_data_name[data_index]), 'apn').symbol)

    signals, fields = wfdb.rdsamp(os.path.join(data_path,train_data_name[data_index]))

    for index in tqdm(range(1, win_num)):

        samp_from = index * 60 * FS # 60 seconds

        samp_to = samp_from + 60 * FS  # 60 seconds

        qrs_ann = wfdb.rdann(data_path + train_data_name[data_index], 'qrs', sampfrom=samp_from - (MARGIN*100), sampto=samp_to + (MARGIN*100)).sample

        apn_ann = wfdb.rdann(data_path + train_data_name[data_index], 'apn', sampfrom=samp_from, sampto=samp_to-1).symbol

        qrs_amp = get_qrs_amp(signals, qrs_ann)

        rri = np.diff(qrs_ann)

        rri_ms = rri.astype('float') / FS * 1000.0

        try:

            rri_filt = moving_median(rri_ms)

            if len(rri_filt) > 5 and (np.min(rri_filt) >= MIN_RRI and np.max(rri_filt) <= MAX_RRI):

                time_intp, rri_intp = interp_cubic_spline(rri_filt, FS_INTP)

                qrs_time_intp, qrs_intp = interp_cubic_spline_qrs(qrs_ann, qrs_amp, FS_INTP)

                rri_intp = rri_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                qrs_intp = qrs_intp[(qrs_time_intp >= MARGIN) & (qrs_time_intp < (60 + MARGIN))]

                #time_intp = time_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                if len(rri_intp) != (FS_INTP * 60):

                    skip = 1

                else:

                    skip = 0

                if skip == 0:

                    rri_intp = rri_intp - np.mean(rri_intp)

                    qrs_intp = qrs_intp - np.mean(qrs_intp)

                    if apn_ann[0] == 'N': # Normal

                        label = 0.0

                    elif apn_ann[0] == 'A': # Apnea

                        label = 1.0

                    else:

                        label = 2.0

                    train_input_array.append([rri_intp, qrs_intp, age[data_index], sex[data_index]])

                    train_label_array.append(label)

        except:

            hrv_module_error = 1

with open('train_input.pickle','wb') as f: 

    pickle.dump(train_input_array, f)

with open('train_label.pickle','wb') as f: 

    pickle.dump(train_label_array, f)

val_input_array = []

val_label_array = []

for data_index in range(len(val_data_name)):

    print (val_data_name[data_index])

    win_num = len(wfdb.rdann(os.path.join(data_path,val_data_name[data_index]), 'apn').symbol)

    signals, fields = wfdb.rdsamp(os.path.join(data_path,val_data_name[data_index]))

    for index in tqdm(range(1, win_num)):

        samp_from = index * 60 * FS # 60 seconds

        samp_to = samp_from + 60 * FS  # 60 seconds

        qrs_ann = wfdb.rdann(data_path + val_data_name[data_index], 'qrs', sampfrom=samp_from - (MARGIN*100), sampto=samp_to + (MARGIN*100)).sample

        apn_ann = wfdb.rdann(data_path + val_data_name[data_index], 'apn', sampfrom=samp_from, sampto=samp_to-1).symbol

        qrs_amp = get_qrs_amp(signals, qrs_ann)

        rri = np.diff(qrs_ann)

        rri_ms = rri.astype('float') / FS * 1000.0

        try:

            rri_filt = moving_median(rri_ms)

            if len(rri_filt) > 5 and (np.min(rri_filt) >= MIN_RRI and np.max(rri_filt) <= MAX_RRI):

                time_intp, rri_intp = interp_cubic_spline(rri_filt, FS_INTP)

                qrs_time_intp, qrs_intp = interp_cubic_spline_qrs(qrs_ann, qrs_amp, FS_INTP)

                rri_intp = rri_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                qrs_intp = qrs_intp[(qrs_time_intp >= MARGIN) & (qrs_time_intp < (60 + MARGIN))]

                #time_intp = time_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                if len(rri_intp) != (FS_INTP * 60):

                    skip = 1

                else:

                    skip = 0

                if skip == 0:

                    rri_intp = rri_intp - np.mean(rri_intp)

                    qrs_intp = qrs_intp - np.mean(qrs_intp)

                    if apn_ann[0] == 'N': # Normal

                        label = 0.0

                    elif apn_ann[0] == 'A': # Apnea

                        label = 1.0

                    else:

                        label = 2.0

                    val_input_array.append([rri_intp, qrs_intp, age[data_index], sex[data_index]])

                    val_label_array.append(label)

        except:

            hrv_module_error = 1

with open('val_input.pickle','wb') as f: 

    pickle.dump(val_input_array, f)

with open('val_label.pickle','wb') as f: 

    pickle.dump(val_label_array, f)

test_input_array = []

test_label_array = []

for data_index in range(len(test_data_name)):

    print (test_data_name[data_index])

    win_num = len(wfdb.rdann(os.path.join(data_path,test_data_name[data_index]), 'apn').symbol)

    signals, fields = wfdb.rdsamp(os.path.join(data_path,test_data_name[data_index]))

    for index in tqdm(range(1, win_num)):

        samp_from = index * 60 * FS # 60 seconds

        samp_to = samp_from + 60 * FS  # 60 seconds

        qrs_ann = wfdb.rdann(data_path + test_data_name[data_index], 'qrs', sampfrom=samp_from - (MARGIN*100), sampto=samp_to + (MARGIN*100)).sample

        apn_ann = wfdb.rdann(data_path + test_data_name[data_index], 'apn', sampfrom=samp_from, sampto=samp_to-1).symbol

        qrs_amp = get_qrs_amp(signals, qrs_ann)

        rri = np.diff(qrs_ann)

        rri_ms = rri.astype('float') / FS * 1000.0

        try:

            rri_filt = moving_median(rri_ms)

            if len(rri_filt) > 5 and (np.min(rri_filt) >= MIN_RRI and np.max(rri_filt) <= MAX_RRI):

                time_intp, rri_intp = interp_cubic_spline(rri_filt, FS_INTP)

                qrs_time_intp, qrs_intp = interp_cubic_spline_qrs(qrs_ann, qrs_amp, FS_INTP)

                rri_intp = rri_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                qrs_intp = qrs_intp[(qrs_time_intp >= MARGIN) & (qrs_time_intp < (60 + MARGIN))]

                #time_intp = time_intp[(time_intp >= MARGIN) & (time_intp < (60+MARGIN))]

                if len(rri_intp) != (FS_INTP * 60):

                    skip = 1

                else:

                    skip = 0

                if skip == 0:

                    rri_intp = rri_intp - np.mean(rri_intp)

                    qrs_intp = qrs_intp - np.mean(qrs_intp)

                    if apn_ann[0] == 'N': # Normal

                        label = 0.0

                    elif apn_ann[0] == 'A': # Apnea

                        label = 1.0

                    else:

                        label = 2.0

                    test_input_array.append([rri_intp, qrs_intp, age[data_index], sex[data_index]])

                    test_label_array.append(label)

        except:

            hrv_module_error = 1

with open('test_input.pickle','wb') as f: 

    pickle.dump(test_input_array, f)

with open('test_label.pickle','wb') as f: 

    pickle.dump(test_label_array, f)