import numpy as np
from tqdm import tqdm
from helper_code import *
from scipy import signal
import pandas as pd
from skmultilearn.model_selection import iterative_train_test_split
import random
import warnings
import sys


def get_nsamp(header):
    return int(header.split('\n')[0].split(' ')[3])

# Adapted from original scoring function code
# For each set of equivalent classes, replace each class with the representative class for the set.


'''
the 'output' will be using the zero-padding if the leads are less than 12.(8 leads)
the 'output_leads' will be [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]

the output will not be changed if the lead is 12
the 'output_leads' will be [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
'''
def mixup(data, mix_data):
    mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
    if data.shape[1] > mix_data.shape[1]:
        data = mix_lambda * data[:, :mix_data.shape[1]] + (1 - mix_lambda) * mix_data
    else:
        data = mix_lambda * data + (1 - mix_lambda) * mix_data[:, :data.shape[1]]
    return data, mix_lambda

def cutmix(data, mix_data):
    cutmix_lambda = np.random.beta(10, 10)
    
    if data.shape[1] < mix_data.shape[1]:
        seq_length = data.shape[1]
        win_len = int(np.ceil(seq_length * cutmix_lambda))
        start = np.random.randint(0, seq_length - win_len + 1)
        end = start + win_len
        data[:,start:end] = mix_data[:,start:end]
    else:
        seq_length = mix_data.shape[1]
        win_len = int(np.ceil(seq_length * cutmix_lambda))
        start = np.random.randint(0, seq_length - win_len + 1)
        end = start + win_len
        data[:,start:end] = mix_data[:,start:end]
    return data, cutmix_lambda

def cutmix_fix_length(data, mix_data):
    # cutmix_lambda = np.random.beta(0.2, 0.2)
    cutmix_lambda = 0.2

    if data.shape[1] < mix_data.shape[1]:
        seq_length = data.shape[1]
        win_len = int(np.ceil(seq_length * cutmix_lambda))
        start = np.random.randint(0, seq_length - win_len + 1)
        end = start + win_len
        data[:,start:end] = mix_data[:,start:end]
    else:
        seq_length = mix_data.shape[1]
        win_len = int(np.ceil(seq_length * cutmix_lambda))
        start = np.random.randint(0, seq_length - win_len + 1)
        end = start + win_len
        data[:,start:end] = mix_data[:,start:end]
    return data, cutmix_lambda

def same_shape_mixup(data, mix_data):
    if data.shape[1] == mix_data.shape[1]:
        mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
        data = mix_lambda * data + (1 - mix_lambda) * mix_data
    return data

def expand_leads(recording, input_leads):
    output = np.zeros((12, recording.shape[1]))
    # recording.shape[1]: 5000
    twelve_leads = ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
    twelve_leads = [k.lower() for k in twelve_leads]
    # ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']

    input_leads = [k.lower() for k in input_leads]
    # Here we can assume: 
    # input_leads:I, II, V1, V2, V3, V4, V5, V6, 
    # so the new input_leads: ['i', 'ii', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']

    output_leads = np.zeros((12,))
    # [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]

    # idx: [0, 1, 6, 7, 8, 9, 10, 11]
    for i,k in enumerate(input_leads):
        idx = twelve_leads.index(k)
        output[idx,:] = recording[i,:]
        output_leads[idx] = 1

    return output, output_leads

'''
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 0. 0. 0. 0. 0. 0.]

'''

class lead_exctractor:
    """
    used to select specific leads or random choice of configurations

    Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
    Eight leads: I, II, V1, V2, V3, V4, V5, V6
    Six leads: I, II, III, aVR, aVL, aVF
    Four leads: I, II, III, V2
    Three leads: I, II, V2
    Two leads: I, II

    """
    L2  = np.array([1,1,0,0,0,0,0,0,0,0,0,0])
    L3  = np.array([1,1,0,0,0,0,0,1,0,0,0,0])
    L4  = np.array([1,1,1,0,0,0,0,1,0,0,0,0])
    L6  = np.array([1,1,1,1,1,1,0,0,0,0,0,0])
    L8  = np.array([1,1,0,0,0,0,1,1,1,1,1,1])
    L12 = np.array([1,1,1,1,1,1,1,1,1,1,1,1])

    @staticmethod
    def get (x, num_leads, lead_indicator):
        if num_leads==None:
            # random choice output
            num_leads = random.choice([12,8,6,4,3,2])

        if num_leads==12:
            # Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
            return x, lead_indicator * lead_exctractor.L12

        if num_leads==8:
            # Six leads:  I, II, V1, V2, V3, V4, V5, V6
            x = x * lead_exctractor.L8.reshape(12,1)
            return x,lead_indicator * lead_exctractor.L8

        if num_leads==6:
            # Six leads: I, II, III, aVL, aVR, aVF
            x = x * lead_exctractor.L6.reshape(12,1)
            return x,lead_indicator * lead_exctractor.L6

        if num_leads==4:
            # Six leads: I, II, III, V2
            x = x * lead_exctractor.L4.reshape(12,1)
            return x,lead_indicator * lead_exctractor.L4

        if num_leads==3:
            # Three leads: I, II, V2
            x = x * lead_exctractor.L3.reshape(12,1)
            return x,lead_indicator * lead_exctractor.L3

        if num_leads==2:
            # Two leads: II, V5
            x = x * lead_exctractor.L2.reshape(12,1)
            return x,lead_indicator * lead_exctractor.L2
        raise Exception("invalid-leads-number")

