import typing
from dataclasses import dataclass
from functools import singledispatch
from typing import Dict, List, Literal, Tuple, Union

import mne
import numpy as np

__all__ = [
    'Event',
    'EventList',
    'filter_event',
    'split_event',
    'merge_events',
    'events_from_array',
    'events_to_dict',
    'events_to_list',
    'events_to_npy',
]


@dataclass
class Event:
    name: Union[str, int]
    onset: float
    duration: float

    def copy(self):
        return Event(self.name, self.onset, self.duration)

    def __eq__(self, other: 'Event'):
        if not isinstance(other, Event):
            return False
        return (
            self.name == other.name
            and np.isclose(self.onset, other.onset, rtol=1e-4)
            and np.isclose(self.duration, other.duration, rtol=1e-4)
        )

    def to_dict(self) -> Dict[str, Union[str, int, float]]:
        """Convert the Event to a dictionary for JSON serialization."""
        return {'name': self.name, 'onset': self.onset, 'duration': self.duration}


class EventList(list):
    """A list-like container for Event objects with additional functionality."""

    def __init__(self, events: typing.Sequence[Event] | None = None):
        super().__init__(events or [])

    def copy(self) -> 'EventList':
        """Return a deep copy of the EventList."""
        return EventList([event.copy() for event in self])

    def filter(self, name: Union[str, int]) -> 'EventList':
        """Filter events by name."""
        return EventList(filter_event(self, name))

    def split(self) -> Dict[str | int, 'EventList']:
        """Split events into multiple EventLists by name."""
        result = {}
        split_dict = split_event(self)
        for name, events in split_dict.items():
            result[name] = EventList(events)
        return result

    def merge(self, distance: float) -> 'EventList':
        """Merge events that are closer than the specified distance."""
        return EventList(merge_events(self, distance))

    def to_dict(self) -> Dict:
        """Convert events to a dictionary format."""
        return events_to_dict(self)

    def to_list(self) -> List[List[float]]:
        """Convert events to a list of [onset, duration] pairs."""
        return events_to_list(self)

    def to_array(self) -> np.ndarray:
        """Convert events to a numpy array of [onset, duration] pairs."""
        return events_to_npy(self)

    @classmethod
    def from_array(
        cls,
        array: Union[typing.Sequence[typing.Sequence[float]], np.ndarray],
        name: Union[str, int],
    ) -> 'EventList':
        """Create an EventList from an array-like object."""
        return cls(events_from_array(array, name))

    @staticmethod
    def from_mne(mne_event: mne.Annotations) -> 'EventList':
        return EventList(
            [
                Event(
                    typing.cast(typing.Any, e).description,
                    typing.cast(typing.Any, e).onset,
                    typing.cast(typing.Any, e).duration,
                )
                for e in mne_event
            ]
        )

    def to_mne(self) -> mne.Annotations:
        return mne.Annotations(
            onset=np.array([event.onset for event in self]),
            duration=np.array([event.duration for event in self]),
            description=np.array([event.name for event in self]),
        )

    def __eq__(self, other: 'EventList') -> bool:
        """Compare two EventLists for equality."""
        if not isinstance(other, EventList):
            return False
        if len(self) != len(other):
            return False
        return all(a == b for a, b in zip(self, other))

    def __str__(self) -> str:
        """String representation of the EventList."""
        return f'EventList({super().__str__()})'

    def __repr__(self) -> str:
        """Detailed string representation of the EventList."""
        return self.__str__()


def filter_event(events: typing.Sequence[Event], name: Union[str, int]) -> List[Event]:
    return [event for event in events if event.name == name]


def split_event(events: typing.Sequence[Event]) -> Dict[str | int, List[Event]]:
    """根据 events 中的事件名称将 events 划分为多个事件列表"""
    events_dict: Dict[str | int, List[Event]] = {}
    for event in events:
        if event.name not in events_dict:
            events_dict[event.name] = []
        events_dict[event.name].append(event)
    return events_dict


