import numpy as np

from breesy import breesy_scipy, constants
from breesy.recording import Recording
from breesy.errors import BreesyError, BreesyInternalError
from breesy.recording import update_data
from breesy.type_hints import enforce_type_hints


BANDPASS_PARAMS = """
:param recording: Input recording
:param low: Lower frequency bound in Hz
:param high: Upper frequency bound in Hz
"""

SINGLE_FREQ_PARAMS = """
:param recording: Input recording
:param frequency: Frequency threshold in Hz
"""

NOTCH_PARAMS = """
:param recording: Input recording
:param frequency: Frequency to remove in Hz
"""

DEFAULT_WINDOW_PARAMS = constants.DEFAULT_WINDOW

_DOC_PARAMS = {
    name: value for name, value in globals().items()
    if name.endswith('_PARAMS') and isinstance(value, str)
}


def _format_doc(func):
    """Decorator that expands shared parameter docs in the docstring."""
    if func.__doc__:
        func.__doc__ = func.__doc__.format(**_DOC_PARAMS)
    return func


# Default filters


@_format_doc
@enforce_type_hints
def notch_filter(recording: Recording, frequency: int | float, smoothness_factor: int | float = constants.NOTCH_DEFAULT_SMOOTHNESS) -> Recording:
    """Apply notch filter to remove a specific frequency from a recording.
    By default, it is a FIR filter.

    {NOTCH_PARAMS}
    :param smoothness_factor: Controls the smoothness of attenuation. Higher value means wider affected (attenuated) band.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    transition_width, notch_width = _get_fir_params_from_smoothness(smoothness_factor=smoothness_factor, notch_freq=frequency)

    new_data = recording.data.copy()
    new_data = _notch_fir(data=new_data, frequency=frequency, sample_rate=recording.sample_rate, transition_width=transition_width, notch_width=notch_width)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def bandpass_filter(recording: Recording, low: float | int, high: float | int,
                    smoothness_factor: int | float = constants.BANDPASS_DEFAULT_SMOOTHNESS) -> Recording:
    """Filter recording to keep only a specific frequency band by applying a bandpass filter.
    By default, it is a FIR filter.

    {BANDPASS_PARAMS}
    :param smoothness_factor: Controls the smoothness of attenuation. Higher value means wider transition region (from not attenuated to attenuated) from each side of the band.

    :return: Filtered recording
    """
    _assert_bandpass_low_high(sample_rate=recording.sample_rate, low=low, high=high)

    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=low, high=high, sample_rate=recording.sample_rate,
                             transition_width=smoothness_factor)  # TODO: check if needs conversion
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def lowpass_filter(recording: Recording, frequency: float | int,
                   smoothness_factor: int | float = constants.BANDPASS_DEFAULT_SMOOTHNESS) -> Recording:
    """Filter recording to keep only frequencies below certain threshold by applying a lowpass filter.
    By default, it is a FIR filter.

    {SINGLE_FREQ_PARAMS}
    :param smoothness_factor: Controls the smoothness of attenuation. Higher value means wider transition region (from not attenuated to attenuated) from each side of the band.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)

    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=None, high=frequency, sample_rate=recording.sample_rate,
                             transition_width=smoothness_factor)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def highpass_filter(recording: Recording, frequency: float | int,
                   smoothness_factor: int | float = constants.BANDPASS_DEFAULT_SMOOTHNESS) -> Recording:
    """Filter recording to keep only frequencies above certain threshold by applying a highpass filter.
    By default, it is a FIR filter.

    {SINGLE_FREQ_PARAMS}
    :param smoothness_factor: Controls the smoothness of attenuation. Higher value means wider transition region (from not attenuated to attenuated) from each side of the band.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)

    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=frequency, high=None, sample_rate=recording.sample_rate,
                             transition_width=smoothness_factor)
    return update_data(recording, new_data) 


# Specific filters


@_format_doc
@enforce_type_hints
def notch_filter_fir(recording: Recording, frequency: int | float, notch_width: int | float, transition_width: int | float,
                     window: str | None = None) -> Recording:
    """Apply FIR notch filter to remove a specific frequency from a recording.

    {NOTCH_PARAMS}
    :param notch_width: Controls the width of removed (strongly attenuated) band
    :param transition_width: Controls the width of transition region (from not attenuated to attenuated) from each side of removed band
    :param window: Name of window type to use (default: {DEFAULT_WINDOW_PARAMS})

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _notch_fir(data=new_data, frequency=frequency, sample_rate=recording.sample_rate, 
                          transition_width=transition_width, notch_width=notch_width, window=window)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def notch_filter_iir(recording: Recording, frequency: int | float, quality_factor: int) -> Recording:
    """Apply IIR notch filter to remove a specific frequency from a recording.

    {NOTCH_PARAMS}
    :param quality_factor: Controls the strength and width of attenuation. Higher value means narrower and deeper removed band.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _notch_iir(data=new_data, frequency=frequency, sample_rate=recording.sample_rate, Q=quality_factor)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def bandpass_filter_fir(recording: Recording, low: float | int, high: float | int, transition_width: int | float,
                        window: str | None = None) -> Recording:
    """Apply FIR bandpass filter to keep only a specific frequency band.

    {BANDPASS_PARAMS}
    :param transition_width: Controls the width of transition region (from not attenuated to attenuated) from each side of the band
    :param window: Name of window type to use (default: {DEFAULT_WINDOW_PARAMS})
    
    :return: Filtered recording
    """
    _assert_bandpass_low_high(sample_rate=recording.sample_rate, low=low, high=high)
    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=low, high=high, sample_rate=recording.sample_rate,
                             transition_width=transition_width, window=window)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def lowpass_filter_fir(recording: Recording, frequency: float | int, transition_width: int | float,
                       window: str | None = None) -> Recording:
    """Apply FIR bandpass filter to keep only frequencies below the threshold.

    {SINGLE_FREQ_PARAMS}
    :param transition_width: Controls the width of transition region (from not attenuated to attenuated)
    :param window: Name of window type to use (default: {DEFAULT_WINDOW_PARAMS})

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=None, high=frequency, sample_rate=recording.sample_rate, transition_width=transition_width)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def highpass_filter_fir(recording: Recording, frequency: float | int, transition_width: int | float,
                        window: str | None = None) -> Recording:
    """Apply FIR bandpass filter to keep only frequencies above the threshold.

    {SINGLE_FREQ_PARAMS}
    :param transition_width: Controls the width of transition region (from attenuated to not attenuated)
    :param window: Name of window type to use (default: {DEFAULT_WINDOW_PARAMS})

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _bandpass_fir(data=new_data, low=frequency, high=None, sample_rate=recording.sample_rate, transition_width=transition_width)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def bandpass_filter_iir(recording: Recording, low: float | int, high: float | int, filter_order: int) -> Recording:
    """Apply IIR bandpass filteer to keep only a specific frequency band.

    {BANDPASS_PARAMS}
    :param filter_order: Controls the strength of filtering. Higher values mean narrower transition regions from each side of the band.

    :return: Filtered recording
    """
    _assert_bandpass_low_high(sample_rate=recording.sample_rate, low=low, high=high)
    new_data = recording.data.copy()
    new_data = _bandpass_iir(data=new_data, low=low, high=high, sample_rate=recording.sample_rate, filter_order=filter_order)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def lowpass_filter_iir(recording: Recording, frequency: float | int, filter_order: int) -> Recording:
    """Apply IIR bandpass filteer to keep only frequencies below the threshold.

    {SINGLE_FREQ_PARAMS}
    :param filter_order: Controls the strength of filtering. Higher values mean narrower transition region.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _bandpass_iir(data=new_data, low=None, high=frequency, sample_rate=recording.sample_rate, filter_order=filter_order)
    return update_data(recording, new_data) 


