import numpy as np
import scipy.signal

#----------------------------------------------------------------
# Sample entropy:
#----------------------------------------------------------------

'''
Inspired in: 

Gao, Q., Yang, Y., Kang, Q. et al. EEG-based Emotion Recognition with Feature Fusion Networks. Int. J. Mach. Learn. & Cyber. 13, 421–429 (2022). 
https://doi.org/10.1007/s13042-021-01414-5


The sample entropy [6] is used to test the complexity of the timing of EEG signals.

Let the original data be a time series of length N, expressed as x(1), x(2), \ldots , x(N)which can compose m vectors in order :
\begin{aligned} \begin{array}{c} X_m\left( i \right) =\left[ x\left( i \right) ,x\left( i+1 \right) ,...,x\left( i+m-1 \right) \right] \\ 1\le i\le N-m+1\\ \end{array} \end{aligned}

The distance of X_m\left( i \right) and X_m\left( j \right) can been defined as the one with the largest difference between the two corresponding elements.

\begin{aligned} \begin{array}{r} d\left[ X_m\left( i \right) ,X_m\left( i \right) \right] =\max {|}x\left( i+k \right) -x\left( j+k \right) |\\ \end{array} \end{aligned}

The statistical value is less than the ratio of the threshold r to the total number N-M-1:

\begin{aligned} B_{i}^{m}\left( r \right) =\frac{1}{N-m-1}{mum}\left\{ d\left[ X_m\left( i \right) ,X_m\left( j \right) \right] <r \right\} \end{aligned}

Find all the average values of i:

\begin{aligned} B^m\left( r \right) =\frac{1}{N-m}\sum _{i=1}^{Nm}{B}_{i}^{m}\left( r \right) \end{aligned}

For m+1-dimensional vectors, there are also :

\begin{aligned} A_{i}^{m}\left( r \right) =\frac{1}{N-m-1}{num}\left\{ d\left[ X_{m+1}\left( i \right) ,X_{m+1}\left( j \right) \right] <r \right\} \end{aligned}

Find all the average values of j:

\begin{aligned} A^m\left( r \right) =\frac{1}{N-m}\sum _{i=1}^{N-m}{A}_{i}^{m}\left( r \right) \end{aligned}

The sequence of this sample entropy:

\begin{aligned} {SampEn}\left( m,r \right) =\lim _{N\rightarrow \infty }\left\{ -\ln \left[ A^m\left( r \right) /B^m\left( r \right) \right] \right\} \end{aligned}
'''


def max_distance(x_i, x_j):
    """
    Calculate the maximum distance between two vectors.
    
    Args:
    x_i, x_j: List of integers or floats.

    Returns:
    The maximum absolute difference between two corresponding elements of the input vectors.
    """
    return max([abs(a - b) for a, b in zip(x_i, x_j)])

def chebyshev_norm(x_i, x_j):
    """
    Calculate the Chebyshev norm (maximum absolute difference) between two vectors.
    
    Args:
    x_i, x_j: List of integers or floats.

    Returns:
    The maximum absolute difference between two corresponding elements of the input vectors.
    """
    return max(np.abs(x_i - x_j))

def _phi(m, r, time_series):
    """
    Calculate the ratio of the number of matches for m-length sequences.
    
    Args:
    m: Integer, length of sequences to be compared.
    r: Float, tolerance for accepting matches.
    time_series: List of integers or floats representing the time series.

    Returns:
    Float representing the ratio.
    """
    x = np.array([time_series[i: i + m] for i in range(len(time_series) - m + 1)])
    B = np.sum([np.sum(np.abs(x[:, None] - x[i, :]) <= r, axis=0) - 1 for i in range(x.shape[0])])
    return B / ((len(time_series) - m + 1) * (len(time_series) - m))

def sample_entropy(time_series, m=2, r=None):
    """
    Calculate the Sample Entropy of a time series.
    
    Args:
    time_series: List of integers or floats representing the time series.
    m: Integer, length of sequences to be compared (default: 2).
    r: Float, tolerance for accepting matches (default: 0.2 * standard deviation of the time series).

    Returns:
    Float representing the Sample Entropy.
    """
    if r is None:
        r = 0.2 * np.std(time_series)
    return -np.log(_phi(m+1, r, time_series) / _phi(m, r, time_series))







#----------------------------------------------------------------
# Differential entropy:
#----------------------------------------------------------------

