# -*- coding: utf-8 -*-

import numpy as np
import numpy.random
import pytest
import neurokit2 as nk
from neurokit2.signal.signal_fixpeaks import _correct_artifacts, _find_artifacts, signal_fixpeaks


def compute_rmssd(peaks):

    rr = np.ediff1d(peaks, to_begin=0)
    rr[0] = np.mean(rr[1:])
    rmssd = np.sqrt(np.mean(rr ** 2))

    return rmssd


@pytest.fixture
def n_peaks():
    return 1000


@pytest.fixture
def k_peaks():
    return 100


@pytest.fixture
def artifact_idcs(k_peaks, n_peaks):
    idcs = np.arange(k_peaks, n_peaks, k_peaks)
    return idcs


@pytest.fixture
def peaks_correct(n_peaks):
    # Simulate sinusoidally changing heart periods.
    rr = np.sin(np.arange(n_peaks))
    # Add some noise.
    rng = numpy.random.default_rng(42)
    rr_noisy = rng.normal(rr, 0.1)
    # Scale to range of 250msec and offset by 1000msec. I.e., heart period
    # fluctuates in a range of 250msec around 1000msec.
    rr_scaled = 1000 + rr_noisy * 125

    peaks = np.cumsum(np.rint(rr_scaled)).astype(int)

    return peaks


@pytest.fixture
def peaks_misaligned(request, peaks_correct, artifact_idcs):

    rmssd = compute_rmssd(peaks_correct)
    displacement = request.param * rmssd

    peaks_misaligned = peaks_correct.copy()
    peaks_misaligned[artifact_idcs] = peaks_misaligned[artifact_idcs] - displacement

    return peaks_misaligned


@pytest.fixture
def peaks_missed(peaks_correct, artifact_idcs):

    peaks_missed = peaks_correct.copy()
    peaks_missed = np.delete(peaks_missed, artifact_idcs)

    return peaks_missed


@pytest.fixture
def peaks_extra(peaks_correct, artifact_idcs):

    extra_peaks = (peaks_correct[artifact_idcs + 1] - peaks_correct[artifact_idcs]) / 15 + peaks_correct[artifact_idcs]

    peaks_extra = peaks_correct.copy()
    peaks_extra = np.insert(peaks_extra, artifact_idcs, extra_peaks)

    return peaks_extra


@pytest.fixture
def artifacts_misaligned(artifact_idcs):
    artifacts = {"ectopic": list(artifact_idcs + 1), "missed": [], "extra": [], "longshort": list(artifact_idcs)}
    return artifacts


@pytest.fixture
def artifacts_missed(artifact_idcs):
    missed_idcs = [
        j - i for i, j in enumerate(artifact_idcs)
    ]  # account for the fact that peak indices are shifted to the left after deletion of peaks
    artifacts = {"ectopic": [], "missed": missed_idcs, "extra": [], "longshort": []}
    return artifacts


@pytest.fixture
def artifacts_extra(artifact_idcs):
    extra_idcs = [
        j + (i + 1) for i, j in enumerate(artifact_idcs)
    ]  # account for the fact that peak indices are shifted to the right after insertion of peaks
    artifacts = {"ectopic": [], "missed": [], "extra": extra_idcs, "longshort": []}
    return artifacts


@pytest.mark.parametrize("peaks_misaligned", [2, 4, 8], indirect=["peaks_misaligned"])
def test_misaligned_detection(peaks_misaligned, artifacts_misaligned):

    artifacts, _ = _find_artifacts(peaks_misaligned, sampling_rate=1)
    assert artifacts == artifacts_misaligned  # check for identical key-value pairs


def test_missed_detection(peaks_missed, artifacts_missed):

    artifacts, _ = _find_artifacts(peaks_missed, sampling_rate=1)
    assert artifacts == artifacts_missed


def test_extra_detection(peaks_extra, artifacts_extra):

    artifacts, _ = _find_artifacts(peaks_extra, sampling_rate=1)
    assert artifacts == artifacts_extra


@pytest.mark.parametrize("peaks_misaligned", [2, 4, 8], indirect=["peaks_misaligned"])
def test_misaligned_correction(peaks_misaligned, artifacts_misaligned):

    peaks_corrected = _correct_artifacts(artifacts_misaligned, peaks_misaligned)

    assert (
        np.unique(peaks_corrected).size == peaks_misaligned.size
    )  # make sure that no peak duplication occurs and that number of peaks doesn't change


def test_missed_correction(peaks_missed, artifacts_missed):

    peaks_corrected = _correct_artifacts(artifacts_missed, peaks_missed)

    assert np.unique(peaks_corrected).size == (peaks_missed.size + len(artifacts_missed["missed"]))


def test_extra_correction(peaks_extra, artifacts_extra):

    peaks_corrected = _correct_artifacts(artifacts_extra, peaks_extra)

    assert np.unique(peaks_corrected).size == (peaks_extra.size - len(artifacts_extra["extra"]))


def idfn(val):
    if isinstance(val, bool):
        return f"iterative_{val}"


