import wfdb as wf
import numpy as np
import glob
import os
import matplotlib.pyplot as plt
from tqdm import tqdm as tqdm
from scipy import signal
from sklearn.preprocessing import StandardScaler 
### Use standard scaler

def data_read(args):
    
    
    patient_no = args.patient_no
    patient_no -= 1

    path_dir = args.path_dir
    file_dirs = sorted(glob.glob(os.path.join(path_dir,'infant*')))   
    patient = {'ecg':[],'resp':[],'r_peaks':[],'resp_peaks':[],'brad_onset':[],'ecg_fs':[],'resp_fs':[]}
    
    for i in range(0,len(file_dirs),7):
        if (i == 7 * patient_no):
            
            print('-------- Acquiring data from patient {} --------'.format(patient_no + 1))

            ecg_file_name = os.path.splitext(file_dirs[i + 1])[0] 
            resp_file_name = os.path.splitext(file_dirs[i + 5])[0]
            qrsc_ext = os.path.splitext(file_dirs[i + 3])[1][1:]
            resp_ext = os.path.splitext(file_dirs[i + 6])[1][1:]
            brad_onset_ext = os.path.splitext(file_dirs[i + 6])[1][1:]
            atr_ext = os.path.splitext(file_dirs[i])[1][1:]
                
            ecg_sample_rate = wf.rdsamp(ecg_file_name)[-1]['fs']
            resp_sample_rate = wf.rdsamp(resp_file_name)[-1]['fs']
            ecg = wf.io.rdrecord(ecg_file_name).p_signal 
            resp = wf.io.rdrecord(resp_file_name).p_signal
            r_peaks_loc = wf.rdann(ecg_file_name,qrsc_ext).sample
            resp_peak_loc = wf.rdann(resp_file_name,resp_ext).sample
            brad_onset = wf.rdann(ecg_file_name,atr_ext).sample
                
            patient['ecg'].append(ecg)
            patient['resp'].append(resp)
            patient['r_peaks'].append(resp_peak_loc)
            patient['resp_peaks'].append(resp_peak_loc)
            patient['ecg_fs'].append(ecg_sample_rate)
            patient['resp_fs'].append(resp_sample_rate)
            patient['brad_onset'].append(brad_onset)

    return patient

def windowing_and_resampling_hr(patient):
    
    no_sec = 10
    final_ecg_sample_rate = 500
    final_resp_sample_rate = 50

    windowed_patient = {'ecg':[[] for i in range(10)],'resp':[[] for i in range(10)]}
    windowed_patient_overlap = {'ecg':[[] for i in range(10)],'resp':[[] for i in range(10)]}
    infant_no = 0

    for patients in patient[ 'ecg']:
        fs = patient['ecg_fs'][infant_no]
        window_len = no_sec * fs
        for i in range(len(patients) // (window_len-fs)):
            windowed_patient_overlap['ecg'][infant_no].append(patients[(window_len-fs) * i : ((window_len-fs) * i + window_len)])
            windowed_patient['ecg'][infant_no].append(patients[window_len * i : window_len * (i+1) ])
            if (final_ecg_sample_rate * no_sec) != window_len:
                windowed_patient_overlap['ecg'][infant_no][i] = signal.resample(windowed_patient_overlap['ecg'][infant_no][i],final_ecg_sample_rate * no_sec)
                windowed_patient['ecg'][infant_no][i] = signal.resample(windowed_patient['ecg'][infant_no][i],final_ecg_sample_rate * no_sec)
        
        infant_no += 1

    infant_no = 0
    for patients in patient['resp']:
        window_len = no_sec * patient['resp_fs'][infant_no] 
        for i in range(len(patients) // window_len):
            windowed_patient['resp'][infant_no].append(patients[window_len * i : window_len * (i+1)])
            if (final_resp_sample_rate * no_sec) != window_len:
                windowed_patient['resp'][infant_no][i] = custom_resample(windowed_patient['resp'][infant_no][i],patient['resp_fs'][infant_no])
        infant_no += 1

    return windowed_patient_overlap,windowed_patient

def windowing_and_resampling_br(patient):
    
    scaler = StandardScaler()
    no_sec = 10
    final_ecg_sample_rate = 500
    final_resp_sample_rate = 50

    windowed_patient = {'ecg':[[] for i in range(10)],'resp':[[] for i in range(10)]}
    windowed_patient_overlap = {'ecg':[[] for i in range(10)],'resp':[[] for i in range(10)]}
    infant_no = 0
    overlap_percent = 50
    overlap = int(50 / 100 * 5000)
    
    for patients in patient[ 'ecg']:
        
        fs = patient['ecg_fs'][infant_no]
        window_len = no_sec * fs
        for i in range(len(patients) // (window_len-overlap)):
            windowed_patient_overlap['ecg'][infant_no].append(scaler.fit_transform(patients[(window_len-overlap) * i : ((window_len-overlap) * i + window_len)]))
            if (final_ecg_sample_rate * no_sec) != window_len:
                windowed_patient_overlap['ecg'][infant_no][i] = signal.resample(windowed_patient_overlap['ecg'][infant_no][i],final_ecg_sample_rate * no_sec)
        
        infant_no += 1
    return windowed_patient_overlap

def custom_resample(resp,fs):
    
    modified_resp = []
    for i in range(int(len(resp) * 50/fs)):
        modified_resp.append(resp[int(fs/50*i)].astype(float)) 
    return np.asarray(modified_resp)