import typing
from dataclasses import InitVar, dataclass, field
from pathlib import Path
from typing import List, Literal, Optional, Protocol, Union

import mne
import numpy as np

from zyplib.annotation.event import Event

SignalUnitType = Literal['V', 'uV', 'mV', '?']


def to_mne_raw(
    data: np.ndarray,
    srate: float,
    channels: typing.Sequence[str],
    events: typing.Sequence[Event] | None = None,
    ch_types: Union[str, typing.Sequence[str]] = 'eeg',
):
    if data.shape[0] != len(channels):
        raise ValueError(
            f'Data shape mismatch: data has {data.shape[0]} channels, '
            f'but {len(channels)} channel names were provided. '
            f'MNE expects data in (n_channels, n_times) format.'
        )
    ch_types_arg = ch_types if isinstance(ch_types, str) else list(ch_types)
    info = mne.create_info(
        channels,
        srate,
        ch_types=typing.cast(str | list[str], ch_types_arg),  # pyright: ignore[reportArgumentType]
        verbose=False,
    )
    raw = mne.io.RawArray(data, info)
    if events is not None:
        onset = np.array([event.onset for event in events])
        duration = np.array([event.duration for event in events])
        description = np.array([event.name for event in events])
        raw.set_annotations(
            mne.Annotations(onset=onset, duration=duration, description=description)
        )
    return raw


class EEGSampleInterface(Protocol):
    """EEG 文件接口"""

    filepath: Path  # 信号文件路径
    meta: dict  # 信号的元数据
    events: list[Event]  # 信号的标注
    name: str  # 信号的名称
    channels: list[str]  # 信号通道标签, 长度等于 `n_channels`
    n_channels: int  # 信号的通道数量
    srate: float  # 信号的采样率
    t_duration: float  # 信号的时间长度; 单位: 秒
    n_samples: int  # 信号的时间样本数量; n_samples = t_duration * sample_rate
    data: np.ndarray  # 信号数据; 形状为 (n_channels, n_samples)
    units: dict[str, SignalUnitType] | SignalUnitType

    def get_epoch(self, t_start: float, t_duration: float) -> np.ndarray:
        """获取指定时间范围内的信号数据

        Example
        ----------
        ```py
        >>> signals.get_epoch(t_start=0, t_duration=1)
        >>> # 获取从0秒到1秒的信号数据
        ```

        Parameters
        ----------
        - `t_start` : `float`
            - 开始时间; 单位: 秒
        - `t_duration` : `float`
            - 时间长度; 单位: 秒

        Returns
        ----------
        - `np.ndarray`
            - 信号数据; 形状为 (n_channels, n_samples)
        """
        ...

    def get_events(
        self, t_start: float, t_duration: float | None = None
    ) -> list[Event] | None:
        """获取指定时间范围内的标注"""
        ...

    def add_event(self, event: Event): ...

    def delete_event(self, event: Event): ...

    def pick_event(self, t_at: float, episilon: float = 0.1) -> list[Event]: ...


class EventMixin:
    def _sort_events(self: typing.Any):
        self.events.sort(key=lambda x: x.onset)

    def get_events(
        self: typing.Any, t_start: float, t_duration: float | None = None
    ) -> list[Event]:
        self._sort_events()
        events = []
        for event in self.events:
            if t_start <= event.onset <= t_start + typing.cast(float, t_duration):
                events.append(event)
        return events

    def pick_event(self: typing.Any, t_at: float, episilon: float = 1e-3) -> list[Event]:
        """获取在 t_at 附近的标注


        Parameters
        ----------
        - `t_at` : `float`
            - 时间点
        - `episilon` : `float`
            - 时间精度; 指在多大程度上我们可以允许 t_at 超过给定的浮点范围

        Returns
        ----------
        - `list[Event]`
            - 在 t_at 附近的标注
        """
        self._sort_events()
        events = []
        for event in self.events:
            if event.duration is None or event.duration <= episilon:
                if t_at - episilon <= event.onset <= t_at + episilon:
                    events.append(event)
            else:
                if (
                    event.onset - episilon
                    <= t_at
                    <= event.onset + event.duration + episilon
                ):
                    events.append(event)
        return events

    def add_event(self: typing.Any, event: Event) -> bool:
        if event in self.events:
            return False
        self.events.append(event)
        return True

    def delete_event(self: typing.Any, event: Event):
        self.events.remove(event)


