"""Functions for processing and interpreting physiological data."""

import matplotlib.pyplot as plt
import numpy as np
from loguru import logger
from scipy import interpolate, signal

from peakdet import editor, utils


@utils.make_operation()
def filter_physio(data, cutoffs, method, *, order=3):
    """
    Apply an `order`-order digital `method` Butterworth filter to `data`.

    Parameters
    ----------
    data : Physio_like
        Input physiological data to be filtered
    cutoffs : int or list
        If `method` is 'lowpass' or 'highpass', an integer specifying the lower
        or upper bound of the filter (in Hz). If method is 'bandpass' or
        'bandstop', a list specifying the lower and upper bound of the filter
        (in Hz).
    method : {'lowpass', 'highpass', 'bandpass', 'bandstop'}
        The type of filter to apply to `data`
    order : int, optional
        Order of filter to be applied. Default: 3

    Returns
    -------
    filtered : :class:`peakdet.Physio`
        Filtered input `data`
    """
    _valid_methods = ["lowpass", "highpass", "bandpass", "bandstop"]

    data = utils.check_physio(data, ensure_fs=True)
    if method not in _valid_methods:
        raise ValueError(
            f"Provided method {method} is not permitted; must be in {_valid_methods}."
        )

    cutoffs = np.array(cutoffs)
    if method in ["lowpass", "highpass"] and cutoffs.size != 1:
        raise ValueError(f"Cutoffs must be length 1 when using {method} filter")
    elif method in ["bandpass", "bandstop"] and cutoffs.size != 2:
        raise ValueError(f"Cutoffs must be length 2 when using {method} filter")

    nyq_cutoff = cutoffs / (data.fs * 0.5)
    if np.any(nyq_cutoff > 1):
        raise ValueError(
            f"Provided cutoffs {cutoffs} are outside of the Nyquist "
            f"frequency for input data with sampling rate {data.fs}."
        )

    if method in ["lowpass", "highpass"]:
        logger.info(
            f"Applying a {method} filter (order: {order}) to the signal, with cutoff "
            f"frequency at {cutoffs} Hz"
        )
    elif method in ["bandpass", "bandstop"]:
        logger.info(
            f"Applying a {method} filter (order: {order}) to the signal, with cutoff "
            f"frequencies at {cutoffs[0]} and {cutoffs[1]} Hz"
        )

    b, a = signal.butter(int(order), nyq_cutoff, btype=method)
    filtered = utils.new_physio_like(data, signal.filtfilt(b, a, data))

    return filtered


@utils.make_operation()
def interpolate_physio(data, target_fs, *, kind="cubic"):
    """
    Interpolate `data` to desired sampling rate `target_fs`.

    Parameters
    ----------
    data : Physio_like
        Input physiological data to be interpolated
    target_fs : float
        Desired sampling rate for `data`
    kind : str or int, optional
        Type of interpolation to perform. Must be one of available kinds in
        :func:`scipy.interpolate.interp1d`. Default: 'cubic'

    Returns
    -------
    interp : :class:`peakdet.Physio`
        Interpolated input `data`
    """
    data = utils.check_physio(data, ensure_fs=True)

    factor = target_fs / data.fs

    # generate original and target "time" series
    t_orig = np.linspace(0, len(data) / data.fs, len(data))
    t_new = np.linspace(0, len(data) / data.fs, int(len(data) * factor))

    # interpolate data and generate new Physio object
    interp = interpolate.interp1d(t_orig, data, kind=kind)(t_new)
    if data.suppdata is None:
        suppinterp = None
    else:
        suppinterp = interpolate.interp1d(t_orig, data.suppdata, kind=kind)(t_new)

    logger.info(
        f"Interpolating the signal at {target_fs} Hz (Interpolation ratio: {factor})."
    )
    interp = utils.new_physio_like(data, interp, fs=target_fs, suppdata=suppinterp)

    return interp


