"""
实现了 EEGSampleInterface 接口的文件读取类
"""

import os
import re
import typing
from pathlib import Path

import numpy as np
import pyedflib
import scipy.io as sio
from mne.io import read_raw_edf

from zyplib.annotation.event import Event
from zyplib.utils.print import print_warning

from ._base import ConvertMixin, EventMixin, PlotMixin


class MNEArgs(typing.TypedDict, total=False):
    eog: list[str] | tuple[str] | None
    misc: list[str] | tuple[str] | None
    stim_channel: str
    exclude: list[str] | str | None
    infer_types: bool
    include: list[str] | str | None
    preload: bool
    units: list[str] | str | None
    encoding: str
    exclude_after_unique: bool
    verbose: bool | None


class EDFFileSample(EventMixin, ConvertMixin, PlotMixin):
    """
    对 EDF 文件实现 EEGSampleInterface 接口

    可以使用 MNE 或 pyedflib 读取 EDF文件。

    注意如果 backend 为 pyedflib 且读入的 EDF 文件中含有多种不同采样率的通道，则会抛出错误。

    Parameters
    ----------
    - `filepath` : str
        EDF 文件的路径。
    - `backend` : str, optional
        要使用的后端，'mne' 或 'pyedflib'。默认为 'mne'。
    - `mne_args` : MNEArgs | None, optional
        传递给后端读取函数的额外关键字参数 (mne.read_raw_edf)。
    """

    def __init__(
        self,
        filepath: str,
        backend: str = 'pyedflib',
        pick_channels: typing.Sequence[str] | None = None,
        mne_args: MNEArgs | None = None,
    ):
        self.filepath = Path(filepath)
        self.backend = backend

        # Internal data storage
        self._data: np.ndarray | None = None
        self._channels: list[str] | None = None
        self._srate: float | list[float] | None = None
        self._units: dict[str, str] = {}
        self.raw: typing.Any | None = None  # for MNE backend
        self.events = []

        mne_args = mne_args or {}
        if pick_channels:
            mne_args['include'] = typing.cast(typing.Any, pick_channels)

        if backend == 'mne':
            self._init_mne(**mne_args)
        elif backend == 'pyedflib':
            self._init_pyedflib(pick_channels)
        else:
            raise ValueError(
                f"Unsupported backend: {backend}. Choose 'mne' or 'pyedflib'."
            )

    def _init_mne(self, **mne_kwargs):
        """使用 MNE 后端进行初始化。"""
        options = {'verbose': False, 'preload': False}
        options.update(mne_kwargs)
        self.raw = read_raw_edf(self.filepath, **typing.cast(typing.Any, options))  # pyright: ignore[reportArgumentType]

        # 从 raw 中提取事件
        for ann in self.raw.annotations:
            self.add_event(
                Event(
                    name=ann['description'], onset=ann['onset'], duration=ann['duration']
                )
            )

    def _init_pyedflib(self, pick_channels: typing.Sequence[str] | None = None):
        """使用 pyedflib 后端进行初始化。"""

        with pyedflib.EdfReader(str(self.filepath)) as f:
            all_channels = f.getSignalLabels()
            n_ch_total = f.signals_in_file

            # 确定要读取的通道
            if pick_channels is not None:
                # 找到指定通道的索引
                channel_indices = []
                selected_channels = []
                for ch_name in pick_channels:
                    if ch_name in all_channels:
                        idx = all_channels.index(ch_name)
                        channel_indices.append(idx)
                        selected_channels.append(ch_name)
                    else:
                        raise ValueError(
                            f"Channel '{ch_name}' not found in file. Available channels: {all_channels}"
                        )

                self._channels = selected_channels
                n_ch = len(channel_indices)
            else:
                # 读取所有通道
                channel_indices = list(range(n_ch_total))
                self._channels = all_channels
                n_ch = n_ch_total

            # 检查采样率一致性（只检查选中的通道）
            srates = [f.getSampleFrequency(i) for i in channel_indices]
            if len(set(srates)) > 1:
                print_warning(
                    f'注意: 导入 EDF 文件 {self.filepath} 内含有多种采样率 {srates}'
                )
            else:
                self._srate = float(srates[0])

            # 检查样本数一致性（只检查选中的通道）
            n_samples_per_channel = [f.getNSamples()[i] for i in channel_indices]
            if len(set(n_samples_per_channel)) > 1:
                raise ValueError(
                    'Selected channels have different number of samples, which is not supported by this reader yet.'
                )
            n_samples = n_samples_per_channel[0]

            # 只读取选中的通道数据
            self._data = np.zeros((n_ch, n_samples), dtype=np.float64)
            for idx, ch_idx in enumerate(channel_indices):
                self._data[idx, :] = f.readSignal(ch_idx)
                unit = str(f.getSignalHeader(ch_idx)['dimension'])
                self._units[self._channels[idx]] = unit

            # Read annotations
            annotations = f.readAnnotations()
            for onset, duration, description in zip(*annotations):
                self.add_event(Event(description, onset, duration))

    @property
    def meta(self) -> dict:
        return {
            'filepath': self.filepath,
            'backend': self.backend,
            'n_channels': self.n_channels,
            'srate': self.srate,
            't_duration': self.t_duration,
        }

    @property
    def name(self) -> str:
        return os.path.basename(self.filepath)

    @property
    def channels(self) -> list[str]:
        if self.backend == 'mne':
            return typing.cast(typing.Any, self.raw).ch_names
        return typing.cast(list[str], self._channels)

    @property
    def n_channels(self) -> int:
        return len(self.channels)

    @property
    def srate(self) -> float | list[float]:
        if self.backend == 'mne':
            return typing.cast(typing.Any, self.raw).info['sfreq']
        return typing.cast(float | list[float], self._srate)

    @property
    def n_samples(self) -> int:
        if self.backend == 'mne':
            return typing.cast(typing.Any, self.raw).n_times
        return typing.cast(np.ndarray, self._data).shape[1]

    @property
    def t_duration(self) -> float:
        return self.n_samples / typing.cast(typing.Any, self.srate)

    @property
    def data(self) -> np.ndarray:
        """
        获取 EEG 数据，单位为伏特 (V)。
        """
        if self.backend == 'mne' and self._data is None:
            data = typing.cast(typing.Any, self.raw).get_data()
            self._data = data
        return typing.cast(np.ndarray, self._data)

    @property
    def units(self) -> dict[str, str]:
        """
        获取每个通道的原始物理单位。
        """
        if self.backend == 'mne':
            # MNE stores units in 'unit' key of chs dict.
            # It also has a 'unit_mul' for scaling, but we'll keep it simple.
            raw = typing.cast(typing.Any, self.raw)
            return {ch['ch_name']: ch.get('unit_str', 'V') for ch in raw.info['chs']}
        return self._units

    def get_epoch(self, t_start: float, t_duration: float) -> np.ndarray:
        """
        获取一个时间段的数据。

        Parameters
        ----------
        t_start : float
            开始时间（秒）。
        t_duration : float
            持续时间（秒）。

        Returns
        -------
        np.ndarray
            请求的数据段。
        """
        srate = typing.cast(typing.Any, self.srate)
        if self.backend == 'mne':
            start_sample = int(t_start * srate)
            end_sample = start_sample + int(t_duration * srate)
            return typing.cast(typing.Any, self.raw).get_data(
                start=start_sample, stop=end_sample
            )

        # pyedflib backend logic
        start_sample = int(t_start * srate)
        if start_sample < 0 or start_sample >= self.n_samples:
            raise ValueError('t_start is out of bounds.')

        end_sample = start_sample + int(t_duration * srate)
        end_sample = min(end_sample, self.n_samples)  # clamp to data length

        epoch_data = typing.cast(np.ndarray, self._data)[:, start_sample:end_sample]

        return epoch_data

    def _require_raw(self):
        if self.raw is None:
            raise RuntimeError('MNE raw object is not initialized.')
        return self.raw

    def _require_data(self) -> np.ndarray:
        if self._data is None:
            raise RuntimeError('EEG data is not initialized.')
        return self._data

    def _require_scalar_srate(self) -> float:
        srate = self.srate
        if isinstance(srate, list):
            raise RuntimeError('Sampling rates are not uniform across channels.')
        return float(srate)