@pytest.mark.parametrize(
    "peaks_misaligned, iterative, rmssd_diff",
    [(2, True, 27), (2, False, 27), (4, True, 113), (4, False, 113), (8, True, 444), (8, False, 444)],
    indirect=["peaks_misaligned"],
    ids=idfn,
)
def test_misaligned_correction_wrapper(peaks_correct, peaks_misaligned, iterative, rmssd_diff):

    _, peaks_corrected = signal_fixpeaks(peaks_misaligned, sampling_rate=1, iterative=iterative)

    rmssd_correct = compute_rmssd(peaks_correct)
    rmssd_corrected = compute_rmssd(peaks_corrected)
    rmssd_uncorrected = compute_rmssd(peaks_misaligned)

    # Assert that correction does not produce peaks that exceed the temporal
    # bounds of the original peaks.
    assert peaks_correct[0] <= peaks_corrected[0]
    assert peaks_correct[-1] >= peaks_corrected[-1]

    # Assert that after artifact correction, the difference in RMSSD to the
    # undistorted signal decreases. This also implicitely tests if the peak
    # distortion affects the RMSSD (manipulation check).
    rmssd_diff_uncorrected = np.abs(rmssd_correct - rmssd_uncorrected)
    rmssd_diff_corrected = np.abs(rmssd_correct - rmssd_corrected)
    assert int(rmssd_diff_uncorrected - rmssd_diff_corrected) == rmssd_diff


@pytest.mark.parametrize("iterative, rmssd_diff", [(True, 3), (False, 3)], ids=idfn)
def test_extra_correction_wrapper(peaks_correct, peaks_extra, iterative, rmssd_diff):

    _, peaks_corrected = signal_fixpeaks(peaks_extra, sampling_rate=1, iterative=iterative)

    rmssd_correct = compute_rmssd(peaks_correct)
    rmssd_corrected = compute_rmssd(peaks_corrected)
    rmssd_uncorrected = compute_rmssd(peaks_extra)

    # Assert that correction does not produce peaks that exceed the temporal
    # bounds of the original peaks.
    assert peaks_correct[0] <= peaks_corrected[0]
    assert peaks_correct[-1] >= peaks_corrected[-1]

    # Assert that after artifact correction, the difference in RMSSD to the
    # undistorted signal decreases. This also implicitely tests if the peak
    # distortion affects the RMSSD (manipulation check).
    rmssd_diff_uncorrected = np.abs(rmssd_correct - rmssd_uncorrected)
    rmssd_diff_corrected = np.abs(rmssd_correct - rmssd_corrected)

    assert int(rmssd_diff_uncorrected - rmssd_diff_corrected) == rmssd_diff


@pytest.mark.parametrize("iterative, rmssd_diff", [(True, 13), (False, 13)], ids=idfn)
def test_missed_correction_wrapper(peaks_correct, peaks_missed, iterative, rmssd_diff):

    _, peaks_corrected = signal_fixpeaks(peaks_missed, sampling_rate=1, iterative=iterative)

    rmssd_correct = compute_rmssd(peaks_correct)
    rmssd_corrected = compute_rmssd(peaks_corrected)
    rmssd_uncorrected = compute_rmssd(peaks_missed)

    # Assert that correction does not produce peaks that exceed the temporal
    # bounds of the original peaks.
    assert peaks_correct[0] <= peaks_corrected[0]
    assert peaks_correct[-1] >= peaks_corrected[-1]

    # Assert that after artifact correction, the difference in RMSSD to the
    # undistorted signal decreases. This also implicitely tests if the peak
    # distortion affects the RMSSD (manipulation check).
    rmssd_diff_uncorrected = np.abs(rmssd_correct - rmssd_uncorrected)
    rmssd_diff_corrected = np.abs(rmssd_correct - rmssd_corrected)

    assert int(rmssd_diff_uncorrected - rmssd_diff_corrected) == rmssd_diff


@pytest.fixture
def testpeaks_for_neurokit_method():
    signal = nk.signal_simulate(duration=20, sampling_rate=1000, frequency=1)
    peaks_true = nk.signal_findpeaks(signal)["Peaks"]
    peaks = np.delete(peaks_true, [5,6,7,8,9,10,15,16,17,19])  # create gaps
    # (I added more than in the example in the function docstring)
    peaks = np.sort(np.append(peaks, [1350, 11350, 18350]))  # add artifacts
    return peaks


@pytest.mark.parametrize("interval_max", [None, 1.5, 2.0])
def test_neurokit_method_returns_only_positive_indices(testpeaks_for_neurokit_method, interval_max):
    peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5,
                                         interval_max=interval_max,
                                         method="neurokit")
    assert np.all(peaks_corrected >= 0)


@pytest.mark.parametrize("interval_max", [None, 1.5, 2.0])
def test_neurokit_method_returns_no_duplicates(testpeaks_for_neurokit_method, interval_max):
    peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5,
                                         interval_max=interval_max,
                                         method="neurokit")
    assert np.unique(peaks_corrected).size == peaks_corrected.size


@pytest.mark.parametrize("interval_max", [None, 1.5, 2.0])
def test_neurokit_method_returns_strictly_increasing_indices(testpeaks_for_neurokit_method, interval_max):
    peaks_corrected = nk.signal_fixpeaks(peaks=testpeaks_for_neurokit_method, interval_min=0.5,
                                         interval_max=interval_max,
                                         method="neurokit")
    assert np.all(np.diff(peaks_corrected) > 0)