import numpy as np
import random
from helper_code import *
from tqdm import tqdm
import pandas as pd
from scipy import signal
import warnings

def standard_to_cabrera(ecg_data):
    """
    Converts 12-lead ECG from Standard order to Cabrera order.
    
    Input shape: (12, N)
    Input order:  I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
    
    Output order: aVL, I, -aVR, II, aVF, III, V1, V2, V3, V4, V5, V6
    """
    
    # 1. Define the index mapping from Standard to Cabrera
    # Standard Indices:
    # I=0, II=1, III=2, aVR=3, aVL=4, aVF=5, V1=6, V2=7 ... V6=11
    
    # Cabrera Target:             aVL, I, aVR, II, aVF, III, V1...V6
    cabrera_indices = np.array([  4,   0,  3,  1,   5,   2,  6, 7, 8, 9, 10, 11])
    
    # 2. Reorder the rows using advanced integer indexing
    # This creates a copy of the data in the new order
    cabrera_data = ecg_data[cabrera_indices, :]
    
    # 3. Invert aVR to become -aVR
    # In the NEW Cabrera array, aVR is now at index 2 (the 3rd row)
    cabrera_data[2, :] = -1 * cabrera_data[2, :]
    
    return cabrera_data


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

def replace_equivalent_classes(classes, equivalent_classes):
    for j, x in enumerate(classes):
        for multiple_classes in equivalent_classes:
            if x in multiple_classes:
                classes[j] = multiple_classes[0]  # Use the first class as the representative class.
    return classes

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")

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

class dataset:
    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']
    normal_class = '426783006'
    equivalent_classes = [['713427006', '59118001'],
                          ['284470004', '63593006'],
                          ['427172004', '17338001'],
                          ['733534002', '164909002']]
    def __init__(self, header_files):
        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,))
            tmp['dx'] = replace_equivalent_classes(tmp['dx'], dataset.equivalent_classes)

            for dx in tmp['dx']:
                # in SNOMED code is in scored classes
                if dx in dataset.classes:
                    idx = dataset.classes.index(dx)
                    tmp['target'][idx] = 1

            self.files.append(tmp)

        # print("This is the target:", tmp['target'])
        # set filter parameters

        # Filtering: Data are filtered using a zero-phase method with 3rd order Butterworth bandpass filter 
        # with frequency band from 1 Hz to 47 Hz.
        self.b, self.a = signal.butter(3, [1 / 250, 47 / 250], 'bandpass')

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


        self.current_epoch = 0  # Initialize the current epoch

    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 Pre_processing(self, fs, leads, data):

        # 1. Expand to a 12-lead setup if the original signal has fewer channels 
        data, lead_indicator = expand_leads(data, input_leads=leads)
        data = np.nan_to_num(data)
 

        # 2. Resample from 1000 Hz and other rates to 100 Hz
        if fs == float(1000):
            data = signal.resample_poly(data, up=1, down=10, axis=-1)  # to 500Hz
            fs = 100
        elif fs == float(500):
            data = signal.resample_poly(data, up=1, down=5, axis=-1)  # to 500Hz
            fs = 100
        else:
            data = signal.resample(data, int(data.shape[1] * 100 / fs), axis=1)
            fs = 100


        # 4. We compute the average and mean for each lead using the code below; this is Z-score normalization 
        mu  = np.nanmean(data, axis=-1, keepdims=True)
        std = np.nanstd(data, axis=-1, keepdims=True)
        
        #std = np.nanstd(data.flatten())
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            data = (data - mu) / std
        
        data = np.nan_to_num(data)

        # 5. Selection of leads for 12-lead ECGs; the default is 12 leads
        data, lead_indicator = lead_exctractor.get(data, self.num_leads, lead_indicator)  

        each_file_length = 1000

        # 6. We filter out samples with a length of 1000 using the code below, similar to random shift windows, and zero-padding as well 
        if self.sample:
            fs = int(fs)
            # random sample signal if len > 8192 samples
            if data.shape[-1] >= each_file_length:
                idx = data.shape[-1] - each_file_length
                data = data[:, idx:idx + each_file_length]
            else:
                # Apply zero-padding along the second dimension
                padding_size = each_file_length - data.shape[-1]
                data = np.pad(data, ((0, 0), (0, padding_size)), mode='constant', constant_values=0) 

        #7 use the orde of cabrera sequence
        # data =standard_to_cabrera(data)
        
        return data


    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']
        leads = self.files.iloc[item]['leads']
        data = load_recording(self.files.iloc[item]['record'])
        
        # the phase of pre-processing:
        data = self.Pre_processing(fs, leads, data)
             

        return data, target