class ConvertMixin:
    def to_mne_raw(self: 'EEGSampleInterface') -> mne.io.RawArray:
        """将当前对象转换为 mne.io.RawArray"""
        return to_mne_raw(self.data, self.srate, self.channels, self.events)

    def write_hdf5_file(self: 'EEGSampleInterface', filepath: str):
        """将当前脑电信号写入 HDF5 文件"""
        from .eeg_hdf5 import save_eeg_hdf5

        return save_eeg_hdf5(filepath, self)

    def write_pkl_file(self: 'EEGSampleInterface', filepath: str):
        """将当前脑电信号写入 HDF5 文件"""
        from .pkl import save_pkl

        values = {
            'filepath': self.filepath,
            'meta': self.meta,
            'name': self.name,
            'channels': self.channels,
            'n_channels': self.n_channels,
            'srate': self.srate,
            't_duration': self.t_duration,
            'n_samples': self.n_samples,
            'data': self.data,
            'units': self.units,
        }
        # filter none
        values = {k: v for k, v in values.items() if v is not None}
        return save_pkl(filepath, **values)

    def write_edf_file(self: 'EEGSampleInterface', filepath: str):
        """将当前脑电信号写入 EDF 文件

        将实现了 EEGSampleInterface 接口的对象数据写入到 EDF 文件中。
        使用 pyedflib.highlevel 的高级 API 进行写入，支持事件注释。

        Parameters
        ----------
        - `filepath` : `str`
            - 输出 EDF 文件的路径

        Raises
        ----------
        - `ImportError`
            - 如果 pyedflib 库未安装
        - `ValueError`
            - 如果数据格式不符合 EDF 要求
        - `RuntimeError`
            - 如果写入过程中发生错误
        """
        try:
            import pyedflib
            from pyedflib import highlevel
        except ImportError:
            raise ImportError(
                'pyedflib is required for writing EDF files. '
                'Please install it with: pip install pyedflib'
            )

        # 检查数据的有效性
        if self.data is None or self.data.size == 0:
            raise ValueError('Data is empty, cannot write EDF file')

        if not isinstance(self.srate, (int, float)):
            raise ValueError('Sample rate must be a single numeric value for EDF format')

        if self.n_channels != len(self.channels):
            raise ValueError(
                f'Number of channels in data ({self.n_channels}) does not match '
                f'number of channel names ({len(self.channels)})'
            )

        # 创建信号头部信息
        signal_headers = highlevel.make_signal_headers(
            self.channels, sample_frequency=int(self.srate)
        )

        # 设置物理单位
        if hasattr(self, 'units') and self.units:
            if isinstance(self.units, dict):
                for i, ch_name in enumerate(self.channels):
                    unit = self.units.get(ch_name, 'uV')
                    signal_headers[i]['physical_dimension'] = unit
            elif isinstance(self.units, str):
                # 如果 units 是单个字符串，应用到所有通道
                for signal_header in signal_headers:
                    signal_header['physical_dimension'] = self.units

        # 创建文件头部信息
        file_header = highlevel.make_header()

        # 设置患者信息（如果有的话）
        if hasattr(self, 'meta') and self.meta and isinstance(self.meta, dict):
            if 'patient_name' in self.meta:
                file_header['patient_name'] = str(self.meta['patient_name'])
            if 'recording_id' in self.meta:
                file_header['recording_id'] = str(self.meta['recording_id'])

        # 设置文件名作为默认记录 ID
        if hasattr(self, 'name') and self.name:
            file_header['recording_id'] = self.name

        # 准备注释数据
        annotations = []
        if hasattr(self, 'events') and self.events:
            for event in self.events:
                onset = float(event.onset)
                duration = float(event.duration) if event.duration is not None else -1
                description = str(event.name)
                annotations.append([onset, duration, description])

        # 将注释添加到header中
        if annotations:
            file_header['annotations'] = annotations

        try:
            # 使用 highlevel API 一次性写入信号和注释
            highlevel.write_edf(filepath, self.data, signal_headers, file_header)
        except Exception as e:
            raise RuntimeError(f'Failed to write EDF file: {str(e)}')


