import numpy as np

from breesy.Event import Event
from breesy.Recording import Recording
from breesy.errors import BreesyError
from breesy.type_hints import enforce_type_hints
from breesy.breesy_scipy import _si_butter_bandpass_sos, _si_sosfiltfilt


@enforce_type_hints
def take_n_seconds(recording: Recording, duration: float | int = 5.0):
    """Take the first N seconds from a recording.

    :param recording: Input recording
    :param duration: Number of seconds to take
    
    :return: New recording with only the first N seconds
    """
    return cut_by_second_range(recording, 0, duration)


@enforce_type_hints
def cut_by_second_range(recording: Recording, start_second: float | int, end_second: float | int):
    """Cut a recording to a specific time range.

    :param recording: Input recording
    :param start_second: Start time in seconds
    :param end_second: End time in seconds
    
    :return: New recording containing only the specified time range
    """
    if start_second > end_second:
        raise BreesyError(
            f"Requested to cut from {start_second} s to {end_second} s, but the start time is higher than the end time.",
            "Make sure that the start time is lower than the end time."
        )

    if end_second > recording.duration:
        raise BreesyError(
            f"Requested to cut from {start_second} s to {end_second} s, but the recording has only {recording.duration}s.",
            "Make sure that the start and end seconds are within the range of the recording duration. "
            "You can see the recording duration like so: recording.duration."
        )

    start_index = recording.sample_rate * start_second
    end_index = recording.sample_rate * end_second

    return cut_by_sample_range(recording, start_index, end_index)


@enforce_type_hints
def take_n_samples(recording: Recording, n_samples: int) -> Recording:
    """Take the first N samples from a recording.

    :param recording: Input recording
    :param n_samples: Number of samples to take
    
    :return: New recording with only the first N samples
    """
    if n_samples > recording.number_of_samples:
        raise BreesyError(
            f"Requested to take {n_samples} samples, but the recording has only {recording.number_of_samples} samples.",
            "Make sure that the number of samples requested is less than the number of samples in the recording. "
            "You can see the number of samples in the recording like so: recording.num_samples."
        )

    return cut_by_sample_range(recording, 0, n_samples - 1)


@enforce_type_hints
def cut_by_sample_range(recording: Recording, start_index: int, end_index: int) -> Recording:
    """Cut a recording to a specific sample range.

    :param recording: Input recording
    :param start_index: Start sample index
    :param end_index: End sample index
    
    :return: New recording containing only the specified sample range
    """
    if start_index < 0 or end_index > recording.number_of_samples:
        raise BreesyError(
            f"Requested to cut from {start_index} to {end_index}, but the recording has only {recording.number_of_samples} samples.",
            "Make sure that the start and end sample indices are within the range of the recording samples. "
            "You can see the number of samples in the recording like so: recording.number_of_samples."
        )

    # Take both ends inclusive so that the events themselves also fit in the epoch
    new_eeg_data = recording.data[:, start_index:end_index + 1]
    events_in_new_data = [event for event in (recording.events or []) if start_index <= event.index <= end_index]
    events_with_offset = [Event(e.name, e.index - start_index) for e in events_in_new_data]

    return Recording(
        data=new_eeg_data,
        channel_names=recording.channel_names,
        sample_rate=recording.sample_rate,
        events=events_with_offset,
        metadata=recording.metadata
    )


@enforce_type_hints
def update_data(recording: Recording, new_eeg_data: np.ndarray) -> Recording:
    """Create a new recording with updated EEG data.

    :param recording: Input recording
    :param new_eeg_data: New EEG data array to use
    
    :return: New recording with updated data but same metadata and events
    """
    if recording.number_of_samples != new_eeg_data.shape[1]:
        raise BreesyError(
            f"New EEG data has a different number of samples ({new_eeg_data.shape[1]}) than the original EEG data ({recording.number_of_samples}).",
            "Make sure that the new EEG data has the same number of samples as the original EEG data. "
            "You can see the number of samples in the original EEG data like so: recording.number_of_samples."
        )

    return Recording(
        data=new_eeg_data,
        channel_names=recording.channel_names,
        sample_rate=recording.sample_rate,
        events=recording.events,
        metadata=recording.metadata
    )


