# © SleepECG developers
#
# License: BSD (3-clause)

"""Heartbeat detection and detector evaluation."""

import warnings
from typing import NamedTuple

import numpy as np
import scipy.interpolate
import scipy.signal
import scipy.stats

_all_backends = ("c", "numba", "python")
_available_backends = list(_all_backends)

try:
    from sleepecg._heartbeat_detection import _squared_moving_integration, _thresholding
except ImportError:
    _available_backends.remove("c")

try:
    from numba import jit
except ImportError:
    _available_backends.remove("numba")


# cache sos-filter created with scipy.signal.butter to reduce runtime
_sos_filters: dict[float, np.ndarray] = {}


def detect_heartbeats(ecg: np.ndarray, fs: float, backend: str = "c") -> np.ndarray:
    """
    Detect heartbeats in an ECG signal.

    This is a modified version of the beat detection algorithm using adaptive thresholds
    described by Pan & Tompkins in 1985.

    Modifications/additions to the
    [original algorithm](https://doi.org/10.1109/TBME.1985.325532) are:

    - Instead of a hardware filter adjusted to the sampling frequency of the MIT-BIH
      Arrhythmia Database, a second order bandpass with cutoff frequencies 5 and 30 Hz
      created via `scipy.signal.butter()` is used.
    - A bidirectional filter is used to remove filter delay.
    - The integration window is centered on the filtered signal, i.e. a peak in the filtered
      signal corresponds to a plateau in the integrated signal, not a saddle in the rising
      edge. This lets the adaptive threshold for the integrated signal remain at a higher
      level, which is less susceptible to noise.
    - Learning phase 1 is not described in detail in the original paper. This implementation
      uses maximum and mean values inside the first two seconds to initialize
      SPKI/SPKF/NPKI/NPKF. Details are provided in the `_thresholding` code.
    - In addition to the original searchback criterion, a searchback is also performed if no
      peak is found during the first second of the signal or no second peak is found 1.5 s
      after the first one. This ensures correct behaviour at signal start in case an
      unusually large peak during learning phase 1 messes up threshold initialization.
    - After an unsuccessful searchback, the procedure is repeated in the same interval with
      further reduced thresholds, up to 16 times.

    Parameters
    ----------
    ecg : np.ndarray
        ECG signal. Note that the unit of the data does not matter (the algorithm will
        return similar results regardless of the scaling of the data).
    fs : float
        Sampling frequency in Hz. For best results, a sampling frequency of at least 100 Hz
        is recommended.
    backend : {'c', 'numba', 'python'}
        Which implementation of the squared moving integration and thresholding algorithm to
        use. If available, `'c'` is the fastest implementation, `'numba'` is about 25%
        slower, and `'python'` is about 20 times slower but provided as a fallback. By
        default `'c'`.

    Returns
    -------
    np.ndarray
        Indices of detected heartbeats.

    Examples
    --------
    Detect heartbeats in a short electrocardiogram:

    >>> from sleepecg import detect_heartbeats, get_toy_ecg
    >>> ecg, fs = get_toy_ecg()  # 5 min of ECG data at 360 Hz
    >>> heartbeats = detect_heartbeats(ecg, fs)
    >>> print(f"{len(heartbeats)} heartbeats detected")
    478 heartbeats detected
    """
    if backend not in _all_backends:
        raise ValueError(
            f"Invalid backend for detect_heartbeats: {backend!r}. "
            f"Possible options are: {_all_backends}."
        )
    if backend not in _available_backends:
        fallback = _available_backends[0]
        warnings.warn(f"Backend {backend!r} not available, using {fallback!r} instead.")
        backend = fallback

    # check for flat data at the beginning to avoid messing up detection thresholds
    if len(ecg) < 2:
        raise ValueError("ECG signal must have more than one sample.")
    if ecg[1] == ecg[0]:
        idx_nonflat = np.nonzero(ecg != ecg[0])[0]
        if not len(idx_nonflat):
            raise ValueError("ECG signal is flat. Please check your data.")
        first_nonflat = idx_nonflat[0]  # first non-flat index
    else:
        first_nonflat = 0

    # cache filter in global variable to speed up runtime (especially for short signals)
    try:
        sos = _sos_filters[fs]
    except KeyError:
        sos = scipy.signal.butter(
            N=2,
            Wn=(5, 30),
            btype="bandpass",
            output="sos",
            fs=fs,
        )
        _sos_filters[fs] = sos

    # filtering bidirectionally removes filter delay. Note that we start
    # filtering the signal at the first non-flat index.
    filtered_ecg = scipy.signal.sosfiltfilt(sos, ecg[first_nonflat:])

    # scipy.signal.sosfiltfilt returns an array with negative strides. Both `np.correlate`
    # and `_thresholding` require contiguity, so ensuring this here once reduces total
    # runtime.
    filtered_ecg = np.ascontiguousarray(filtered_ecg)

    # five-point derivative as described by Pan & Tompkins
    derivative = np.correlate(filtered_ecg, np.array([-1, -2, 0, 2, 1]), mode="same")

    moving_window_width = int(0.15 * fs)  # 150ms

    if backend == "c":
        integrated_ecg = _squared_moving_integration(derivative, moving_window_width)
        beat_mask = _thresholding(filtered_ecg, integrated_ecg, fs)
    elif backend == "numba":
        integrated_ecg = _squared_moving_integration_numba(derivative, moving_window_width)
        beat_mask = _thresholding_numba(filtered_ecg, integrated_ecg, fs)
    elif backend == "python":
        integrated_ecg = np.convolve(
            derivative**2,
            np.ones(moving_window_width),
            mode="same",
        )
        beat_mask = _thresholding_py(filtered_ecg, integrated_ecg, fs)

    heartbeats = np.where(beat_mask)[0]
    heartbeats += first_nonflat  # Add back the number of flat samples
    return heartbeats


