"""Module for processing Respiration data."""

from collections.abc import Sequence
from typing import Optional, Union

import neurokit2 as nk
import numpy as np
import pandas as pd
import scipy.signal as ss

from biopsykit.signals._base import _BaseProcessor
from biopsykit.utils.array_handling import sanitize_input_1d

__all__ = ["RspProcessor"]


class RspProcessor(_BaseProcessor):
    """Class for processing Respiration data."""

    def __init__(
        self,
        data: Union[pd.DataFrame, dict[str, pd.DataFrame]],
        sampling_rate: Optional[float] = None,
        time_intervals: Optional[Union[pd.Series, dict[str, Sequence[str]]]] = None,
        include_start: Optional[bool] = False,
    ):
        """Initialize an ``RspProcessor`` instance.

        You can either pass a data dictionary 'data_dict' containing Respiration data or dataframe containing
        Respiration data. For the latter, you can additionally supply time information via ``time_intervals`` parameter
        to automatically split the data into single phases.


        Parameters
        ----------
        data : :class:`~pandas.DataFrame` or dict
            dataframe (or dict of such) with Respiration data
        sampling_rate : float, optional
            sampling rate of recorded data
        time_intervals : dict or :class:`~pandas.Series`, optional
            time intervals indicating how ``data`` should be split.
            Can either be a :class:`~pandas.Series` with the `start` times of the single phases
            (the phase names are then derived from the index) or a dictionary with tuples indicating
            `start` and `end` times of phases (the phase names are then derived from the dict keys).
            Default: ``None`` (data is not split further)
        include_start : bool, optional
            ``True`` to include the data from the beginning of the recording to the first time interval as the
            first phase (then named ``Start``), ``False`` otherwise. Default: ``False``

        """
        super().__init__(
            data=data, sampling_rate=sampling_rate, time_intervals=time_intervals, include_start=include_start
        )

    @classmethod
    def rsp_compute_rate(cls, rsp_signal: pd.DataFrame, sampling_rate: Optional[float] = 256.0) -> float:
        """Compute respiration rate for given interval from respiration signal.

        Parameters
        ----------
        rsp_signal : :class:`~pandas.DataFrame`
            Raw respiration signal (1D). Can be a 'true' respiration signal (e.g. from bioimpedance or Radar)
            or an 'estimated' respiration signal (e.g. from ECG-derived respiration)
        sampling_rate : float, optional
            Sampling rate of recorded data


        Returns
        -------
        float
            Respiration rate during the given interval in bpm (breaths per minute)


        References
        ----------
        Schäfer, A., & Kratky, K. W. (2008). Estimation of Breathing Rate from Respiratory Sinus Arrhythmia:
        Comparison of Various Methods. *Annals of Biomedical Engineering*, 36(3), 476-485.
        https://doi.org/10.1007/s10439-007-9428-1


        Examples
        --------
        >>> from biopsykit.signals.ecg import EcgProcessor
        >>> # initialize EcgProcessor instance
        >>> ecg_processor = EcgProcessor(...)

        >>> # Extract respiration signal estimated from ECG using the 'peak_trough_diff' method
        >>> rsp_signal = ecg_processor.ecg_estimate_rsp(ecg_processor, key="Data", edr_type='peak_trough_diff')
        >>> # Compute respiration rate from respiration signal
        >>> rsp_rate = ecg_processor.rsp_compute_rate(rsp_signal)

        """
        # find peaks: minimal distance between peaks: 1 seconds
        rsp_signal = sanitize_input_1d(rsp_signal)
        edr_maxima = ss.find_peaks(rsp_signal, height=0, distance=sampling_rate)[0]
        edr_minima = ss.find_peaks(-1 * rsp_signal, height=0, distance=sampling_rate)[0]
        # threshold: 0.2 * Q3 (= 75th percentile)
        max_threshold = 0.2 * np.percentile(rsp_signal[edr_maxima], 75)
        # find all maxima that are above the threshold
        edr_maxima = edr_maxima[rsp_signal[edr_maxima] > max_threshold]

        # rearrange maxima into 2D array to that each row contains the start and end of one cycle
        rsp_cycles_start_end = np.vstack([edr_maxima[:-1], edr_maxima[1:]]).T

        # check for
        valid_resp_phases_mask = np.apply_along_axis(cls._check_contains_trough, 1, rsp_cycles_start_end, edr_minima)
        rsp_cycles_start_end = rsp_cycles_start_end[valid_resp_phases_mask]

        edr_signal_split: list = np.split(rsp_signal, rsp_cycles_start_end.flatten())[1::2]
        rsp_cycles_start_end = rsp_cycles_start_end.T
        rsp_peaks = rsp_cycles_start_end[0]
        rsp_troughs = np.array(list(map(np.argmin, edr_signal_split))) + rsp_cycles_start_end[0]

        rsp_rate_peaks = cls._rsp_rate(rsp_peaks, int(sampling_rate), len(rsp_signal))
        rsp_rate_troughs = cls._rsp_rate(rsp_troughs, int(sampling_rate), len(rsp_signal))
        return np.concatenate([rsp_rate_peaks, rsp_rate_troughs]).mean()

    @classmethod
    def _check_contains_trough(cls, start_end: np.ndarray, minima: np.ndarray) -> bool:
        """Check whether exactly one minima is in the interval given by the array ``start_end``: [start, end].

        Parameters
        ----------
        start_end : :class:`numpy.array`
            Array with start and end index
        minima : :class:`numpy.array`
            Array containing minima to be checked


        Returns
        -------
        bool
            ``True`` if exactly one minima is in the ``[start, end]`` interval, ``False`` otherwise

        """
        start, end = start_end
        return minima[(minima > start) & (minima < end)].shape[0] == 1

    @classmethod
    def _rsp_rate(cls, extrema: np.array, sampling_rate: int, desired_length: int) -> np.array:
        """Compute continuous respiration rate from extrema values.

        Parameters
        ----------
        extrema: :class:`numpy.array`
            List of respiration extrema (peaks or troughs)
        sampling_rate : float
            Sampling rate of recorded data
        desired_length : int
            Desired length of the output signal


        Returns
        -------
        :class:`numpy.array`
            Respiration rate array interpolated to desired length

        """
        rsp_rate_raw = (sampling_rate * 60) / np.ediff1d(extrema)
        # remove last sample
        x_old = extrema[:-1]
        x_new = np.linspace(x_old[0], x_old[-1], desired_length)
        return nk.signal_interpolate(x_old, rsp_rate_raw, x_new, method="linear")
