import os
import glob
import wfdb as wf
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np

from scipy.signal import resample
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

import torch

from utils import dist_transform,getWindow

def data_preprocess(args):
    
    PATH = args.data_path
    all_paths = sorted(glob.glob(PATH + '/*.dat'))
    all_paths = [paths[:-4] for paths in all_paths]

    windowed_data,windowed_beats = getWindow(all_paths)

    scaler = StandardScaler()
    mod_windowed_data = []
    dist_tran_data = []

    for window in range(len(windowed_data)):
        
        beats = ((windowed_beats[window] * 500/360).astype(int))
        if(len(beats) != 0):
            mod_windowed_data.append(scaler.fit_transform(resample(windowed_data[window],5000).reshape(-1,1)))
            dist_tran_data.append(dist_transform(5000,beats))

    X_train,X_test,y_train,y_test = train_test_split(np.array(mod_windowed_data)[:,:,0],np.array(dist_tran_data)[:,:,0],test_size = 0.1, random_state = 42)

    X_train_toTensor = torch.Tensor(X_train).type(torch.float)
    print(X_train_toTensor.size())
    X_test_toTensor = torch.Tensor(X_test).type(torch.float)
    print(X_test_toTensor.size())
    y_train_toTensor = torch.Tensor(y_train).type(torch.float)
    print(y_train_toTensor.size())
    y_test_toTensor = torch.Tensor(y_test).type(torch.float)
    print(y_test_toTensor.size())

    if not(os.path.exists('data')):
        os.mkdir('data')

    torch.save(X_train_toTensor, "data/ecgtoHR_train_data.pt")
    torch.save(y_train_toTensor, "data/ecgtoHR_train_labels.pt")
    torch.save(X_test_toTensor, "data/ecgtoHR_test_data.pt")
    torch.save(y_test_toTensor, "data/ecgtoHR_test_labels.pt")