import numpy as np

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


@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.

    :param recording: Input recording
    :param frequency: Frequency to remove in Hz
    :param smoothness_factor: Controls the smoothness of attenuation. Higher value means wider affected (attenuated) band.

    :return: Filtered recording
    """
    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) 


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

    :param recording: Input recording
    :param frequency: Frequency to remove in Hz
    :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

    :return: Filtered recording
    """    
    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) 


@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.

    :param recording: Input recording
    :param frequency: Frequency to remove in Hz
    :param quality_factor: Controls the strength and width of attenuation. Higher value means narrower and deeper removed band.

    :return: Filtered recording
    """
    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) 


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")


@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. This is done by applying a bandpass filter.

    :param recording: Input recording
    :param low: Lower frequency bound in Hz
    :param high: Upper frequency bound in Hz
    :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) 


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

    :param recording: Input recording
    :param low: Lower frequency bound in Hz
    :param high: Upper frequency bound in Hz
    :param transition_width: Controls the width of 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=transition_width)
    return update_data(recording, new_data) 


@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.

    :param recording: Input recording
    :param low: Lower frequency bound in Hz
    :param high: Upper frequency bound in Hz
    :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()

    bandpass = _si_butter_bandpass_sos(low=low, high=high, sample_rate=recording.sample_rate, filter_order=filter_order)
    new_data = _si_sosfiltfilt(sos=bandpass, data=new_data)
    return update_data(recording, new_data) 


def _notch_fir(data: np.ndarray, frequency: int | float, sample_rate: int | float,
               transition_width: int | float, notch_width: int | float) -> 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
    )
    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, high: int | float, sample_rate: int | float,
                  transition_width: float) -> np.ndarray:
    numtaps = int(sample_rate / transition_width * 3.3)
    numtaps = numtaps if numtaps % 2 == 1 else numtaps + 1
    fir_coeff = breesy_scipy._si_firwin_bandpass_multibands(
        numtaps=numtaps,
        bands=[low, high],
        sample_rate=sample_rate
    )
    filtered_data = breesy_scipy._si_filtfilt(b=fir_coeff, a=[1.0], data=data)
    return filtered_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)