@enforce_type_hints
def add_events(recording: Recording, events: list[Event]) -> Recording:
    """Create a new recording with additional events.

    :param recording: Input recording
    :param events: List of events to add
    
    :return: New recording with added events
    """
    new_events = recording.events or [] + events
    return Recording(
        data=recording.data,
        channel_names=recording.channel_names,
        sample_rate=recording.sample_rate,
        events=new_events,
        metadata=recording.metadata
    )


@enforce_type_hints
def remove_channels(recording: Recording, channels_to_remove: list[str]) -> Recording:
    """Create a new recording with specified channels removed.

    :param recording: Input recording
    :param channels_to_remove: List of channel names to remove
    
    :return: New recording without the specified channels
    """
    missing_channels = [channel for channel in channels_to_remove if channel not in recording.channel_names]
    if missing_channels:
        raise BreesyError(
            f"Channels to remove contain channel names that are not present in the recording, namely: {missing_channels}.",
            "Make sure that the channel names provided to remove are present in the recording. "
            "Double check that you haven't removed those channels for this data already. "
            "You can see the channel names in the recording with: recording.channel_names."
        )

    channel_indices_to_remove = [recording.channel_names.index(channel) for channel in channels_to_remove]
    new_eeg_data = np.delete(recording.data, channel_indices_to_remove, axis=0)
    new_channel_names = [channel for channel in recording.channel_names if channel not in channels_to_remove]
    return Recording(
        data=new_eeg_data,
        channel_names=new_channel_names,
        sample_rate=recording.sample_rate,
        events=recording.events,
        metadata=recording.metadata
    )


@enforce_type_hints
def select_channels(recording: Recording, channels_to_select: list[str]) -> Recording:
    """Create a new recording with only the specified channels.

    :param recording: Input recording
    :param channels_to_select: List of channel names to keep
    
    :return: New recording with only the specified channels
    """
    missing_channels = [channel for channel in channels_to_select if channel not in recording.channel_names]
    if missing_channels:
        raise BreesyError(
            f"Channels to select contain channel names that are not present in the recording, namely: {missing_channels}.",
            "Make sure that the channel names provided to select are present in the recording. "
            "You can see the channel names in the recording with: recording.channel_names."
        )

    channel_indices_to_select = [recording.channel_names.index(channel) for channel in channels_to_select]
    new_eeg_data = recording.data[channel_indices_to_select]
    return Recording(
        data=new_eeg_data,
        channel_names=channels_to_select,
        sample_rate=recording.sample_rate,
        events=recording.events,
        metadata=recording.metadata
    )


@enforce_type_hints
def set_channel_names(recording: Recording, ch_names: list[str]) -> Recording:
    """Set new channel names for a recording.

    :param recording: Input recording
    :param ch_names: List of new channel names
    
    :return: New recording with updated channel names
    """
    if len(ch_names) != recording.number_of_channels:
        raise BreesyError(
            f"{len(ch_names)} channel names are provided, but data has {recording.number_of_channels} channels. "
            "Please recheck your list of channel names. "
        )
    return Recording(
        data=recording.data,
        channel_names=ch_names,
        sample_rate=recording.sample_rate,
        events=recording.events,
        metadata=recording.metadata
    )


@enforce_type_hints
def get_events_by_name(recording: Recording, event_name: str) -> list[Event]:
    """Return a list of events with a specific name.

    :param recording: Input recording
    :param event_name: Name of events to find
    
    :return: List of events with the specified name
    """
    return [event for event in recording.events if event.name == event_name]


