#!/usr/bin/python3

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

import scipy.stats as stats

def str_join(a, b, c):

    a = np.array(a)

    b = np.array(b)

    c = np.array(c)

    str_array = np.core.defchararray.add(np.core.defchararray.add(a, b), c)

    return str_array

def get_sensivities(results_df, detector_name, experiment=None):

    if experiment!=None:

        tp_col_names = str_join(detector_name+' ', [experiment], ' TP')

        fp_col_names = str_join(detector_name+' ', [experiment], ' FP')

        fn_col_names = str_join(detector_name+' ', [experiment], ' FN')

        tn_col_names = str_join(detector_name+' ', [experiment], ' TN')

        total_tp = (results_df.loc[:, tp_col_names].values)[:,0]

        total_fn = (results_df.loc[:, fn_col_names].values)[:,0]

    else:

        tp_col_names = detector_name+' '+'TP'

        fp_col_names = detector_name+' '+'FP'

        fn_col_names = detector_name+' '+'FN'

        tn_col_names = detector_name+' '+'TN'      

        total_tp = results_df.loc[:, tp_col_names].values

        total_fn = results_df.loc[:, fn_col_names].values

    se = []

    for tp,fn in zip(total_tp,total_fn):

            if (tp + fn) > 0:

                s = tp/(tp+fn)*100.0

                if s > 0:

                        se.append(s)

    return np.array(se)

def get_result(results_df, det_names, experiment=None):

    m = []

    s = []

    for det in det_names:

        m.append(np.mean(get_sensivities(results_df, det, experiment)))

        s.append(np.std(get_sensivities(results_df, det, experiment)))

    return np.array(m),np.array(s)

def compare_det_test(results_df, detector_name1, detector_name2, experiment=None):

    se1 = get_sensivities(results_df, detector_name1, experiment)

    if len(se1) < 2:

            return 0

    se2 = get_sensivities(results_df, detector_name2, experiment)

    if len(se2) < 2:

            return 0

    l = min(len(se1),len(se2))

    #print("1:",se1[:l])

    #print("2:",se2[:l])

    try:

        t,p = stats.wilcoxon(se1[:l],se2[:l])

        return p

    except:

        return None

def single_plot(data, std, y_label, title = None):

    fig, ax = plt.subplots()

    plot_names = ['Elgendi et al', 'Matched Filter', 'Kalidas and Tamil', 'Engzee Mod', 'Christov', 'Hamilton', 'Pan and Tompkins']

    x_pos = np.arange(len(plot_names))

    fig.set_size_inches(10, 7)

    rects1 = ax.bar(x_pos, data, yerr=std, width = 0.65, align='center', alpha=0.5, ecolor='black', capsize=10)

    ax.set_ylim([0,100])

    ax.set_ylabel(y_label)

    ax.set_xlabel('Detector')

    ax.set_xticks(x_pos)

    ax.set_xticklabels(plot_names)

    if title!=None:

        ax.set_title(title)

    plt.tight_layout()

    return rects1

def double_plot(data1, std1, data2, std2, y_label, legend1, legend2, title=None):

    fig, ax = plt.subplots()

    plot_names = ['Elgendi et al', 'Matched Filter', 'Kalidas and Tamil', 'Engzee Mod', 'Christov', 'Hamilton', 'Pan and Tompkins']

    x_pos = np.arange(len(plot_names))

    fig.set_size_inches(10, 7)

    width = 0.4

    rects1 = ax.bar(x_pos, data1, width, yerr=std1, alpha=0.5, ecolor='black', capsize=10)

    rects2 = ax.bar(x_pos+width, data2, width, yerr=std2, alpha=0.5, ecolor='black', capsize=10)

    ax.set_ylim([0,100])

    ax.set_ylabel(y_label)

    ax.set_xlabel('Detector')

    ax.set_xticks(x_pos + width / 2)

    ax.set_xticklabels(plot_names)

    ax.legend((rects1[0], rects2[0]), (legend1, legend2))

    if title!=None:

        ax.set_title(title)

    plt.tight_layout()

    return rects1, rects2

def print_stat(p):

    if p == None:

        print('--- & ',end='')

        return

    s = ""

    if p < 0.05:

        s = "*"

    print('{:03.2f}{} & '.format(p,s),end='')

# GUDB

gudb_cs_results = pd.read_csv('results_GUDB_chest_strap.csv', dtype=int, index_col=0)

gudb_cable_results = pd.read_csv('results_GUDB_loose_cables.csv', dtype=int, index_col=0)

# MITDB

mitdb_results = pd.read_csv('results_MITDB.csv', dtype=int, index_col=0)

det_names = ['two_average', 'matched_filter', 'swt', 'engzee', 'christov', 'hamilton', 'pan_tompkins']

plot_names = ['Elgendi et al', 'Matched Filter', 'Kalidas and Tamil', 'Engzee Mod', 'Christov', 'Hamilton', 'Pan and Tompkins']

experiment_names = ['sitting','maths','walking','hand_bike','jogging']

output_names = ['TP', 'FP', 'FN', 'TN']

print("MIT")

for det1 in det_names:

    for det2 in det_names:

        p = compare_det_test(mitdb_results, det1, det2)

        print_stat(p)

    print("\\\\")

print("CHEST STRAP SITTING")

for det1 in det_names:

    for det2 in det_names:

        p = compare_det_test(gudb_cs_results, det1, det2, 'sitting')

        print_stat(p)

    print("\\\\")

print("CHEST STRAP JOGGING")

for det1 in det_names:

    for det2 in det_names:

        p = compare_det_test(gudb_cs_results, det1, det2, 'jogging')

        print_stat(p)

    print("\\\\")

print("LOOSE CABLE SITTING")

for det1 in det_names:

    for det2 in det_names:

        p = compare_det_test(gudb_cable_results, det1, det2, 'sitting')

        print_stat(p)

    print("\\\\")

print("LOOSE CABLE JOGGING")

for det1 in det_names:

    for det2 in det_names:

        p = compare_det_test(gudb_cable_results, det1, det2, 'jogging')

        print_stat(p)

    print("\\\\")

# calculating results

mitdb_avg,mitdb_std = get_result(mitdb_results, det_names)

print("mitdb:",mitdb_avg)

gudb_cs_sitting_avg,gudb_cs_sitting_std = get_result(gudb_cs_results, det_names, 'sitting')

print("chest strap sitting:",gudb_cs_sitting_avg)

gudb_cable_sitting_avg,gudb_cable_sitting_std = get_result(gudb_cable_results, det_names, 'sitting')

print("lose cables sitting:",gudb_cable_sitting_avg)

gudb_cs_jogging_avg,gudb_cs_jogging_std = get_result(gudb_cs_results, det_names, 'jogging')

gudb_cable_jogging_avg,gudb_cable_jogging_std = get_result(gudb_cable_results, det_names, 'jogging')

# plotting

single_plot(mitdb_avg, mitdb_std, 'Sensitivity (%)', 'MITDB')

double_plot(gudb_cs_sitting_avg, gudb_cs_sitting_std,

            gudb_cable_sitting_avg, gudb_cable_sitting_std,

            'Sensitivity (%)', 'Chest Strap', 'Loose Cables', 'GUDB: cable, sitting')

double_plot(gudb_cs_jogging_avg, gudb_cs_jogging_std,

            gudb_cable_jogging_avg, gudb_cable_jogging_std,

            'Sensitivity (%)', 'Chest Strap', 'Loose Cables', 'GUDB: cable, jogging')

plt.show()