"""
WARNING
-------
This module is deprecated, please use `torch_ecg._preprocessors` instead.


preprocess of (single lead) ECG signal:
    band pass (dep. on purpose?) --> remove baseline (?) --> find rpeaks
    --> wave delineation (?, put into several stand alone files)

NOTE:
-----
1. ALL signals are assumed to have units in mV

References:
-----------
[1] https://github.com/PIA-Group/BioSPPy

"""

import multiprocessing as mp
from collections import Counter
from numbers import Real
from typing import Dict, List, Optional

import numpy as np

# from scipy.signal import medfilt
# https://github.com/scipy/scipy/issues/9680
from biosppy.signals.tools import filter_signal
from scipy.ndimage.filters import median_filter

from ..cfg import CFG
from .misc import list_sum, ms2samples
from .rpeaks import (
    christov_detect,
    engzee_detect,
    gamboa_detect,
    gqrs_detect,
    hamilton_detect,
    ssf_detect,
    xqrs_detect,
)


__all__ = [
    "preprocess_multi_lead_signal",
    "preprocess_single_lead_signal",
    "rpeaks_detect_multi_leads",
    "merge_rpeaks",
]


QRS_DETECTORS = {
    "xqrs": xqrs_detect,
    "gqrs": gqrs_detect,
    "hamilton": hamilton_detect,
    "ssf": ssf_detect,
    "christov": christov_detect,
    "engzee": engzee_detect,
    "gamboa": gamboa_detect,
}

# ECG signal preprocessing configurations
PreprocCfg = CFG()
# PreprocCfg.fs = 500
PreprocCfg.rpeak_mask_radius = 50  # ms
PreprocCfg.rpeak_lead_num_thr = (
    8 / 12
)  # ratio of leads, used for merging rpeaks detected from multiple leads
PreprocCfg.beat_winL = 250
PreprocCfg.beat_winR = 250


def preprocess_multi_lead_signal(
    raw_sig: np.ndarray,
    fs: Real,
    sig_fmt: str = "channel_first",
    bl_win: Optional[List[Real]] = None,
    band_fs: Optional[List[Real]] = None,
    rpeak_fn: Optional[str] = None,
    verbose: int = 0,
) -> Dict[str, np.ndarray]:
    """
    perform preprocessing for multi-lead ECG signal (with units in mV),
    preprocessing may include median filter, bandpass filter, and rpeaks detection, etc.

    Parameters
    ----------
    raw_sig: ndarray,
        the raw ECG signal, with units in mV
    fs: numbers.Real,
        sampling frequency of `raw_sig`
    sig_fmt: str, default "channel_first",
        format of the multi-lead ECG signal,
        "channel_last" (alias "lead_last"), or
        "channel_first" (alias "lead_first", original)
    bl_win: list (of 2 numbers.Real), optional,
        window (units in second) of baseline removal using `median_filter`,
        the first is the shorter one, the second the longer one,
        a typical pair is [0.2, 0.6],
        if is None or empty, baseline removal will not be performed
    band_fs: list (of 2 numbers.Real), optional,
        frequency band of the bandpass filter,
        a typical pair is [0.5, 45],
        be careful when detecting paced rhythm,
        if is None or empty, bandpass filtering will not be performed
    rpeak_fn: str, optional,
        name of the function detecting rpeaks,
        can be one of keys of `QRS_DETECTORS`, case insensitive
    verbose: int, default 0,
        print verbosity

    Returns
    -------
    retval: dict,
        with items
        - "filtered_ecg": the array of the processed ECG signal
        - "rpeaks": the array of indices of rpeaks; empty if `rpeak_fn` is not given

    """
    assert sig_fmt.lower() in [
        "channel_first",
        "lead_first",
        "channel_last",
        "lead_last",
    ]
    if sig_fmt.lower() in ["channel_last", "lead_last"]:
        filtered_ecg = raw_sig.T
    else:
        filtered_ecg = raw_sig.copy()
    rpeaks_candidates = []
    cpu_num = max(1, mp.cpu_count() - 3)
    with mp.Pool(processes=cpu_num) as pool:
        results = pool.starmap(
            func=preprocess_single_lead_signal,
            iterable=[
                (filtered_ecg[lead, ...], fs, bl_win, band_fs, rpeak_fn)
                for lead in range(filtered_ecg.shape[0])
            ],
        )
    for lead in range(filtered_ecg.shape[0]):
        filtered_metadata = results[lead]
        filtered_ecg[lead, ...] = filtered_metadata["filtered_ecg"]
        rpeaks_candidates.append(filtered_metadata["rpeaks"])
        if verbose >= 1:
            print(
                f"for the {lead}-th lead, rpeaks_candidates = {filtered_metadata['rpeaks']}"
            )

    if sig_fmt.lower() in ["channel_last", "lead_last"]:
        rpeaks = merge_rpeaks(rpeaks_candidates, raw_sig.T, fs, verbose)
    else:
        rpeaks = merge_rpeaks(rpeaks_candidates, raw_sig, fs, verbose)
    retval = CFG(
        {
            "filtered_ecg": filtered_ecg,
            "rpeaks": rpeaks,
        }
    )
    return retval


