from dataclasses import InitVar, dataclass
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: List[str],
    events: List[Event] = None,
    ch_types: Union[str, List[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.'
        )
    info = mne.create_info(channels, srate, ch_types=ch_types, 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):
        self.events.sort(key=lambda x: x.onset)

    def get_events(self, 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 + t_duration:
                events.append(event)
        return events

    def pick_event(self, 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, event: Event) -> bool:
        if event in self.events:
            return False
        self.events.append(event)
        return True

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


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

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

        return save_eeg_hdf5(filepath, self)

    def write_pkl_file(self, 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, 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,
        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=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:
        return self._duration or self.data.shape[1] / self.srate

    @property
    def n_samples(self) -> int:
        return self._duration * self.srate if self._duration else 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 = raw.get_data()
        srate = raw.info['sfreq']
        channels = list(raw.ch_names)
        filepath = Path(raw.filenames[0]) if getattr(raw, 'filenames', None) 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:
                events.append(
                    Event(ann['description'], float(ann['onset']), float(ann['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)
