import numpy as np
import scipy.signal
import scipy.io
import pickle
from glob import glob
import os
from tqdm import tqdm
import pandas as pd
from sklearn.preprocessing import StandardScaler

import torch
from torch.utils.data import TensorDataset,DataLoader
from torch.autograd import Variable

def custom_resample(ECG,fs):
    modified_ECG = []
    for i in range(int(len(ECG) * 500/fs)):
        modified_ECG.append(ECG[int(fs/500*i)].astype(float))
    return modified_ECG

def peak_correction(peak_locs,ecg_records):
    
    mod_peak_locs = []
    ecg_records = ecg_records.cpu().numpy()
    for j in range(len(peak_locs)): 
        mod_peak_locs.append(np.asarray([peak_locs[j][i] - 37 + np.argmax(ecg_records[j,0,peak_locs[j][i]-37:peak_locs[j][i]+37]) for i in range(len(peak_locs[j])) if(peak_locs[j][i]>37 and peak_locs[j][i]<5000-37)]))
    return mod_peak_locs
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              
def peak_finder(y_pred_array,x_test):
   
    fs_ = 500 
    peak_locs = []
    for i in range(y_pred_array.shape[0]):
        peak_locs.append(scipy.signal.find_peaks(-y_pred_array[i,:],distance = 120,height = -0.4, prominence = 0.035)[0])
        peak_locs[i] = peak_locs[i][(peak_locs[i] >= 0.5*fs_) & (peak_locs[i] <= 9.5*fs_)]
    modified_peak_locs = peak_correction(peak_locs,x_test)
    return modified_peak_locs

def compute_heart_rate(r_peaks):
    fs_ = 500
    r_peaks = r_peaks[(r_peaks >= 0.5*fs_) & (r_peaks <= 9.5*fs_)]
    return round( 60 * fs_ / np.mean(np.diff(r_peaks)))