'''
Differential entropy  is used to measure the complexity of a continuous random variable, which extends to Shannon entropy. The original calculation formula of differential entropy is defined as:

\begin{aligned} \begin{array}{c} h\left( X \right) =-\int _X{f}\left( x \right) \log \left( f\left( x \right) \right) dx \end{array} \end{aligned}

If a random variable obeys the Gaussian distribution N\left( u,\sigma ^2 \right) , the differential entropy can simply be calculated by the following formulation:

\begin{aligned} \begin{array}{c} h\left( X \right) =-\int _{\infty }^{\infty }{\frac{1}{\sqrt{2\pi \sigma ^2}}}\exp \frac{\left( x-\mu \right) ^2}{2\sigma ^2}\log \frac{1}{\sqrt{2\pi \sigma ^2}}\exp \frac{\left( x-\mu \right) ^2}{2\sigma ^2}dx\\ =\frac{1}{2}\log 2\pi e\sigma ^2\\ \end{array} \end{aligned}
'''

def differential_entropy(time_series):
    """
    Calculate the Differential Entropy of a Gaussian-distributed random variable.
    
    Args:
    time_series: List of integers or floats representing the time series.

    Returns:
    Float representing the Differential Entropy.
    """
    sigma = np.std(time_series)
    return 0.5 * np.log(2 * np.pi * np.e * sigma**2)



#----------------------------------------------------------------
# Power spectral density (PSD):
#----------------------------------------------------------------
'''
We apply 256-point Short-Time Fourier Transform(STFT) without overlapping a Hanning window, which can get a two-dimensional function of frequency and time to extract useful information. Short-time Fourier transform (STFT) of discrete-time signal x[n] is calculated as a discrete-time Fourier transform (DTFT) of windowed sequence.

\begin{aligned} X_{n}\left( e^{j \Omega }\right) =\sum _{m=-\infty }^{\infty } x[m] w[n-m] e^{-j \omega m} \end{aligned}


The Short-Time Fourier Transform (STFT) is a commonly used tool for analyzing the frequency of a signal over time. 
The Power Spectral Density (PSD) is a measure of a signal's power per unit of frequency and can be obtained by 
squaring the magnitude of the STFT.


'''

def calculate_psd(time_series, fs=1.0, nperseg=256, noverlap=0):
    """
    Calculate the Power Spectral Density (PSD) of a time series.
    
    Args:
    time_series: List of integers or floats representing the time series.
    fs: Float, the sample rate of the time series (default: 1.0).
    nperseg: Integer, the length of each segment for the STFT (default: 256).
    noverlap: Integer, the number of points to overlap between segments (default: 0).

    Returns:
    freqs: Array of floats, the frequencies at which the PSD was calculated.
    psd: Array of floats, the calculated PSD.
    """
    _, _, Zxx = scipy.signal.stft(time_series, fs=fs, nperseg=nperseg, noverlap=noverlap, window='hann')
    psd = np.abs(Zxx)**2
    freqs = np.fft.fftfreq(nperseg, 1/fs)
    
    return freqs, psd



#----------------------------------------------------------------
# Hjorth:
#----------------------------------------------------------------
'''
Hjorth provides a fast time-domain signal calculation approach for three important features, including Activity, Mobility, and Complexity. It has a wide range of applications in the field of physiological signal processing. The Activity parameter represents the signal power, which can be the spectral value in the frequency domain.

\begin{aligned} Activity\ =\ var\left( y\left( t \right) \right) \end{aligned}

The Mobility parameter represents the ratio of the average frequency or the standard deviation of the power spectrum.

\begin{aligned} Mobility\ =\ \sqrt{\frac{var\left( \frac{dy\left( t \right) }{dt} \right) }{var\left( y\left( t \right) \right) }} \end{aligned}

The Complexity parameter indicates the similarity of signal and sinusoidal signal.

\begin{aligned} Complexity\ =\ \sqrt{\frac{Mobility\left( \frac{dy\left( t \right) }{dt} \right) }{Mobility\left( y\left( t \right) \right) }} \end{aligned}
'''

def hjorth_params(time_series):
    """
    Calculate Hjorth parameters (Activity, Mobility, and Complexity) of a time series.
    
    Args:
    time_series: List of integers or floats representing the time series.

    Returns:
    activity: Float, the signal power.
    mobility: Float, the standard deviation of the power spectrum.
    complexity: Float, indicating the similarity of the signal and a sinusoidal signal.
    """
    # First derivative of the time series
    first_derivative = np.diff(time_series)

    # Second derivative of the time series
    second_derivative = np.diff(time_series, 2)

    # Activity
    activity = np.var(time_series)

    # Mobility
    mobility = np.sqrt(np.var(first_derivative) / activity)

    # Complexity
    complexity = np.sqrt(np.var(second_derivative) / np.var(first_derivative)) / mobility

    return activity, mobility, complexity