@enforce_type_hints
def split_by_continuous_events(recording: Recording, start_event_name: str, end_event_name: str) -> list[Recording]:
    """Split a recording into segments between start and end events.

    :param recording: Input recording
    :param start_event_name: Name of events marking segment starts
    :param end_event_name: Name of events marking segment ends
    
    :return: List of recording segments between each start-end event pair
    """
    if not recording.events:
        raise BreesyError(
            "No events found in the recording.",
            "Make sure that the recording has events. You can see the number of events like so: len(recording.events)."
        )

    start_events = get_events_by_name(recording, start_event_name)
    end_events = get_events_by_name(recording, end_event_name)

    if len(start_events) != len(end_events):
        raise BreesyError(
            f"Found {len(start_events)} start events and {len(end_events)} end events. "
            "Expected the number of start and end events to be the same.",
            "Make sure each start event has a corresponding end event. "
            "You can see the number of start events like so: get_events_by_name(recording, start_event_name); "
            "number of end events: get_events_by_name(recording, end_event_name)."
        )

    epoch_recordings = []
    for start_event, end_event in zip(start_events, end_events):
        if start_event.index >= end_event.index:
            pair_number = start_events.index(start_event) + 1
            raise BreesyError(
                f"When going pairwise through events, found a pair (nr. {pair_number}) where start event index ({start_event.index}) is higher than the end event index ({end_event.index}).",
                "Inspect the events manually, for start events: get_events_by_name(recording, start_event_name); "
                "end events: get_events_by_name(recording, end_event_name). "
                "Make sure that the start event index is lower than the end event index. "
            )

        epoch_recording = cut_by_sample_range(recording, start_event.index, end_event.index)
        epoch_recordings.append(epoch_recording)

    return epoch_recordings


@enforce_type_hints
def split_by_events(recording: Recording, event_name: str, duration: float | int = 5.0) -> list[Recording]:
    """Split a recording into segments starting at specified events.

    :param recording: Input recording
    :param event_name: Name of events to use as segment starts
    :param duration: Duration of each segment in seconds
    
    :return: List of recording segments starting at each event
    """
    events = get_events_by_name(recording, event_name)

    epoch_recordings = []
    for event in events:
        end_index = event.index + recording.sample_rate * duration
        epoch_recording = cut_by_sample_range(recording, event.index, end_index)
        epoch_recordings.append(epoch_recording)

    return epoch_recordings


@enforce_type_hints
def split_by_window_duration(recording: Recording, window_duration: float | int, overlap_duration: float | int = 0.0) -> list[Recording]:
    """Split a recording into equal-sized segments of specified duration with optional overlap.

    :param recording: Input recording
    :param window_duration: Duration of each segment in seconds
    :param overlap_duration: Duration of overlap between consecutive segments in seconds (default: 0.0)
    
    :return: List of equal-sized recording segments
    """
    if overlap_duration < 0:
        raise BreesyError(
            f"Overlap duration ({overlap_duration}s) cannot be negative.",
            "Provide a non-negative overlap duration value."
        )
    
    if overlap_duration >= window_duration:
        raise BreesyError(
            f"Overlap duration ({overlap_duration}s) must be less than window duration ({window_duration}s).",
            "Provide an overlap duration that is smaller than the window duration."
        )

    num_samples_per_epoch = int(window_duration * recording.sample_rate)
    
    if num_samples_per_epoch > recording.number_of_samples:
        raise BreesyError(
            f"Epoch duration ({window_duration}s) is longer than the recording duration ({recording.duration}s).",
            "Make sure that the epoch duration is shorter than the recording duration. "
            "You can see the recording duration like so: recording.duration."
        )

    step_duration = window_duration - overlap_duration
    step_samples = int(step_duration * recording.sample_rate)
    
    max_start_sample = recording.number_of_samples - num_samples_per_epoch
    num_epochs = (max_start_sample // step_samples) + 1

    epoch_recordings = []
    for i in range(num_epochs):
        start_index = i * step_samples
        end_index = start_index + num_samples_per_epoch
        
        # skip window if out of bounds
        if end_index > recording.number_of_samples:
            break
            
        epoch_data = recording.data[:, start_index:end_index]
        epoch_recording = Recording(
            data=epoch_data,
            channel_names=recording.channel_names,
            sample_rate=recording.sample_rate,
            events=recording.events,
            metadata=recording.metadata
        )
        epoch_recordings.append(epoch_recording)

    return epoch_recordings


# TODO: update function to FIR and unhardcode
@enforce_type_hints
def select_bandwidth(recording: Recording, low: float | int, high: float | int, filter_order: int = 5) -> Recording:
    """Filter recording 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: Order of the filter
    
    :return: Filtered recording
    """
    data = recording.data
    sample_rate = recording.sample_rate

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

    bandpass = _si_butter_bandpass_sos(low, high, sample_rate, filter_order)

    new_eeg_data = _si_sosfiltfilt(bandpass, data)

    return update_data(recording, new_eeg_data)