import numpy as np

import os

import pandas as pd

import wfdb as wf

import argparse

from glob import glob

from tqdm import tqdm

from scipy.signal import resample

from sklearn.preprocessing import StandardScaler, MinMaxScaler

from sklearn.model_selection import train_test_split

import wget

import torch

from utils import dist_transform,getWindow

def data_preprocess(args):

    """ Preprocess data and create train - validate split

    """

    dat_path = os.path.join(args.data_path,'*.dat')

    paths = glob(dat_path)

    paths= sorted([path[:-4] for path in paths if path[-5] != "n"] )

    fs = args.sampling_freq

    fs_upsample = args.upsample_freq

    WINDOWS = args.window_length 

    ecgSignals = []

    BRSignals = []

    BRAnn1 = []

    BRAnn2 = []

    for path in tqdm(paths):

        ann    = wf.rdann(path,'breath')

        samples = np.array(ann.sample)

        aux_note = np.array(ann.aux_note)

        ann1 = samples[(aux_note == "ann1")]

        ann2 = samples[(aux_note == "ann2")]

        record = wf.io.rdrecord(path)

        ecgSignals.append(record.p_signal[:,record.sig_name.index('II,')])

        BRSignals.append(record.p_signal[:,record.sig_name.index('RESP,')])

        BRAnn1.append(ann1)

        BRAnn2.append(ann2)

    ecgSignals = np.array(ecgSignals,ndmin = 2)

    BRSignals = np.array(BRSignals, ndmin = 2)

    signals = np.stack([ecgSignals,BRSignals], axis= -1 )

    inputECG = []

    groundTruth = []

    for i in tqdm(range(len(signals))):

        generateSignals = getWindow(signals[i],BRAnn2[i],windows=WINDOWS)

        for sig, ann in generateSignals:

            ecg = sig[:,0]

            br = sig[:,1]

            if len(ecg) == 1 or len(ann) == 0:

                break

            resampled = resample(ecg, WINDOWS*fs_upsample)

            scaler = StandardScaler()

            resampled = scaler.fit_transform(resampled.reshape((-1,1)))

            transform = dist_transform(br,ann)

            if resampled.shape == (fs_upsample*WINDOWS,1) and transform.shape == (WINDOWS*fs,1):

                inputECG.append(resampled.reshape((1,-1)))

                groundTruth.append(transform.reshape((1,-1)))

    X_train,X_test,y_train,y_test = train_test_split(np.array(inputECG),np.array(groundTruth),test_size = 0.2, random_state = 42)

    X_train_toTensor = torch.Tensor(X_train).type(torch.float)

    X_test_toTensor = torch.Tensor(X_test).type(torch.float)

    y_train_toTensor = torch.Tensor(y_train).type(torch.float)

    y_test_toTensor = torch.Tensor(y_test).type(torch.float)

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

        os.mkdir('data')

    torch.save(X_train_toTensor, "data/ecgtoBR_train_data.pt")

    torch.save(y_train_toTensor, "data/ecgtoBR_train_labels.pt")

    torch.save(X_test_toTensor, "data/ecgtoBR_test_data.pt")

    torch.save(y_test_toTensor, "data/ecgtoBR_test_labels.pt")