class _CompareHeartbeatsResult(NamedTuple):
    TP: np.ndarray
    FP: np.ndarray
    FN: np.ndarray


def compare_heartbeats(
    detection: np.ndarray,
    annotation: np.ndarray,
    max_distance: int = 0,
) -> _CompareHeartbeatsResult:
    """
    Determine correctness of detection results.

    Determine true positives (TP), false positives (FP) and false negatives (FN) for an
    array of detected heartbeat indices based on an array of annotated heartbeat indices.
    Since neither annotations nor automated detectors usually hit the peak perfectly,
    detected peaks no further than `max_distance` in both directions from an annotated peak
    are considered true positives.

    Parameters
    ----------
    detection : np.ndarray
        Detected heartbeat indices.
    annotation : np.ndarray
        Annotated heartbeat indices.
    max_distance : int, optional
        Maximum distance between indices to consider as the same peak, by default `0`.

    Returns
    -------
    TP : np.ndarray
        True positives, i.e. actual heartbeats detected as heartbeats.
    FP : np.ndarray
        False positives, i.e. non-heartbeats detected as heartbeats.
    FN : np.ndarray
        False negatives, i.e. actual heartbeats not detected as heartbeats.
    """
    if len(detection) == 0:
        return _CompareHeartbeatsResult(
            TP=np.array([]),
            FP=np.array([]),
            FN=annotation.copy(),
        )

    max_len = max(np.max(detection), np.max(annotation)) + 1
    detection_mask = np.zeros(max_len, dtype=bool)
    detection_mask[detection] = 1
    annotation_mask = np.zeros(max_len, dtype=bool)
    annotation_mask[annotation] = 1

    fuzzy_filter = np.ones(max_distance * 2 + 1, dtype=bool)
    detection_mask_fuzzy = np.convolve(detection_mask, fuzzy_filter, mode="same")
    annotation_mask_fuzzy = np.convolve(annotation_mask, fuzzy_filter, mode="same")

    return _CompareHeartbeatsResult(
        TP=np.where(detection_mask & annotation_mask_fuzzy)[0],
        FP=np.where(detection_mask & ~annotation_mask_fuzzy)[0],
        FN=np.where(annotation_mask & ~detection_mask_fuzzy)[0],
    )


class _RRISimilarityResult(NamedTuple):
    pearsonr: float
    spearmanr: float
    rmse: float