class PlotMixin:
    def plot(
        self: 'EEGSampleInterface',
        t_start: float,
        t_duration: float,
        show=True,
        save_path=None,
        figsize=None,
        **kwargs,
    ):
        """绘制信号数据"""
        epoch = self.get_epoch(t_start, t_duration)
        from zyplib.vis.plot_eeg import plot_signals

        return plot_signals(
            epoch,
            fs=typing.cast(int, self.srate),
            labels=self.channels,
            show=show,
            save_path=save_path,
            figsize=figsize,
            **kwargs,
        )


@dataclass
class BaseEEGSample(EventMixin, ConvertMixin, PlotMixin):
    """一个最简单的实现了 EEGSampleInterface 接口的类"""

    data: np.ndarray
    srate: float
    channels: Optional[list[str]]
    name: str | None
    filepath: Path | None
    units: SignalUnitType | dict[str, SignalUnitType] | None
    meta: dict | None
    events: list[Event] | None = None
    duration: InitVar[float | None] = None

    def __post_init__(self, duration: float | None):
        if self.channels is None:
            self.channels = [f'ch{i}' for i in range(self.n_channels)]
        self._duration = duration
        # self.events 会自动调用 EventMixin 的属性, 所以不要去设置

    @property
    def n_channels(self) -> int:
        return self.data.shape[0]

    @property
    def t_duration(self) -> float:
        if self._duration:
            return float(self._duration)
        return self.data.shape[1] / self.srate

    @property
    def n_samples(self) -> int:
        if self._duration:
            return int(self._duration * self.srate)
        return int(self.data.shape[1])

    def get_epoch(self, t_start: float, t_duration: float) -> np.ndarray:
        idx_start = int(t_start * self.srate)
        idx_end = idx_start + int(t_duration * self.srate)
        return self.data[:, idx_start:idx_end]

    @staticmethod
    def from_mne_raw(raw: mne.io.Raw) -> 'BaseEEGSample':
        """根据 mne.io.Raw 对象构建 BaseEEGSample"""
        data = typing.cast(np.ndarray, raw.get_data())
        srate = float(raw.info['sfreq'])
        channels = list(raw.ch_names)
        filenames = getattr(raw, 'filenames', None)
        filepath = Path(typing.cast(str, filenames[0])) if filenames else None
        units = {ch['ch_name']: ch.get('unit_str', '?') for ch in raw.info['chs']}
        # 将 MNE 注释转换为 Event
        events: list[Event] = []
        if raw.annotations is not None:
            for ann in raw.annotations:
                ann_dict = typing.cast(typing.Mapping[str, typing.Any], ann)
                name = str(ann_dict['description'])
                onset = float(ann_dict['onset'])
                duration = float(ann_dict['duration'])
                events.append(Event(name, onset, duration))
        return BaseEEGSample(
            data=data,
            srate=float(srate),
            channels=channels,
            name=filepath.name if filepath else None,
            filepath=filepath,
            units=units,
            meta={'mne_info': raw.info},
            events=events,
        )

    # def to_mne_raw(self) -> mne.io.Raw:
    #     """将当前对象转换为 mne.io.RawArray"""
    #     return to_mne_raw(self.data, self.srate, self.channels, self.events)