class dataset:
    # #ECG 2021
    # classes = ['164889003','164890007','6374002','426627000','733534002',
    #            '713427006','270492004','713426002','39732003','445118002',
    #            '164947007','251146004','111975006','698252002','426783006',
    #            '284470004','10370003','365413008','427172004','164917005',
    #            '47665007','427393009','426177001','427084000','164934002',
    #            '59931005']
    
    # ECG 2020

    

    normal_class = '426783006'
    
    def __init__(self, header_files, Mixup = 0, amount = 0, cutMix = 0, Mixup_no_label_interpolate =0, progressive_switch = False):
        self.files = []
        self.sample = True
        self.num_leads = None

        for h in tqdm(header_files):
            tmp = dict()
            tmp['header'] = h
            tmp['record'] = h.replace('.hea','.mat')
            hdr = load_header(h)
            tmp['nsamp'] = get_nsamp(hdr)
            tmp['leads'] = get_leads(hdr)
            tmp['age'] = get_age(hdr)
            tmp['sex'] = get_sex(hdr)
            tmp['dx'] = get_labels(hdr)
            tmp['fs'] = get_frequency(hdr)
            # tmp['target'] = np.zeros((26,))


            self.files.append(tmp)

        # print("This is the target:", tmp['target'])
        # set filter parameters
        self.b, self.a = signal.butter(3, [1 / 250, 47 / 250], 'bandpass')

        self.files = pd.DataFrame(self.files)

        self.Mixup = Mixup
        self.cutMix = cutMix
        self.Mixup_no_label_interplate = Mixup_no_label_interpolate
        self.amount = amount
        self.progressive = progressive_switch

        self.current_epoch = 0  # Initialize the current epoch

        address_2020 = "./csv-file_2020_challenge/training_validation_testing/group1/train_validation_testing_fold5.csv"
        self.data_df_2020 = pd.read_csv(address_2020, index_col=0)
        self.data_df_2020.set_index('Patient', inplace=True)
        self.classes_2020 = ['270492004', '164889003', '164890007', '426627000', '713427006', 
                  '713426002', '445118002', '39732003', '164909002', '251146004', 
                  '698252002', '10370003', '284470004', '427172004', '164947007', 
                  '111975006', '164917005', '47665007', '59118001', '427393009', 
                  '426177001', '426783006', '427084000', '63593006', '164934002', 
                  '59931005', '17338001']

    
    def set_epoch(self, epoch):
        self.current_epoch = epoch
    
    def summary(self, output):
        if output=='pandas':
            # print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
            # print("This is the target:",self.files['target'],type(self.files['target']))
            return pd.Series(np.stack(self.files['target'].to_list(),axis=0).sum(axis=0),index=dataset.classes)

        if output=='numpy':
            return np.stack(self.files['target'].to_list(),axis=0).sum(axis=0)

    def __len__(self):
        return len(self.files)
    
    '''
    fs: 500.0   sampling rate
    target: [0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
    leads: ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
    self.files.iloc[item]['record']: ../../python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00518.mat
    data: shape->(12 X 5000)
    '''

    def __getitem__(self, item):

        fs = self.files.iloc[item]['fs']
        # target = self.files.iloc[item]['target']
        header = self.files.iloc[item]['header']
        name = header.split('/')[-1]
        row = self.data_df_2020.loc[name[0:-4]]
        target = row[self.classes_2020].values.astype(np.int_)

        leads = self.files.iloc[item]['leads']
        data = load_recording(self.files.iloc[item]['record'])
        
        # print("This is the target:", target)
        # Set your threshold
        
        threshold = 0.5
        # print("This is the original target:", target)
        # print("This is the type of original target:", type(target))
        
        # if random.random() < self.Mixup :
        #     mix_sample_idx = random.randint(0, self.amount-1)
        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
        #     mix_target = self.files.iloc[mix_sample_idx]['target']
        #     data, alpha_value = mixup(data, mix_datum)
        #     target = alpha_value * target + (1-alpha_value) * mix_target
        #     # Binarize the labels based on the threshold
        #     target[target >= threshold] = 1
        #     target[target < threshold] = 0
            
        
        # if random.random() < self.cutMix:
        #     mix_sample_idx = random.randint(0, self.amount-1)
        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
        #     mix_target = self.files.iloc[mix_sample_idx]['target']
        #     data, alpha_value = cutmix_fix_length(data, mix_datum)
        #     target = alpha_value * target + (1-alpha_value) * mix_target
        #     target[target >= threshold] = 1
        #     target[target < threshold] = 0
        
        # if random.random()< self.Mixup_no_label_interplate:
        #     mix_sample_idx = random.randint(0, self.amount-1)
        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
        #     data= same_shape_mixup(data, mix_datum)

        # proressive_index = 0
        
        # # # below is the progessive data augmentation method
        # # if self.progressive: 
        # #     if self.current_epoch < 20:
        # #         if self.current_epoch % 4 == 0 or 3:
        # #             pass
        # #         elif self.current_epoch % 4 == 1:
        # #             proressive_index = 0.2
        # #         else:
        # #             proressive_index = 0.5

        # #     elif self.current_epoch < 40:
        # #         if self.current_epoch % 4 == 0 or 3:
        # #             pass
        # #         elif self.current_epoch % 4 == 1:
        # #             proressive_index = 0.7
        # #         else:
        # #             proressive_index = 0.8
        # #     else:
        # #         if self.current_epoch % 4 == 0 or 3:
        # #             pass
        # #         elif self.current_epoch % 4 == 1 or 2:
        # #             proressive_index = 1
                    
        # # below is the wave-mix data augmentation method
        # if self.progressive: 
        #     if self.current_epoch < 5:
        #         match self.current_epoch:
        #             case 0:
        #                 self.progressive=0
        #             case 1:
        #                 self.progressive=0.2
        #             case 2:
        #                 self.progressive=0.4
        #             case 3:
        #                 self.progressive=0.6
        #             case 4:
        #                 self.progressive=0.8
                    
        #     else :
        #         proressive_index = 0.8
            
        # if random.random() < proressive_index :
        #     mix_sample_idx = random.randint(0, self.amount-1)
        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
        #     mix_target = self.files.iloc[mix_sample_idx]['target']
        #     data, alpha_value = mixup(data, mix_datum)
        #     target = alpha_value * target + (1-alpha_value) * mix_target
        #     # Binarize the labels based on the threshold
        #     target[target >= threshold] = 1
        #     target[target < threshold] = 0


        # expand to 12 lead setup if original signal has less channels
        data, lead_indicator = expand_leads(data, input_leads=leads)
        data = np.nan_to_num(data)
 

        # resample to 500hz
        if fs == float(1000):
            data = signal.resample_poly(data, up=1, down=2, axis=-1)  # to 500Hz
            fs = 500
        elif fs == float(500):
            pass
        else:
            data = signal.resample(data, int(data.shape[1] * 500 / fs), axis=1)
            fs = 500

        # below is the Butterworth digital and analog filter design.
        data = signal.filtfilt(self.b, self.a, data)

        '''
        we filter out the sample which the length is 8192 via the below code.
        just like the random shift windows.  
        '''
        if self.sample:
            fs = int(fs)
            # random sample signal if len > 8192 samples
            if data.shape[-1] >= 8192:
                idx = data.shape[-1] - 8192-1
                idx = np.random.randint(idx)
                data = data[:, idx:idx + 8192]

        '''
        we can obtain the average and mean for each lead via below code, 90% conclusion.
        '''
        mu  = np.nanmean(data, axis=-1, keepdims=True)
        std = np.nanstd(data, axis=-1, keepdims=True)

        '''
        Z-Score Normalization
        '''
        #std = np.nanstd(data.flatten())
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            data = (data - mu) / std
        
        data = np.nan_to_num(data)

        # random choose number of leads to keep
        data, lead_indicator = lead_exctractor.get(data, self.num_leads, lead_indicator)
        
        # print("This is the lead_indicator:(l)", type(lead_indicator), lead_indicator)
        

        return data, target