##### ------------------- klasifikace paroxysmální a perzistentní fibrilace síní -----------------

### importy
import numpy as np
import random
import matplotlib.pyplot as plt
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from bayes_opt import BayesianOptimization
from scipy.signal import medfilt, find_peaks, peak_prominences, peak_widths
from scipy import interpolate
from sklearn.metrics import mean_squared_error
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPClassifier

from train_valid import classification_metrics
from config import Config

## ----------- Klasifikace paroxysmálních a perzistentní FiS - rozhodovací pravidla s prohledáváním stavového prostoru pro volbu parametrů -----------

def parox_afib_classification(heat_maps, heat_map_ind, ground_truth_parox, input=None, threshold_std=None, threshold_max=None):

    if input == "training":

        ind_std = np.arange(0.5, 3.551, 0.05)
        ind_max = np.arange(2, 8.01, 0.1)
        conf_mat_parox_all = list()
        f1_parox_all = list()
        std_max_all = list()
        clas_parox_all = list()
        index_parox_all = list()
        index_non_parox_all = list()

        for istd in ind_std:
            for imax in ind_max:

                clas_parox_temp = list()
                index_parox_temp = list()
                index_non_parox_temp = list()

                for i in heat_map_ind:

                    ## klasifikace parox vs. persi FiS
                    std_hm = np.max(np.std(heat_maps[i, :, :], axis=1))
                    max_hm = np.max(heat_maps[i, :, :])
                    threshold_std_hm = istd
                    threshold_max_hm = imax
                    if max_hm > threshold_max_hm:
                        clas_parox_temp.append(1)
                        index_parox_temp.append(i)
                    elif std_hm > threshold_std_hm:
                        clas_parox_temp.append(1)
                        index_parox_temp.append(i)
                    else:
                        clas_parox_temp.append(0)
                        index_non_parox_temp.append(i)

                ## výpočet matice záměn
                conf_mat_parox_temp, f1_parox_temp = classification_metrics(clas_parox_temp, ground_truth_parox)

                conf_mat_parox_all.append(conf_mat_parox_temp)
                f1_parox_all.append(f1_parox_temp)
                std_max_all.append((istd, imax))
                clas_parox_all.append(clas_parox_temp)
                index_parox_all.append(index_parox_temp)
                index_non_parox_all.append(index_non_parox_temp)

        f1_parox_all = np.array(f1_parox_all)
        ind_best = np.argmax(f1_parox_all)
        f1_parox = f1_parox_all[ind_best]
        threshold_std, threshold_max = std_max_all[ind_best]
        conf_mat_parox = conf_mat_parox_all[ind_best]

        print(f'Min Max Standard deviation: {threshold_std}, Min Maximum: {threshold_max}')
        print(f'Paroxysmal afib classification:\n{conf_mat_parox}, F1 = {f1_parox}')

        clas_parox = clas_parox_all[ind_best]
        index_parox = index_parox_all[ind_best]
        index_non_parox = index_non_parox_all[ind_best]

        return clas_parox, index_parox, index_non_parox, threshold_std, threshold_max

    elif input == "testing":

        clas_parox = list()
        index_parox = list()
        index_non_parox = list()

        for i in heat_map_ind:

            std_hm = np.max(np.std(heat_maps[i, :, :], axis=1))
            max_hm = np.max(heat_maps[i, :, :])
            threshold_std_hm = threshold_std
            threshold_max_hm = threshold_max
            if max_hm > threshold_max_hm:
                clas_parox.append(1)
                index_parox.append(i)
            elif std_hm > threshold_std_hm:
                clas_parox.append(1)
                index_parox.append(i)
            else:
                clas_parox.append(0)
                index_non_parox.append(i)

        ## výpočet matice záměn
        conf_mat_parox, f1_parox = classification_metrics(clas_parox, ground_truth_parox)
        print(f'Paroxysmal afib classification:\n{conf_mat_parox}, F1 = {f1_parox}')

        return clas_parox, index_parox, index_non_parox
