import torch

import torch.nn as nn

import torch.nn.functional as functional

import numpy as np

import pandas as pd

import scipy.signal

from sklearn.preprocessing import MinMaxScaler, StandardScaler

def testDataEval(model, loader, criterion):

    """Test model on dataloader

    Arguments:

        model {torch object} -- Model   

        loader {torch object} -- Data Loader  

        criterion {torch object} -- Loss Function

    Returns:

        float -- total loss

    """

    model.eval()

    with torch.no_grad():

        total_loss = 0

        for (x,y) in loader:

            ecg,BR = x.cuda(),y.cuda()

            BR_pred = model(ecg)

            loss = criterion(BR_pred, BR)

            total_loss += loss

    return total_loss

def smooth(signal,window_len=50):

    """Compute moving average of specified window length

    Arguments:

        signal {ndarray} -- signal to smooth

    Keyword Arguments:

        window_len {int} -- size of window over which average is to be computed (default: {50})

    Returns:

        ndarray   -- smoothed signal

    """

    y = pd.DataFrame(signal).rolling(window_len,center = True, min_periods = 1).mean().values.reshape((-1,))

    return y

def findValleys(signal, prominence = 0.07):

    """Find valleys of distance transform to estimate breath positions   

    Arguments:

        signal {ndarray} -- transform to get breath positions

    Keyword Arguments:

        prominence {int} -- threshold prominence to detect peaks (default: {0.07})

    Returns:

        ndarray -- valley locations in signal

    """

    smoothened = smooth(-1*signal)

    valley_loc = scipy.signal.find_peaks(smoothened, prominence= prominence)[0]

    return valley_loc

def getBR(signal, model):

    """ Get Breathing Rate after passing ECG through Model

    Arguments:

        signal {torch tensor} -- input ECG signal

        model  -- ECG to BR model

    Returns:

        ndarray -- position of predicted valley and corresponding predicted transform

    """

    model.eval()

    with torch.no_grad():

        transformPredicted = model(signal)

    transformPredicted = transformPredicted.cpu().numpy().reshape((-1,))

    valleys = findValleys(transformPredicted)

    return valleys, transformPredicted

def save_model(exp_dir, epoch, model, optimizer,best_dev_loss):

    """ save checkpoint of model 

    Arguments:

        exp_dir {string} -- Path to checkpoint

        epoch {int} -- epoch at which model is checkpointed

        model -- model state to be checkpointed

        optimizer {torch optimizer object} -- optimizer state of model to be checkpoint

        best_dev_loss {float} -- loss of model to be checkpointed

    """

    out = torch.save(

        {

            'epoch': epoch,

            'model': model.state_dict(),

            'optimizer': optimizer.state_dict(),

            'best_dev_loss': best_dev_loss,

            'exp_dir':exp_dir

        },

        f=exp_dir + '/best_model.pt'

    )

def dist_transform(signal, ann):

    """ Compute distance transform of Respiration signaal based on breath positions

    Arguments:

        signal{ndarray} -- The ECG signal  

        ann{ndarray} -- The ground truth R-Peaks

    Returns:

       ndarray -- transformed signal

    """

    length = len(signal)

    transform = []

    sample = 0

    if len(ann) == 0:

        return None

    if len(ann) ==1:

        for i in range(length):

            transform.append(abs(i-ann[sample]))

    else:

        for i in range(length):

            if sample+1 == len(ann):

                for j in range(i,length):

                    transform.append(abs(j - nextAnn))

                break

            prevAnn = ann[sample]

            nextAnn = ann[sample+1]

            middle = int((prevAnn + nextAnn )/2) 

            if i < middle:

                transform.append(abs(i - prevAnn))

            elif i>= middle:

                transform.append(abs(i- nextAnn))

            if i == nextAnn:

                sample+=1

    transform = np.array(transform)

    minmaxScaler = MinMaxScaler()

    transform = minmaxScaler.fit_transform(transform.reshape((-1,1)))

    return transform

def getWindow(signal,ann, windows = 10, freq  = 125, overlap = 0.5):

    """Generate ECG and Respiration signals with annotations of specified window length

    Arguments:

        signal {2-D array} -- array containing ecg at index 0 and resp at index 1

        ann {list} -- annotations within specified window

    Keyword Arguments:

        windows {int} -- size of window in seconds (default: {5})

        freq {int} -- sampling rate in Hz (default: {125})

        overlap {float} -- percentage of overlap between windows (default: {0.5})

    Yields:

        tuple -- signals and correspoinding annotations

    """

    for start in range(0,len(signal),int((1-overlap)*freq*windows)):

        yield (signal[start: start + windows*freq, :],[x-start for x in ann if x >= start and x < start+windows*freq])