"""EEG connectivity feature extraction algorithms.

Provides various connectivity measures for EEG signals analysis:

- Pearson Correlation
- Phase Locking Value (PLV)
- Phase Lag Index (PLI)
- Coherence
- Weighted Phase Lag Index (wPLI)
- Debiased Squared wPLI
- Cross-correlation
- Coherency

#### References

- [1] https://mne.tools/mne-connectivity/stable/index.html
- [2] Bastos, A. M., & Schoffelen, J. M. (2016). A tutorial review of functional
    connectivity analysis methods and their interpretational pitfalls.
    Frontiers in systems neuroscience, 9, 175.
"""

import typing

from typing import Any, Literal, Optional, Tuple, Union, cast

import numpy as np
from mne_connectivity import spectral_connectivity_epochs
from mne_connectivity.viz import plot_connectivity_circle

from zyplib.utils.ensure import ensure_ndarray

ConnMethod = Literal[
    'coh',  # Coherence
    'cohy',  # Coherency
    'plv',  # Phase Locking Value
    'pli',  # Phase Lag Index
    'wpli',  # Weighted Phase Lag Index
    'wpli2_debiased',  # Debiased Squared wPLI
    'ppc',  # Pairwise Phase Consistency
    'ccd',  # Canonical Correlation Density
]


def compute_spectra_connectivity(
    data: np.ndarray,
    sfreq: float,
    method: ConnMethod = 'pli',
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    mode: str = 'multitaper',
    average: bool = True,
    **kwargs,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """计算 EEG 数据的连接性度量

    本函数封装了 mne-connectivity 以简化使用。

    Parameters
    ----------
    - `data` : `np.ndarray`
        - EEG 数据数组，形状为 [n_trials, n_channels, n_samples] 或 [n_channels, n_samples]
    - `sfreq` : `float`
        - 采样频率，单位为 Hz
    - `method` : `ConnMethod`
        - 计算的连接性度量：
        - 'coh': Coherence
        - 'cohy': Coherency
        - 'plv': Phase Locking Value
        - 'pli': Phase Lag Index
        - 'wpli': Weighted Phase Lag Index
        - 'wpli2_debiased': Debiased Squared wPLI
        - 'ppc': Pairwise Phase Consistency
        - 'ccd': Canonical Correlation Density
    - `fmin` : `float | None`
        - 感兴趣的最小频率。如果为 None，则使用最低可用频率
    - `fmax` : `float | None`
        - 感兴趣的最大频率。如果为 None，则使用最高可用频率
    - `average` : `bool`
        - 是否在频率和试验之间取平均，默认为 True

    Returns
    -------
    - `Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]`
        - 如果 average=True:
            - 连接性矩阵，形状为 [n_channels, n_channels]
        - 如果 average=False:
            - (连接性, 频率) 的元组：
            - 连接性: 形状为 [n_channels, n_channels, n_freqs]
            - 频率: 形状为 [n_freqs,]

    Notes
    -----
    对于基于相位的方法（PLV, PLI），建议进行预处理（例如过滤）。
    """
    data = cast(np.ndarray, ensure_ndarray(data))
    if data.ndim == 2:
        data = data[np.newaxis, ...]  # Add trial dimension
    default_kwargs = {
        'fmin': fmin,
        'fmax': fmax,
        'faverage': average,
        'n_jobs': 1,  # Avoid multiprocessing issues
        'verbose': False,
        'mode': mode,
    }
    kwargs = {**default_kwargs, **kwargs}
    # Compute connectivity using mne-connectivity
    conn = spectral_connectivity_epochs(data, method=method, sfreq=sfreq, **kwargs)

    return conn


def plot_connectivity(
    conn_matrix: np.ndarray,
    node_names: Optional[typing.Sequence] = None,
    title: Optional[str] = None,
    **kwargs,
) -> Any:
    """绘制连接性圆图。

    Parameters
    ----------
    - `conn_matrix` : `np.ndarray`
        - 形状为 (n_channels, n_channels) 的连接性矩阵
    - `node_names` : `list | None`
        - 节点名称列表。如果为 None，则使用数字索引
    - `title` : `str | None`
        - 图表的标题
    - `**kwargs` : `dict`
        - 传递给 mne_connectivity.viz.plot_connectivity_circle 的额外参数

    Notes
    -----
    此函数封装了 mne_connectivity.viz.plot_connectivity_circle，以简化使用。
    """
    if node_names is None:
        node_names = [str(i) for i in range(conn_matrix.shape[0])]

    fig = plot_connectivity_circle(conn_matrix, node_names, title=title, **kwargs)
    return fig


# Convenience functions for common connectivity measures
def coherence(
    data: np.ndarray,
    sfreq: float,
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    average: bool = True,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """Calculate coherence between channels.

    See compute_connectivity for parameter details.
    """
    return compute_spectra_connectivity(data, sfreq, 'coh', fmin, fmax, average=average)


def plv(
    data: np.ndarray,
    sfreq: float,
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    average: bool = True,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """Calculate Phase Locking Value between channels.

    See compute_connectivity for parameter details.
    """
    return compute_spectra_connectivity(data, sfreq, 'plv', fmin, fmax, average=average)


def pli(
    data: np.ndarray,
    sfreq: float,
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    average: bool = True,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """Calculate Phase Lag Index between channels.

    See compute_connectivity for parameter details.
    """
    return compute_spectra_connectivity(data, sfreq, 'pli', fmin, fmax, average=average)


def wpli(
    data: np.ndarray,
    sfreq: float,
    fmin: Optional[float] = None,
    fmax: Optional[float] = None,
    average: bool = True,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
    """Calculate Weighted Phase Lag Index between channels.

    See compute_connectivity for parameter details.
    """
    return compute_spectra_connectivity(data, sfreq, 'wpli', fmin, fmax, average=average)