@utils.make_operation()
def peakfind_physio(data, *, thresh=0.2, dist=None):
    """
    Perform peak and trough detection on `data`.

    Parameters
    ----------
    data : Physio_like
        Input data in which to find peaks
    thresh : float [0,1], optional
        Relative height threshold a data point must surpass to be classified as
        a peak. Default: 0.2
    dist : int, optional
        Distance in indices that peaks must be separated by in `data`. If None,
        this is estimated. Default: None

    Returns
    -------
    peaks : :class:`peakdet.Physio`
        Input `data` with detected peaks and troughs
    """
    ensure_fs = True if dist is None else False
    data = utils.check_physio(data, ensure_fs=ensure_fs, copy=True)
    # first pass peak detection to get approximate distance between peaks
    cdist = data.fs // 4 if dist is None else dist

    # check if data is negative, if so make it all positive and continue with signal
    phys_signal = data.data - data.data.min() if data.data.min() < 0 else data.data
    logger.debug(
        f"Negative signal detected (min = {data.data.min()}), workgin with positive "
        "signal for peak detection."
    )

    thresh = np.squeeze(np.diff(np.percentile(phys_signal, [5, 95]))) * thresh
    locs, heights = signal.find_peaks(phys_signal, distance=cdist, height=thresh)
    logger.debug(
        f"First peak detection iteration. Acquiring approximate distance between peaks "
        f"(Number of peaks: {len(locs)})"
    )

    # second, more thorough peak detection
    cdist = np.diff(locs).mean() // 2
    heights = np.percentile(heights["peak_heights"], 1)
    locs, heights = signal.find_peaks(phys_signal, distance=cdist, height=heights)
    data._metadata["peaks"] = locs
    logger.debug(
        f"Second peak detection iteration. Acquiring more precise peak locations "
        f"(Number of peaks: {len(locs)})"
    )
    # perform trough detection based on detected peaks
    data._metadata["troughs"] = utils.check_troughs(data, data.peaks)
    logger.debug(
        f"Trough detection based on detected peaks (Number of troughs: "
        f"{len(data.troughs)})"
    )

    return data


@utils.make_operation()
def delete_peaks(data, remove):
    """
    Delete peaks in `remove` from peaks stored in `data`.

    Parameters
    ----------
    data : Physio_like
    remove : array_like

    Returns
    -------
    data : Physio_like
    """
    data = utils.check_physio(data, ensure_fs=False, copy=True)
    data._metadata["peaks"] = np.setdiff1d(data._metadata["peaks"], remove)
    data._metadata["troughs"] = utils.check_troughs(data, data.peaks, data.troughs)

    return data


@utils.make_operation()
def reject_peaks(data, remove):
    """
    Mark peaks in `remove` as rejected artifacts in `data`.

    Parameters
    ----------
    data : Physio_like
    remove : array_like

    Returns
    -------
    data : Physio_like
    """
    data = utils.check_physio(data, ensure_fs=False, copy=True)
    data._metadata["reject"] = np.append(data._metadata["reject"], remove)
    data._metadata["troughs"] = utils.check_troughs(data, data.peaks, data.troughs)

    return data


@utils.make_operation()
def add_peaks(data, add):
    """
    Add `newpeak` to add them in `data`.

    Parameters
    ----------
    data : Physio_like
    add : list of int

    Returns
    -------
    data : Physio_like
    """
    data = utils.check_physio(data, ensure_fs=False, copy=True)
    peaks = data._metadata["peaks"]

    if np.isscalar(add):
        add = [add]

    add = np.asarray(add, dtype=int)

    for val in add:
        if val in peaks or np.any(np.abs(peaks - val) == 1):
            logger.warning(
                f"Added peak {val} is too close to an existing peak. Skipping."
            )
        else:
            idx = np.searchsorted(data._metadata["peaks"], val)
            data._metadata["peaks"] = np.insert(data._metadata["peaks"], idx, val)
            data._metadata["troughs"] = utils.check_troughs(data, data.peaks)

    return data


def edit_physio(data):
    """
    Open interactive plot with `data` to permit manual editing of time series.

    Parameters
    ----------
    data : Physio_like
        Physiological data to be edited

    Returns
    -------
    edited : :class:`peakdet.Physio`
        Input `data` with manual edits
    """
    data = utils.check_physio(data, ensure_fs=True)

    # no point in manual edits if peaks/troughs aren't defined
    if not (len(data.peaks) and len(data.troughs)):
        return

    # perform manual editing
    logger.info("Opening interactive peak editor")
    edits = editor._PhysioEditor(data)
    plt.show(block=True)

    # replay editing on original provided data object
    if len(edits.rejected) > 0:
        data = reject_peaks(data, remove=sorted(edits.rejected))
    if len(edits.deleted) > 0:
        data = delete_peaks(data, remove=sorted(edits.deleted))
    if len(edits.included) > 0:
        data = add_peaks(data, add=sorted(edits.included))

    return data


def plot_physio(data, *, ax=None):
    """
    Plot `data` and associated peaks / troughs.

    Parameters
    ----------
    data : Physio_like
        Physiological data to plot
    ax : :class:`matplotlib.axes.Axes`, optional
        Axis on which to plot `data`. If None, a new axis is created. Default:
        None

    Returns
    -------
    ax : :class:`matplotlib.axes.Axes`
        Axis with plotted `data`
    """
    logger.debug(f"Plotting {data}")
    # generate x-axis time series
    fs = 1 if np.isnan(data.fs) else data.fs
    time = np.arange(0, len(data) / fs, 1 / fs)
    if ax is None:
        fig, ax = plt.subplots(1, 1)
    # plot data with peaks + troughs, as appropriate
    ax.plot(
        time,
        data,
        "b",
        time[data.peaks],
        data[data.peaks],
        ".r",
        time[data.troughs],
        data[data.troughs],
        ".g",
    )

    return ax
