"""
Signal-processing utilities: smoothing, peak/min finding, baseline,
baseline subtraction, sign inversion.
"""

import numpy as np
from scipy.signal import savgol_filter, find_peaks
from scipy.interpolate import interp1d, pchip_interpolate

try:
    from scipy.interpolate import make_smoothing_spline
except ImportError:
    make_smoothing_spline = None

# from skimage.restoration import rolling_ball


def smooth(y, win_size=None):
    """
    Savitzky–Golay smoothing. If win_size is None, choose
    ~5% of len + 0.5*(range).
    """
    cr = y.max() - y.min()
    if win_size is None:
        w = int(0.05 * len(y) + 0.5 * cr)
    else:
        w = win_size
    p = 3
    if w <= p:
        w = p + 1
    if w % 2 == 0:
        w += 1
    return savgol_filter(y, w, p)


def find_minima_scipy(y, distance, prominence):
    """
    Find minima by running find_peaks on -y, then refining
    with rolling_ball background subtraction.
    """
    p, _ = find_peaks(-y, distance=distance, prominence=prominence)
    if len(p) > 1:
        # TODO: check, was -(y - bg), but seems wrong
        # radius = p[0]
        # scale = radius / y.max()
        # bg = rolling_ball(y * scale, radius=radius) / scale
        p2, _ = find_peaks(-(y), distance=int((p[1] - p[0]) * 0.75))
        return p2
    return p


def baseline(y, q, minima_idx, method="pchip"):
    """
    Estimate baseline through minima points.
    method: "linear", "spline", or "pchip".
    """
    xs = np.insert(q[minima_idx], 0, 0)
    ys = np.insert(y[minima_idx], 0, 0)
    if method == "linear":
        f = interp1d(
            xs, ys, kind="linear", bounds_error=False, fill_value="extrapolate"
        )
        b = f(q)
    elif method == "spline" and make_smoothing_spline and len(xs) >= 5:
        b = make_smoothing_spline(xs, ys)(q)
    else:
        b = pchip_interpolate(xs, ys, q)
    b = np.clip(b, 0, y)
    return b


def subtract_baseline(y, b):
    """
    Return |y - b|.
    """
    return np.abs(y - b)


def invert_signs(y, minima_idx):
    """
    Flip the sign of y after each minima index.
    """
    y2 = y.copy()
    for idx in minima_idx:
        y2[idx:] *= -1
    return y2
