import numpy as np


def gauss(N, mu, sigma):
    x = np.arange(N)
    norm = sigma * np.sqrt(2*np.pi)
    return np.exp(-1/2 * (x - mu)**2 / sigma**2) / norm


def shift(N, n, x):
    """shift the center of the signal 'x' to time index 'n'."""
    y = np.zeros(N)
    xl = x.shape[0]
    s = n - xl // 2
    x_bi, x_ei = np.clip(xl // 2 - n, a_min=0, a_max=xl), np.clip(N + xl - xl // 2 - n - xl % 2, a_min=0, a_max=xl)
    y_bi, y_ei = np.clip(n - xl // 2, a_min=0, a_max=N), np.clip(n + xl // 2 + xl % 2, a_min=0, a_max=N)
    y[y_bi:y_ei] = x[x_bi:x_ei]
    return y


class SigQuality:
    """
    Compute the Signal-to-Noise Ratio of beats
    """
    gauss_beat_template_win_sec = 0.25542  #: gauss window width (as compared to beats in ssf function)
    gauss_beat_template_sigma_sec = 0.027  #: gauss bump half-width parameter (as compared to beats in ssf function)
    #gauss_amplitude = 2.0

    def __init__(self): pass

    def get_snr(self, fs, ssf, ssf_threshold, est_zxings):
        """Compute the Signal-to-Noise Ratio of beats, based on SSF function and detected beat locations."""
        sigma = fs * self.gauss_beat_template_sigma_sec
        W = int(fs * self.gauss_beat_template_win_sec)
        gb = gauss(W, W//2, sigma)
        # place gaussians on estimated beat locations
        ssf_est = np.zeros(ssf.shape[0])
        for i in est_zxings:
            ssf_est += shift(ssf.shape[0], i, gb)
        ssf_est /= gb[W//2]  # normalize amplitude to 1.0
        ssf_est = np.roll(ssf_est, int(sigma))  # shift to right (beat loc = gauss beginning, not center)
        #sqi_ref = ssf_est * (self.gauss_amplitude * ssf_threshold)  # set amplitude

        # penalty term = (where there should be no amplitude, according to gaussians)
        sqi_pen = 1.0 - ssf_est

        # sqi = ratio of signal energy in goal vs. in noise
        sqi_goal = np.sum(ssf_est * (ssf**2))
        sqi_noise = np.sum(sqi_pen * (ssf**2))

        # noise is everywhere, while signal is only around detected peaks - correct for this.
        goal_density = np.mean(np.clip(2*sigma / np.diff(est_zxings), a_min=0, a_max=1))
        sqi_goal /= goal_density
        sqi = 10 * (np.log10(sqi_goal) - np.log10(sqi_noise))

        return sqi