# -*- coding: utf-8 -*-
import os.path

import numpy as np
import pandas as pd

# Trick to directly access internal functions for unit testing.
#
# Using neurokit2.ecg.ecg_findpeaks._ecg_findpeaks_MWA doesn't
# work because of the "from .ecg_findpeaks import ecg_findpeaks"
# statement in neurokit2/ecg/__init.__.py.
from neurokit2.ecg.ecg_findpeaks import _ecg_findpeaks_MWA, _ecg_findpeaks_peakdetect


def _read_csv_column(csv_name, column):
    csv_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "ecg_data", csv_name)
    csv_data = pd.read_csv(csv_path, header=None)
    return csv_data[column].to_numpy()


def test_ecg_findpeaks_MWA():
    np.testing.assert_array_equal(
        _ecg_findpeaks_MWA(np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=float), 3), [0, 0.5, 1, 2, 3, 4, 5, 6, 7, 8]
    )


# This test case is intentionally a "change aversion" test that simply
# verifies that the output of the _ecg_findpeaks_peakdetect function
# on two different test datasets remains unchanged.
#
# Most notably the assertions here don't necessarily document the
# "correct" output of the function, just what the output used to be earlier.
# Potential bug fixes could legitimately require updates to this test case.
#
# Instead the main purpose of this test case is to give extra confidence
# that optimizations or other refactorings won't accidentally introduce
# new bugs into the function.
def test_ecg_findpeaks_peakdetect():
    good_4000 = _read_csv_column("good_4000.csv", 1)
    expected_good_4000_peaks = _read_csv_column("expected_ecg_findpeaks_peakdetect_good_4000.csv", 0)
    np.testing.assert_array_equal(_ecg_findpeaks_peakdetect(good_4000, sampling_rate=4000), expected_good_4000_peaks)

    bad_500 = _read_csv_column("bad_500.csv", 1)
    expected_bad_500_peaks = _read_csv_column("expected_ecg_findpeaks_peakdetect_bad_500.csv", 0)
    np.testing.assert_array_equal(_ecg_findpeaks_peakdetect(bad_500, sampling_rate=500), expected_bad_500_peaks)
