# A collection of heartrate variability algorithms for
# both the timedomain and frequency domain.
#
# Copyright (C) 2019 Luis Howell & Bernd Porr
# GPL GNU GENERAL PUBLIC LICENSE Version 3, 29 June 2007
#


import numpy as np
import random
import subprocess
from scipy.interpolate import interp1d
from gatspy.periodic import LombScargleFast

class HRV:
    """
    Heartrate variability class which calcualtes the standard HRV
    parameters with the help of Python functions and for cross
    validation also via the physionet's get_hrv script.
    """

    def __init__(self, sampling_frequency):
        """
        Constructor takes the sampling frequency.
        All rr_sample data is in sample number and
        will assume it's at this sampling rate.
        """

        self.fs = float(sampling_frequency)
        self.period = 1.0/sampling_frequency

    
    def _intervals(self, rr_samples):
        """Calculate the RR intervals in ms from sample numbers.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: RR intervals in milliseconds
        :rtype: ndarray
        """

        rr_intervals = np.diff(np.array(rr_samples)*self.period*1000)

        return rr_intervals

    
    def _timestamps(self, rr_samples):
        """Calculate the timestamps in ms from sample locations.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: The timestamps in milliseconds
        :rtype: array_like
        """

        ts = np.array(rr_samples)*self.period*1000

        return ts

    
    def _succ_diffs(self, rr_samples):
        """Calculate the successive differences of R peaks.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: The successive differences of R peaks
        :rtype: ndarray
        """

        rr_ints = self._intervals(rr_samples)

        succ_diffs = []

        for i in range(len(rr_ints)-1):

            diff = rr_ints[i+1] - rr_ints[i]            
            succ_diffs.append(diff)

        return np.array(succ_diffs)

   
    def SDNN(self, rr_samples, normalise=False):
        """Calculate SDNN, the standard deviation of NN intervals.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :param normalise: normalise the SDNN against the average RR interval, defaults to False
        :type normalise: bool, optional
        :return: SDNN, the standard deviation of NN intervals
        :rtype: float
        """

        rr_intervals = self._intervals(rr_samples) 
        rr_std = np.std(rr_intervals)

        if normalise:
            rr_mean_interval = np.mean(rr_intervals)
            rr_std = rr_std/rr_mean_interval

        return rr_std

    
    def SDANN(self, rr_samples, average_period=5.0, normalise=False):
        """Calculate SDANN, the standard deviation of the average RR intervals calculated over short periods.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :param average_period: The averging period in minutes, defaults to 5.0
        :type average_period: float, optional
        :param normalise: normalise the SDANN against the average RR interval, defaults to False
        :type normalise: bool, optional
        :return: SDANN, the standard deviation of the average RR intervals calculated over short periods
        :rtype: float
        """

        average_period_samples = int(self.fs*average_period*60)
        average_rr_intervals = []

        sections = int((np.max(rr_samples)/average_period_samples)+0.5)

        if sections<1:
                sections = 1

        for i in range(sections):

                idx = np.where((rr_samples>=(i*average_period_samples)) &
                               (rr_samples<((i+1)*average_period_samples)))
                idx = idx[0]
                section_rr = rr_samples[idx[0]:idx[-1]+1]

                avg_rr_int = np.mean(self._intervals(section_rr))
                average_rr_intervals.append(avg_rr_int)

        rr_std = np.std(average_rr_intervals)

        if normalise:
                rr_mean_interval = np.mean(average_rr_intervals)
                rr_std = rr_std/rr_mean_interval

        return rr_std


    def RMSSD(self, rr_samples, normalise = False):
        """Calculate RMSSD (root mean square of successive differences).
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :param normalise: normalise the RMSSD against the average RR interval, defaults to False
        :type normalise: bool, optional
        :return: RMSSD (root mean square of successive differences)
        :rtype: float
        """

        succ_diffs = self._succ_diffs(rr_samples)
        succ_diffs = succ_diffs*succ_diffs

        rms = np.sqrt(np.mean(succ_diffs))

        if normalise:
            rms = rms / np.mean(self._intervals(rr_samples))

        return rms


    def SDSD(self, rr_samples):
        """Calculate SDSD (standard deviation of successive differences), the standard deviation of the successive differences between adjacent NNs.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: SDSD (standard deviation of successive differences)
        :rtype: float
        """

        succ_diffs = self._succ_diffs(rr_samples)

        return np.std(succ_diffs)        

    
    def NN50(self, rr_samples):
        """Calculate NN50, the number of pairs of successive NNs that differ by more than 50 ms.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: NN50
        :rtype: float
        """

        succ_diffs = self._succ_diffs(rr_samples)

        over_50 = np.where(abs(succ_diffs)>50)
        over_50 = over_50[0]

        return len(over_50)


    def pNN50(self, rr_samples):
        """Calculate pNN50, the proportion of NN50 divided by total number of NNs.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: pNN50
        :rtype: float
        """

        return self.NN50(rr_samples)/(len(rr_samples)-1)


    def NN20(self, rr_samples):
        """Calculate NN20, the number of pairs of successive NNs that differ by more than 20 ms.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: NN20
        :rtype: float
        """

        succ_diffs = self._succ_diffs(rr_samples)

        over_20 = np.where(abs(succ_diffs)>20)
        over_20 = over_20[0]

        return len(over_20)


    def pNN20(self, rr_samples):
        """Calculate pNN20, the proportion of NN20 divided by total number of NNs.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: pNN20
        :rtype: float
        """

        return self.NN20(rr_samples)/(len(rr_samples)-1)

    
    def HR(self, rr_samples):
        """Calculate heart-rates from R peak samples.
        
        :param rr_samples: R peak sample locations
        :type rr_samples: array_like
        :return: Heart-rates in BPM
        :rtype: ndarray
        """
        
        rr_intervals = np.diff(rr_samples)
        heart_rates = 60.0/(rr_intervals/float(self.fs))

        return heart_rates


    def add_rr_error(self, rr_samples, error):
        """
        Adds jitter to the heartrate timestamps. 
        The error and the rr_samples are in timestamps.
        Returns the noisy timestamps in samples.
        """
        if error==0:
            return rr_samples

        error_values = np.random.randint(-abs(error), abs(error)+1, len(rr_samples))
        noisy_rr_samples = rr_samples+error_values

        return noisy_rr_samples

    
    def fAnalysis(self, rr_samples):
        """
        Frequency analysis to calc self.lf, self.hf, returns the LF/HF-ratio and
        also calculates the spectrum as pairs of (self.f_hr_axis,self.f_hr).
        The input arrary is in sample points where R peaks have been detected.
        """
        # discrete timestamps
        self.hr_discrete = self._intervals(rr_samples) / 1000
        # hr positions in time
        self.t_hr_discrete  = [i/self.fs for i in rr_samples[1:]]
        # now let's create function which approximates the hr(t) relationship
        self.hr_func = interp1d(self.t_hr_discrete, self.hr_discrete)
        # we take 1024 samples for a linear time array for hr(t)
        nsamp = 1000
        # linear time array for the heartrate
        self.t_hr_linear = np.linspace(self.t_hr_discrete[1],
                                       self.t_hr_discrete[len(self.t_hr_discrete)-2],
                                       num=nsamp)
        # duration in secs of the heartrate array minus the ends b/c of the approx
        duration = self.t_hr_discrete[len(self.t_hr_discrete)-2] - self.t_hr_discrete[1];
        # heartrate linearly approximated between discrete samples
        self.hr_linear = self.hr_func(self.t_hr_linear)
        model = LombScargleFast().fit(self.t_hr_discrete, self.hr_discrete, 1E-2)
        fmax = 1
        fmin = 0.01
        df = (fmax - fmin) / nsamp
        self.f_hr = model.score_frequency_grid(fmin, df, nsamp)
        self.f_hr_axis = fmin + df * np.arange(nsamp)
        # lf
        self.lf = 0
        # hf
        self.hf = 0
        for i in range(0,int(nsamp/2)):
            if (self.f_hr_axis[i] >= 0.04) and (self.f_hr_axis[i] <= 0.15):
                self.lf = self.lf + self.f_hr[i]
            if (self.f_hr_axis[i] >= 0.15) and (self.f_hr_axis[i] <= 0.4):
                self.hf = self.hf + self.f_hr[i]
        # hf
        return self.lf/self.hf

    
    def hrv_toolkit(self, rr_samples):
        """
        Calculating the HRV parameters with the HRV toolkit.
        It calls the commandline file get_hrv and then parses the output.
        To be able to run check out the subdir HRV.src and compile/install
        the binaries in /usr/local/bin.
        rr_samples are the samples numbers of the R-peaks.
        Returns the lf/hf ratio and also provides self.sdnn, 
        self.rmssd, self.lf and self.hf.
        """
        m = np.hstack((np.vstack(self._timestamps(rr_samples[1:])/1000),
                       np.vstack(self._intervals(rr_samples)/1000)))
        np.savetxt("hr.txt",m,delimiter = ' ', newline=" N\n")
        r = subprocess.check_output(['get_hrv', '-s', '-R', 'hr.txt'])
        r = str(r)
        a = r.split("\\n")
        for b in a:
            x = b.split("=")
            #print(x)
            if "SDNN" in x[0]:
                self.sdnn = float(x[1])
            elif "rMSSD" in x[0]:
                self.rmssd = float(x[1])
            elif "LF PWR" in x[0]:
                self.lf = float(x[1])
            elif "HF PWR" in x[0]:
                self.hf = float(x[1])
            elif "LF/HF" in x[0]:
                self.lfhf = float(x[1])
        return self.lfhf