@_format_doc
@enforce_type_hints
def highpass_filter_iir(recording: Recording, frequency: float | int, filter_order: int) -> Recording:
    """Apply IIR bandpass filteer to keep only frequencies above the threshold.

    {SINGLE_FREQ_PARAMS}
    :param filter_order: Controls the strength of filtering. Higher values mean narrower transition region.

    :return: Filtered recording
    """
    _assert_frequency_exists(sample_rate=recording.sample_rate, freq=frequency)
    new_data = recording.data.copy()
    new_data = _bandpass_iir(data=new_data, low=frequency, high=None, sample_rate=recording.sample_rate, filter_order=filter_order)
    return update_data(recording, new_data) 


# Internal


def _notch_fir(data: np.ndarray, frequency: int | float, sample_rate: int | float,
               transition_width: int | float, notch_width: int | float, window: str | None = None) -> np.ndarray:
    low_stop = frequency - notch_width / 2
    high_stop = frequency + notch_width / 2
    numtaps = int(sample_rate / transition_width)
    numtaps = numtaps if numtaps % 2 == 1 else numtaps + 1
    fir_coeff = breesy_scipy._si_firwin_bandstop_multibands(
        numtaps=numtaps,
        bands=[low_stop, high_stop],
        sample_rate=sample_rate,
        window=window
    )
    filtered_data = breesy_scipy._si_filtfilt(b=fir_coeff, a=[1.0], data=data)
    return filtered_data