def rri_similarity(
    detection: np.ndarray,
    annotation: np.ndarray,
    fs_resample: float = 4,
) -> _RRISimilarityResult:
    """
    Calculate measures of similarity between RR intervals.

    RR intervals are calculated from detected and annotated heartbeat indices. The RR time
    series is then resampled to frequency `fs_resample` in the timespan common to both
    detection and annotation. Pearson's and Spearman's correlation coefficients as well as
    the root mean square error are returned.

    Parameters
    ----------
    detection : np.ndarray
        Detected heartbeat indices.
    annotation : np.ndarray
        Annotated heartbeat indices.

    Returns
    -------
    pearsonr : float
        Pearson correlation coefficient between resampled RR time series.
    spearmanr : float
        Spearman correlation coefficient between resampled RR time series.
    rmse : float
        Root mean square error between resampled RR time series.
    """
    rr_ann = np.diff(annotation)
    rr_det = np.diff(detection)

    t_ann = annotation[1:]
    t_det = detection[1:]

    start = max(min(t_ann), min(t_det))
    end = min(max(t_ann), max(t_det))
    t_new = np.arange(start, end, 1 / fs_resample)

    interp_det = scipy.interpolate.interp1d(t_det, rr_det)
    interp_ann = scipy.interpolate.interp1d(t_ann, rr_ann)
    det_resampled = interp_det(t_new)
    ann_resampled = interp_ann(t_new)

    return _RRISimilarityResult(
        scipy.stats.pearsonr(det_resampled, ann_resampled)[0],
        scipy.stats.spearmanr(det_resampled, ann_resampled)[0],
        np.sqrt(np.mean((det_resampled - ann_resampled) ** 2)),
    )


def _squared_moving_integration_py(x: np.ndarray, window_length: int) -> np.ndarray:
    """
    Perform squaring and moving integration of an array.

    Parameters
    ----------
    x : np.ndarray
        The input array.
    window_length : int
        Length of the moving window in samples.

    Returns
    -------
    np.ndarray
        The squared and integrated array.
    """
    signal_len = len(x)
    if not 0 < window_length <= signal_len:
        raise ValueError("window_length has to be 0 < window_length <= len(x)")

    output = np.empty_like(x)

    # create a circular buffer to store values inside integration window
    integration_buffer = np.zeros(window_length)
    sum = 0

    # the integration window is centered on the original signal, for even window_length the
    # behavior of np.convolve with a constant window of even length is replicated (i.e the
    # window is off-center to the left)
    window_length_half = (window_length + 1) // 2

    # during the first `window_length/2` samples there is no output since the integration
    # window's center would be at a negative index of the input
    for i in range(window_length_half):
        square = x[i] * x[i]
        integration_buffer[i % window_length] = square
        sum += square

    for i in range(window_length_half, signal_len):
        output[i - window_length_half] = sum  # write to 'window center'
        sum -= integration_buffer[i % window_length]
        square = x[i] * x[i]
        integration_buffer[i % window_length] = square
        sum += square

    # the end of the x signal is reached, so the integration window is built down and the
    # last `window_length/2` entries of the output are filled
    for i in range(signal_len, signal_len + window_length_half):
        output[i - window_length_half] = sum
        sum -= integration_buffer[i % window_length]

    return output


