#!/usr/bin/python3
# Performs heartrate variation timedomain analysis
#
# It calculates the normalised RMSSD during sitting
# and maths.
#
# This comparison is then run with
# - ground truth (hand corrected R time stamps)
# - Wavelet detector
# - Pan Tompkins detector
#
# Via the commandline argument one can choose
# Einthoven II or the ECG from the Chest strap
#

import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as stats
from hrv import HRV
from ecgdetectors import Detectors

path_gu_ecg_database = '../dataset_716'

import sys
sys.path.insert(0, path_gu_ecg_database + r'/example_code')
from ecg_gla_database import Ecg


data_path = path_gu_ecg_database + r'/experiment_data'

maths_rr_sd = []
maths_error_rr_sd = []
maths_true_sd = []

sitting_rr_sd = []
sitting_error_rr_sd = []
sitting_true_sd = []

total_subjects = 25
subject = []

if len(sys.argv) < 2:
    print("Specify 'e' for Einthoven or 'v' for chest strap ECG.")
    exit(1)

for i in range(total_subjects):
#for i in range(2):
    print(i)
    sitting_class = Ecg(data_path, i, 'sitting')
    sitting_class.filter_data()
    maths_class = Ecg(data_path, i, 'maths')
    maths_class.filter_data()

    detectors = Detectors(sitting_class.fs)

    if sitting_class.anno_cs_exists and maths_class.anno_cs_exists and (i != 11):
        subject.append(i)

        hrv_class = HRV(sitting_class.fs)

        if "e" in sys.argv[1]:
            ecg_channel_sitting = sitting_class.einthoven_II
            ecg_channel_maths = maths_class.einthoven_II
        elif "v" in sys.argv[1]:
            ecg_channel_sitting = sitting_class.cs_V2_V1
            ecg_channel_maths = maths_class.cs_V2_V1
        else:
            print("Bad argument. Specify 'e' for Einthoven or 'v' for the Chest strap.")
            exit(1)

        r_peaks = detectors.swt_detector(ecg_channel_sitting)
        sitting_rr_sd.append(hrv_class.RMSSD(r_peaks,True))
        r_peaks = detectors.swt_detector(ecg_channel_maths)
        maths_rr_sd.append(hrv_class.RMSSD(r_peaks,True))

        sitting_error_rr = detectors.pan_tompkins_detector(ecg_channel_sitting)
        sitting_error_rr_sd.append(hrv_class.RMSSD(sitting_error_rr,True))

        maths_error_rr = detectors.pan_tompkins_detector(ecg_channel_maths)
        maths_error_rr_sd.append(hrv_class.RMSSD(maths_error_rr,True))

        maths_true_rr = maths_class.anno_cs
        maths_true_sd.append(hrv_class.RMSSD(maths_true_rr,True))
        
        sitting_true_rr = sitting_class.anno_cs
        sitting_true_sd.append(hrv_class.RMSSD(sitting_true_rr,True))


subject = np.array(subject)
width = 0.4

fig, ax = plt.subplots()
rects1 = ax.bar(subject+(0*width), sitting_true_sd, width)
rects2 = ax.bar(subject+(1*width), maths_true_sd, width)

ax.set_ylabel('SDNN (s)')
ax.set_xlabel('Subject')
ax.set_title('HRV for sitting and maths test')
ax.set_xticks(subject + width)
ax.set_xticklabels(subject)
ax.legend((rects1[0], rects2[0]), ('sitting', 'maths' ))

plt.figure()

# now let's do stats with no error

avg_sitting_rr_sd = np.average(sitting_rr_sd)
sd_sitting_rr_sd = np.std(sitting_rr_sd)

avg_maths_rr_sd = np.average(maths_rr_sd)
sd_maths_rr_sd = np.std(maths_rr_sd)

plt.bar(['sitting','maths'],[avg_sitting_rr_sd,avg_maths_rr_sd],yerr=[sd_sitting_rr_sd,sd_maths_rr_sd],align='center', alpha=0.5, ecolor='black', capsize=10)
plt.ylim([0,0.15])
plt.title("WAVELET: Sitting vs Maths")
plt.ylabel('nRMSSD')

# and stats with error

avg_sitting_error_rr_sd = np.average(sitting_error_rr_sd)
sd_sitting_error_rr_sd = np.std(sitting_error_rr_sd)

avg_maths_error_rr_sd = np.average(maths_error_rr_sd)
sd_maths_error_rr_sd = np.std(maths_error_rr_sd)

avg_sitting_true_sd = np.average(sitting_true_sd)
sd_sitting_true_sd = np.std(sitting_true_sd)

avg_maths_true_sd = np.average(maths_true_sd)
sd_maths_true_sd = np.std(maths_true_sd)

plt.figure()

plt.bar(['sitting','maths'],[avg_sitting_error_rr_sd,avg_maths_error_rr_sd],yerr=[sd_sitting_error_rr_sd,sd_maths_error_rr_sd],align='center', alpha=0.5, ecolor='black', capsize=10)
plt.ylim([0,0.15])
plt.title("Pan Tompkins DETECTOR: Sitting vs Maths")
plt.ylabel('nRMSSD')

plt.figure()

plt.bar(['sitting','maths'],[avg_sitting_true_sd,avg_maths_true_sd],yerr=[sd_sitting_true_sd,sd_maths_true_sd],align='center', alpha=0.5, ecolor='black', capsize=10)
plt.ylim([0,0.15])
plt.title("GROUND TRUTH: Sitting vs Maths")
plt.ylabel('nRMSSD')

t,p = stats.ttest_ind(sitting_true_sd,maths_true_sd,equal_var=False)
print("GROUND TRUTH (sitting vs maths): p=",p)

t,p = stats.ttest_ind(sitting_rr_sd,maths_rr_sd,equal_var=False)
print("WAVELET (sitting vs maths): p=",p)

t,p = stats.ttest_ind(sitting_error_rr_sd,maths_error_rr_sd,equal_var=False)
print("Pan Tompkins DETECTOR: (sitting vs maths): p=",p)

t,p = stats.ttest_ind(sitting_true_sd,sitting_rr_sd,equal_var=False)
print("Sitting: Wavelet vs ground truth, p=",p)

t,p = stats.ttest_ind(sitting_true_sd,sitting_error_rr_sd,equal_var=False)
print("Sitting: PT vs ground truth, p=",p)

t,p = stats.ttest_ind(maths_true_sd,maths_rr_sd,equal_var=False)
print("Maths: Wavelet vs ground truth, p=",p)

t,p = stats.ttest_ind(maths_true_sd,maths_error_rr_sd,equal_var=False)
print("Maths: PT vs ground truth, p=",p)

plt.show()