def _notch_iir(data: np.ndarray, frequency: int | float, sample_rate: int | float, Q: int) -> np.ndarray:
    b, a = breesy_scipy._si_iirnotch(freq=frequency, Q=Q, sample_rate=sample_rate)
    filtered_data = breesy_scipy._si_filtfilt(b=b, a=a, data=data)
    return filtered_data


def _bandpass_fir(data: np.ndarray, low: int | float | None, high: int | float | None, sample_rate: int | float, 
                  transition_width: float, window: str | None = None) -> np.ndarray:

    numtaps = int(sample_rate / transition_width * 3.3)
    numtaps = numtaps if numtaps % 2 == 1 else numtaps + 1

    if low and high:
        fir_coeff = breesy_scipy._si_firwin_bandpass_multibands(
            numtaps=numtaps,
            bands=[low, high],
            sample_rate=sample_rate, 
            window=window
        )
    elif (not low) and high:
        fir_coeff = breesy_scipy._si_firwin_lowpass(numtaps=numtaps, band=high, 
                                                     sample_rate=sample_rate, window=window)
    elif (not high) and low:
        fir_coeff = breesy_scipy._si_firwin_highpass(numtaps=numtaps, band=low, 
                                                     sample_rate=sample_rate, window=window)
    else:
        raise BreesyInternalError(f"At least one frequency should be not None, received: {low=}, {high=}")

    filtered_data = breesy_scipy._si_filtfilt(b=fir_coeff, a=[1.0], data=data)
    return filtered_data


def _bandpass_iir(data: np.ndarray, low: int | float | None, high: int | float | None, sample_rate: int, filter_order: int):
    if low and high:
        the_filter = breesy_scipy._si_butter_bandpass_sos(low=low, high=high, sample_rate=sample_rate, filter_order=filter_order)
    elif (not low) and high:
        the_filter = breesy_scipy._si_butter_lowpass_sos(freq=high, sample_rate=sample_rate, filter_order=filter_order)
    elif (not high) and low:
        the_filter = breesy_scipy._si_butter_highpass_sos(freq=low, sample_rate=sample_rate, filter_order=filter_order)
    else:
        raise BreesyInternalError(f"At least one frequency should be not None, received: {low=}, {high=}")

    return breesy_scipy._si_sosfiltfilt(sos=the_filter, data=data)


def _get_fir_params_from_smoothness(smoothness_factor: int | float, notch_freq: int | float) -> tuple[int | float, int | float]:
    """Uses empirical scaling suggested by Claude Opus 4.5, needs rechecking."""
    bandwidth_3db = notch_freq * smoothness_factor
    notch_width = max(bandwidth_3db * 0.65, 0.1)
    transition_width = max(bandwidth_3db * 0.5, 0.1)
    return round(transition_width, 3), round(notch_width, 3)


def _assert_bandpass_low_high(sample_rate: int | float, low: float | int, high: float | int) -> None:
    low_high_range_valid = low <= high
    if not low_high_range_valid:
        raise BreesyError("Low frequency bound is higher than the high frequency bound",
                          "Provide a low frequency bound value that is less than or equal to the high frequency bound")
    high_bound_exceeds_nyquist = high >= (sample_rate / 2)
    if high_bound_exceeds_nyquist:
        raise BreesyError("High frequency bound exceeds Nyquist frequency",
                          "Provide a high frequency bound value that is less than half of the sample rate")
    if (high <= 0) or (low <= 0):
        raise BreesyError("One or both frequency bounds are not positive",
                          "Ensure both frequency bounds are positive numbers")


def _assert_frequency_exists(sample_rate: int | float, freq: float | int) -> None:
    if freq is None or freq <= 0:  # TODO: test if None check is needed
        raise BreesyError("Frequency bound should be a positive value.",
                          "Provide a frequency bound value that is a positive number.")
    freq_exists = (sample_rate // 2) >= freq
    if not freq_exists:
        raise BreesyError("Frequency bound exceeds Nyquist frequency",
                          "Provide a frequency bound value that is less than half of the sample rate")