class MatFileSample(EventMixin, ConvertMixin, PlotMixin):
    """对 MAT 文件实现 EEGSampleInterface 接口"""

    DATA_PAT = r'^data(.*)'
    LABEL_PAT = r'^(channel|label)s?'
    SRATE_PAT = r'^(srate|sfreq|freq|sample_rate)s?'

    def __init__(
        self,
        filepath: str,
        srate: float | None = None,
        channels: typing.Sequence[str] | None = None,
        events: typing.Sequence[Event] | None = None,
        key_data: str | None = None,
        key_channels: str | None = None,
        key_srate: str | None = None,
        **kwargs,
    ):
        self.filepath = Path(filepath)

        self.data = None
        self.srate = srate
        self.channels = channels
        self.n_channels = 0
        self.n_samples = 0
        self.t_duration = 0
        self.events = []

        self.name = os.path.basename(filepath)
        self.meta = {}

        try:
            mat = sio.loadmat(filepath)
            mat = {k: v for k, v in mat.items() if not k.startswith('__')}
            self._resolve_mat(mat, key_data, key_channels, key_srate)
        except Exception as e:
            raise RuntimeError(f'文件读取错误: {e}')

        self.n_samples = typing.cast(np.ndarray, self.data).shape[1]
        self.t_duration = self.n_samples / typing.cast(typing.Any, self.srate)

    def _resolve_mat(
        self,
        mat: typing.Mapping,
        key_data: str | None,
        key_channels: str | None,
        key_srate: str | None,
    ):
        keys = list(mat.keys())
        data: np.ndarray

        # Use provided keys if available
        if key_data and key_data in mat:
            data = mat[key_data]
        else:
            for key in keys:
                if re.match(self.DATA_PAT, key):
                    data = mat[key]
                    break

        if key_channels and key_channels in mat:
            self.channels = mat[key_channels]
        else:
            for key in keys:
                if re.match(self.LABEL_PAT, key):
                    self.channels = mat[key]
                    break

        if key_srate and key_srate in mat:
            self.srate = mat[key_srate]
        else:
            for key in keys:
                if re.match(self.SRATE_PAT, key):
                    self.srate = mat[key]
                    break

        # =============== Section Beg: Data =============== #
        if data.ndim == 1:
            data = data.reshape(1, -1)
        # 数据默认应该是的第一维是通道数，第二维是时间
        CH, T = data.shape
        # 如果通道数大于时间，那么说明数据的第一维是时间，第二维是通道数
        if CH > T:
            data = data.T
            CH, T = data.shape
        self.data = data
        # --------------- Section End: Data --------------- #

        if self.channels is None:
            self.channels = [f'ch{i}' for i in range(CH)]

        if self.srate is None:
            raise RuntimeError('未找到采样率信息')

        self.n_channels = CH
        self.n_samples = T
        self.t_duration = T / typing.cast(typing.Any, self.srate)

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

    @property
    def units(self):
        channels = typing.cast(typing.Sequence[str], self.channels)
        return {str(c): '?' for c in channels}