def preprocess_single_lead_signal(
    raw_sig: np.ndarray,
    fs: Real,
    bl_win: Optional[List[Real]] = None,
    band_fs: Optional[List[Real]] = None,
    rpeak_fn: Optional[str] = None,
    verbose: int = 0,
) -> Dict[str, np.ndarray]:
    """
    perform preprocessing for single lead ECG signal (with units in mV),
    preprocessing may include median filter, bandpass filter, and rpeaks detection, etc.

    Parameters
    ----------
    raw_sig: ndarray,
        the raw ECG signal, with units in mV
    fs: numbers.Real,
        sampling frequency of `raw_sig`
    bl_win: list (of 2 numbers.Real), optional,
        window (units in second) of baseline removal using `median_filter`,
        the first is the shorter one, the second the longer one,
        a typical pair is [0.2, 0.6],
        if is None or empty, baseline removal will not be performed
    band_fs: list (of 2 numbers.Real), optional,
        frequency band of the bandpass filter,
        a typical pair is [0.5, 45],
        be careful when detecting paced rhythm,
        if is None or empty, bandpass filtering will not be performed
    rpeak_fn: str, optional,
        name of the function detecting rpeaks,
        can be one of keys of `QRS_DETECTORS`, case insensitive
    verbose: int, default 0,
        print verbosity

    Returns
    -------
    retval: dict,
        with items
        - "filtered_ecg": the array of the processed ECG signal
        - "rpeaks": the array of indices of rpeaks; empty if `rpeak_fn` is not given

    """
    filtered_ecg = raw_sig.copy()

    # remove baseline
    if bl_win:
        window1 = 2 * (int(bl_win[0] * fs) // 2) + 1  # window size must be odd
        window2 = 2 * (int(bl_win[1] * fs) // 2) + 1
        baseline = median_filter(filtered_ecg, size=window1, mode="nearest")
        baseline = median_filter(baseline, size=window2, mode="nearest")
        filtered_ecg = filtered_ecg - baseline

    # filter signal
    if band_fs:
        filtered_ecg = filter_signal(
            signal=filtered_ecg,
            ftype="FIR",
            # ftype="butter",
            band="bandpass",
            order=int(0.3 * fs),
            sampling_rate=fs,
            frequency=band_fs,
        )["signal"]

    if rpeak_fn:
        rpeaks = QRS_DETECTORS[rpeak_fn.lower()](filtered_ecg, fs).astype(int)
    else:
        rpeaks = np.array([], dtype=int)

    retval = CFG(
        {
            "filtered_ecg": filtered_ecg,
            "rpeaks": rpeaks,
        }
    )

    return retval


def rpeaks_detect_multi_leads(
    sig: np.ndarray,
    fs: Real,
    sig_fmt: str = "channel_first",
    rpeak_fn: str = "xqrs",
    verbose: int = 0,
) -> np.ndarray:
    """
    detect rpeaks from the filtered multi-lead ECG signal (with units in mV)

    Parameters
    ----------
    sig: ndarray,
        the (better be filtered) ECG signal, with units in mV
    fs: numbers.Real,
        sampling frequency of `sig`
    sig_fmt: str, default "channel_first",
        format of the multi-lead ECG signal,
        "channel_last" (alias "lead_last"), or
        "channel_first" (alias "lead_first", original)
    rpeak_fn: str,
        name of the function detecting rpeaks,
        can be one of keys of `QRS_DETECTORS`, case insensitive
    verbose: int, default 0,
        print verbosity

    Returns
    -------
    rpeaks: np.ndarray,
        array of indices of the detected rpeaks of the multi-lead ECG signal

    """
    assert sig_fmt.lower() in [
        "channel_first",
        "lead_first",
        "channel_last",
        "lead_last",
    ]
    if sig_fmt.lower() in ["channel_last", "lead_last"]:
        s = sig.T
    else:
        s = sig.copy()
    rpeaks = []
    for lead in range(s.shape[0]):
        rpeaks.append((QRS_DETECTORS[rpeak_fn.lower()](s[lead], fs)).astype(int))
    rpeaks = merge_rpeaks(rpeaks, sig, fs, verbose)
    return rpeaks


def merge_rpeaks(
    rpeaks_candidates: List[np.ndarray], sig: np.ndarray, fs: Real, verbose: int = 0
) -> np.ndarray:
    """
    merge rpeaks that are detected from each lead of multi-lead signals (with units in mV),
    using certain criterion merging qrs masks from each lead

    Parameters
    ----------
    rpeaks_candidates: list of ndarray,
        each element (ndarray) is the array of indices of rpeaks of corr. lead
    sig: ndarray,
        the multi-lead ECG signal, with units in mV
    fs: numbers.Real,
        sampling frequency of `sig`
    verbose: int, default 0,
        print verbosity

    Returns
    -------
    final_rpeaks: np.ndarray
        the final rpeaks obtained by merging the rpeaks from all the leads

    """
    rpeak_masks = np.zeros_like(sig, dtype=int)
    sig_len = sig.shape[1]
    radius = ms2samples(PreprocCfg.rpeak_mask_radius, fs)
    if verbose >= 1:
        print(f"sig_len = {sig_len}, radius = {radius}")
    for lead in range(sig.shape[0]):
        for r in rpeaks_candidates[lead]:
            rpeak_masks[lead, max(0, r - radius) : min(sig_len - 1, r + radius)] = 1
    rpeak_masks = (
        rpeak_masks.sum(axis=0) >= int(PreprocCfg.rpeak_lead_num_thr * sig.shape[0])
    ).astype(int)
    rpeak_masks[0], rpeak_masks[-1] = 0, 0
    split_indices = np.where(np.diff(rpeak_masks) != 0)[0]
    if verbose >= 1:
        print(
            f"split_indices = {split_indices}, with total number = {len(split_indices)}"
        )
        if verbose >= 2:
            print(
                f"the corresponding intervals are {[[split_indices[2*idx], split_indices[2*idx+1]] for idx in range(len(split_indices)//2)]}"
            )

    final_rpeaks = []
    for idx in range(len(split_indices) // 2):
        start_idx = split_indices[2 * idx]
        end_idx = split_indices[2 * idx + 1]
        rc = list_sum(  # `lr`: list of rpeaks
            [
                lr[np.where((lr >= start_idx) & (lr <= end_idx))].tolist()
                for lr in rpeaks_candidates
            ]
        )
        if verbose >= 2:
            print(
                f"at the {idx}-th interval, start_idx = {start_idx}, end_idx = {end_idx}"
            )
            print(f"rpeak candidates = {rc}")
        counter = Counter(rc).most_common()
        if len(counter) > 0 and counter[0][1] >= len(rc) // 2 + 1:
            # might have the case where
            # rc is empty
            final_rpeaks.append(counter[0][0])
        elif len(counter) > 0:
            final_rpeaks.append(int(np.mean(rc)))
    final_rpeaks = np.array(final_rpeaks)
    return final_rpeaks