def _thresholding_py(
    filtered_ecg: np.ndarray,
    integrated_ecg: np.ndarray,
    fs: float,
) -> np.ndarray:
    """
    Perform adaptive thresholding based on Pan and Tompkin's algorithm.

    Parameters
    ----------
    filtered_ecg : np.ndarray
        Bandpass filtered ECG.
    integrated_ecg : np.ndarray
        Squared and moving window integrated bandpass filtered ECG.
    fs : float
        Sampling frequency in Hz.

    Returns
    -------
    np.ndarray
        Array containing `1` for every sample in `filtered_ecg` identified as a heartbeat.
    """
    signal_len = len(filtered_ecg)
    beat_mask = np.zeros_like(filtered_ecg)  # numba can't work with dtype=bool

    REFRACTORY_SAMPLES = int(0.2 * fs)  # 200ms
    T_WAVE_WINDOW = int(0.36 * fs)  # 360ms

    # --------------------------------------------------------------------------------------
    # Learning Phase 1
    # --------------------------------------------------------------------------------------
    # Pan & Tompkins mention a learning phase to initialize detection thresholds based upon
    # signal and noise peaks detected during the first two seconds. The exact initialization
    # process is not described. The adaptive thresholds are calculated based on running
    # estimates of signal and noise peaks (`SPKF` and `NPKF` for the filtered signal).
    # Assuming constant peak amplitudes, those values converge towards the signal peak
    # amplitude and noise peak amplitude, respectively. Therefore, SPKF/SPKI are assumed to
    # be the maximum values of the filtered/integrated signal during the learning phase.
    # Accordingly, NPKF/NPKI are initialized to the mean values during the first two
    # seconds.
    SPKF = np.max(filtered_ecg[: int(2 * fs)])
    NPKF = np.mean(filtered_ecg[: int(2 * fs)])
    SPKI = np.max(integrated_ecg[: int(2 * fs)])
    NPKI = np.mean(integrated_ecg[: int(2 * fs)])
    threshold_I1 = NPKI + 0.25 * (SPKI - NPKI)
    threshold_F1 = NPKF + 0.25 * (SPKF - NPKF)

    # According to the original paper, `RR AVERAGE2` is the average of the last 8 RR
    # intervals that lie in a certain interval. In the worst case, this requires going back
    # to the very first RR interval. Therefore, all RR intervals are stored. As the
    # algorithm enforces a refractory period, the maximum number of heartbeats is equal to
    # signal_len / refractory_samples.
    RR_intervals = np.zeros(signal_len // REFRACTORY_SAMPLES)

    # tracking the number of peaks found is required to calculate the average RR intervals
    # correctly during the first 8 beats (also needed for access to RR_intervals)
    num_peaks_found = 0

    RR_missed_limit = 0.0

    # in case a searchback was unsuccessful, no new searchback will be performed until
    # another signal peak has been found regularly
    do_searchback = True

    # initialize, so searchback before any peak has been detected works
    peak_index = -REFRACTORY_SAMPLES + 1
    previous_peak_index = -REFRACTORY_SAMPLES + 1

    index = 1
    while index < signal_len - 1:
        PEAKF = 0
        PEAKI = 0

        signal_peak_found = False
        noise_peak_found = False
        # ----------------------------------------------------------------------------------
        # Searchback
        # ----------------------------------------------------------------------------------
        # During a "searchback", detection thresholds are reduced by one half. The peak with
        # highest amplitude between 200ms (i.e. the refractory period) after the previous
        # detected peak and the current index is considered as a peak candidate.
        # Modifications compared to Pan & Tompkins' original method:
        # - The original paper states that a searchback peak's amplitude has to be between
        #   the original threshold and the reduced one. It can happen that this is the case
        #   for the filtered signal, but not for the integrated one (if the raw signal
        #   amplitude is suddenly considerably lower). Therefore, this implementation
        #   requires both signals (filtered and integrated) to be above the reduced
        #   threshold, but only one of them to be below the original threshold.
        # - No further steps are specified for the case that no peak is found during
        #   searchback. Since a searchback is triggered because (physiologically) there has
        #   to be a heartbeat during the searchback interval, this implementation repeats
        #   the process with further reduced thresholds. Up to 16 searchback runs are
        #   performed, each time the thresholds are further reduced by 1/2. A hard limit of
        #   16 runs avoids an endless loop in case there is really only noise.
        # - Since the criterion for triggering a searchback is based on the average RR
        #   interval, in the original form this could only happen after at least two
        #   detected heartbeats. An exceptionally large peak during the first learning phase
        #   can throw the initial thresholds off, so peaks at the beginning are ignored -
        #   which in turn invalidates learning phase 2. Therefore, in addition to the
        #   original searchback criterion (no peak during 1.66 * "the average RR interval"),
        #   a searchback is triggered in two cases: (1) if there is no peak during the first
        #   second, and (2) if there is no peak 1.5s after the first peak.
        if (
            (
                num_peaks_found > 1
                and index - previous_peak_index > RR_missed_limit
                and do_searchback
            )
            or (num_peaks_found == 0 and index > fs)  # original criterion
            or (num_peaks_found == 1 and index - previous_peak_index > 1.5 * fs)  # (1)
        ):  # (2)
            for i in range(1, 16):
                found_a_candidate = False

                searchback_divisor = 1 << i  # 2^i
                best_searchback_index = previous_peak_index + REFRACTORY_SAMPLES
                best_candidate_amplitude = -1
                searchback_index = best_searchback_index

                while searchback_index < index:
                    PEAKF = filtered_ecg[searchback_index]
                    if PEAKF > filtered_ecg[searchback_index + 1]:
                        # next one is lower
                        if PEAKF > filtered_ecg[searchback_index - 1]:
                            # it is a peak
                            PEAKI = integrated_ecg[searchback_index]
                            # one signal is between the reduced and original threshold, the
                            # other one above the reduced threshold
                            if (
                                threshold_F1 / searchback_divisor < PEAKF
                                and PEAKF < threshold_F1
                                and threshold_I1 / searchback_divisor < PEAKI
                            ) or (
                                threshold_I1 / searchback_divisor < PEAKI
                                and PEAKI < threshold_I1
                                and threshold_F1 / searchback_divisor < PEAKF
                            ):
                                if PEAKF > best_candidate_amplitude:
                                    # highest one so far
                                    best_searchback_index = searchback_index
                                    best_candidate_amplitude = filtered_ecg[
                                        searchback_index
                                    ]
                                    found_a_candidate = True

                        # the amplitude of the next sample is lower, so it can't be a peak
                        # -> skip it
                        searchback_index += 1

                    searchback_index += 1

                if found_a_candidate:
                    SPKI = 0.25 * PEAKI + 0.75 * SPKI
                    SPKF = 0.25 * PEAKF + 0.75 * SPKF
                    signal_peak_found = True
                    peak_index = best_searchback_index

                    # don't perform a searchback until the next signal peak has been found
                    # to avoid endless loops
                    do_searchback = False
                    break

        elif filtered_ecg[index] > filtered_ecg[index + 1]:
            if filtered_ecg[index] > filtered_ecg[index - 1]:
                # a local maximum in the filtered signal was found
                PEAKF = filtered_ecg[index]
                PEAKI = integrated_ecg[index]
                if PEAKF > threshold_F1 and PEAKI > threshold_I1:
                    # Both the filtered and the integrated signal are above their respective
                    # thresholds. Thus the current peak is classified as a signal peak and
                    # the running estimates SPKF and SPKI are updated.
                    SPKF = 0.125 * PEAKF + 0.875 * SPKF
                    SPKI = 0.125 * PEAKI + 0.875 * SPKI

                    signal_peak_found = True
                    peak_index = index
                else:
                    noise_peak_found = True

            # The next sample's amplitude is lower, meaning it can't be a peak, so we skip
            # it. This is why there are two separate if-clauses for this block.
            index += 1

        # Calculating the RR interval and comparing slopes only makes sense if there has
        # already been a signal peak in the past.
        if signal_peak_found and num_peaks_found > 0:
            RR = peak_index - previous_peak_index

            # ------------------------------------------------------------------------------
            # T Wave Identification
            # ------------------------------------------------------------------------------
            # "When an RR interval is less than 360 ms (it must be greater than the 200 ms
            # latency), a judgment is made to determine whether the current QRS complex has
            # been correctly identified or whether it is really a T wave. If the maximal
            # slope that occurs during this waveform is less than half that of the QRS
            # waveform that preceded it, it is identified to be a T wave; otherwise, it is
            # called a QRS complex." (from Pan & Tompkins, 1985)
            if RR < T_WAVE_WINDOW:
                reverse_index = peak_index
                max_slope_in_this_peak = -1
                while reverse_index > 0:
                    amplitude_here = filtered_ecg[reverse_index]
                    amplitude_before = filtered_ecg[reverse_index - 1]
                    if amplitude_before > amplitude_here:
                        break
                    slope = amplitude_here - amplitude_before
                    if slope > max_slope_in_this_peak:
                        max_slope_in_this_peak = slope
                    reverse_index -= 1

                reverse_index = previous_peak_index
                max_slope_in_previous_peak = -1
                while reverse_index > 0:
                    amplitude_here = filtered_ecg[reverse_index]
                    amplitude_before = filtered_ecg[reverse_index - 1]
                    if amplitude_before > amplitude_here:
                        break
                    slope = amplitude_here - amplitude_before
                    if slope > max_slope_in_previous_peak:
                        max_slope_in_previous_peak = slope
                    reverse_index -= 1

                if max_slope_in_this_peak < max_slope_in_previous_peak / 2.0:
                    # based on the slope, this peak should be a T wave
                    signal_peak_found = False
                    noise_peak_found = True

        if signal_peak_found:
            # What we know so far: we are at a local maximum, both thresholds are exceeded
            # and it is not a T wave. Thus, the current sample can be considered as a
            # "signal peak" and the adaptive thresholds are updated.
            num_peaks_found += 1
            beat_mask[peak_index] = 1

            threshold_I1 = NPKI + 0.25 * (SPKI - NPKI)
            threshold_F1 = NPKF + 0.25 * (SPKF - NPKF)

            # calculating RR averages only makes sense once 2 peaks have been found
            if num_peaks_found > 1:
                RR_intervals[num_peaks_found] = peak_index - previous_peak_index

                # --------------------------------------------------------------------------
                # Learning phase 2
                # --------------------------------------------------------------------------
                # "Learning phase 2 requires two heartbeats to initialize RR interval
                # average and RR interval limit values." (from Pan & Tompkins, 1985)
                if num_peaks_found == 2:
                    RR_low_limit = 0.92 * RR_intervals[num_peaks_found]
                    RR_high_limit = 1.16 * RR_intervals[num_peaks_found]

                # --------------------------------------------------------------------------
                # RR Average 1 / RR Average 2
                # --------------------------------------------------------------------------
                # RR Average 2 is the average of the 8 most recent RR intervals which fell
                # between RR_low_limit and RR_high_limit. In case of a regular heart rate,
                # this equals RR Average 1 (the average over the 8 most recent RR intervals,
                # independent of any limits). Therefore, RR Average 1 does not need to be
                # calculated separately.
                RR_sum = 0
                RR_count = 0
                irregular = False
                for i in range(num_peaks_found, 1, -1):
                    RR_n = RR_intervals[i]
                    if RR_low_limit < RR_n and RR_n < RR_high_limit:
                        RR_sum += RR_n
                        RR_count += 1
                        if RR_count >= 8:
                            break
                    else:
                        irregular = True

                RR_average = RR_sum / RR_count
                RR_low_limit = 0.92 * RR_average
                RR_high_limit = 1.16 * RR_average
                RR_missed_limit = 1.66 * RR_average

                if irregular:
                    # "For irregular heart rates, the first threshold of each set is reduced
                    # by half so as to increase the detection sensitivity and to avoid
                    # missing beats."
                    threshold_F1 /= 2
                    threshold_I1 /= 2

            # A signal peak has been found, so performing a searchback makes sense.
            do_searchback = True

            # previous peak index is required to calculate the RR interval
            previous_peak_index = peak_index

            # no peak can happen during the refractory period, so skip it
            index = peak_index + REFRACTORY_SAMPLES

        elif noise_peak_found:
            NPKI = 0.125 * PEAKI + 0.875 * NPKI
            NPKF = 0.125 * PEAKF + 0.875 * NPKF
            threshold_I1 = NPKI + 0.25 * (SPKI - NPKI)
            threshold_F1 = NPKF + 0.25 * (SPKF - NPKF)

        index += 1

    # return an array containing `1`s at beat positions and `0`s elsewhere
    return beat_mask


if "numba" in _available_backends:
    _squared_moving_integration_numba = jit(_squared_moving_integration_py, nopython=True)
    _thresholding_numba = jit(_thresholding_py, nopython=True)