def create_events_with_same_duration(
    position: typing.Sequence[float],
    duration: float,
    name: Union[str, typing.Sequence[str]] | None = None,
    at: Literal['start', 'center', 'end'] = 'start',
) -> List[Event]:
    """快速创建一组具有相同持续时间的事件

    Parameters
    ----------
    - `position` : `List[float]`
        - 事件的起始时间列表
    - `duration` : `float`
        - 每个事件的持续时间
    - `name` : `Union[str, int]`
        - 事件名称

    Returns
    -------
    - `List[Event]`
        - 创建的事件列表
    """
    if name is None:
        name = [f'Event_{i}' for i in range(len(position))]
    elif isinstance(name, (str, int)):
        name = [name] * len(position)
    elif isinstance(name, list) and len(name) != len(position):
        raise ValueError(
            f'事件名称列表长度 ({len(name)}) 必须与起始时间列表长度 ({len(position)}) 相同.'
        )
    start = []
    if at == 'start':
        start = position
    elif at == 'center':
        start = [a - duration / 2 for a in position]
    elif at == 'end':
        start = [a - duration for a in position]
    return [Event(n, s, duration) for n, s in zip(name, start)]


def merge_events(events: typing.Sequence[Event], distance: float) -> list[Event]:
    """将距离小于 distance 的事件合并为一个事件

    本方法无视事件名称, 仅根据 onset 和 duration 进行合并

    Parameters
    ----------
    - `events` : list[Event]
        - 事件列表
    - `distance` : float
        - 事件合并的距离阈值, 和 duration onset 等为同单位

    Returns
    ----------
    - `list[Event]`
        - 合并后的事件列表
    """
    # 1. 将事件按 onset 排序
    events = sorted(events, key=lambda x: x.onset)
    # 2. 合并事件
    merged_events: List[Event] = []
    for event in events:
        if len(merged_events) == 0:
            merged_events.append(event)
        else:
            last_event = merged_events[-1]
            last_event_end = last_event.onset + last_event.duration
            if event.onset - last_event_end > distance:
                merged_events.append(event)
            else:
                last_event.duration = event.onset + event.duration - last_event.onset
    return merged_events


def events_from_array(
    array: Union[typing.Sequence[typing.Sequence[float]], np.ndarray], name
) -> List[Event]:
    """从数组中提取事件列表

    Example
    ----------
    ```py
    >>> events_from_array([[0, 1], [1, 2]], 'spindle')
    >>> events_from_array(np.array([[0, 1], [1, 2]]), 'spindle')
    ```

    Parameters
    ----------
    - `array` : Union[List[List[float]], np.ndarray]
        - 原始的事件列表

    Returns
    ----------
    - `List[Event]`
        - 事件对象列表
    """
    events = []
    for l in array:
        events.append(Event(name, l[0], l[1]))
    return events


def events_to_dict(events: typing.Sequence[Event]) -> Dict:
    dic = {
        'Onset': np.array([event.onset for event in events]),
        'Duration': np.array([event.duration for event in events]),
        'Name': np.array([event.name for event in events]),
    }
    return dic


def events_to_list(events: typing.Sequence[Event]) -> List[List[float]]:
    """将事件列表转换为数组

    Parameters
    ----------
    - `events` : `List[Event]`, 事件列表, Evnet 为 zyplib.annotation.type.Event 类型; 单位为秒

    Returns
    ----------
    - `List[[float, float]]`: 事件列表, [[onset, duration], ...]
    """
    return [[event.onset, event.duration] for event in events]


def events_to_npy(events: typing.Sequence[Event]) -> np.ndarray:
    """将事件列表转换为数组

    Parameters
    ----------
    - `events` : `List[Event]`, 事件列表, Evnet 为 zyplib.annotation.type.Event 类型; 单位为秒

    Returns
    ----------
    - `np.ndarray` : 事件列表, [[onset, duration], ...]
    """
    return np.array(events_to_list(events))
