import matplotlib
matplotlib.use('Agg')
import os
import io
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import matplotlib.cm as cm
import matplotlib.colors as mcolors
from scipy.spatial import Delaunay
from scipy.interpolate import Rbf
from scipy.signal import welch
from scipy.stats import sem
from scipy.signal import butter, filtfilt, hilbert
import base64

# 位置坐标
def read_ch_pos(file_path=r'xy_64.xlsx'):
    """
    将电极位置信息转换为Dict

    参数:
    file_path: 电极位置存储文件, 必须包含'channel', 'x', 'y', 'z'列

    """
    script_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(script_dir,file_path )
    df = pd.read_excel(file_path)
    # 确保列名正确
    if not all(col in df.columns for col in ['channel', 'x', 'y', 'z']):
        raise ValueError("DataFrame必须包含'channel', 'x', 'y', 'z'列")
    # 创建电极位置字典
    ch_pos = {}
    for _, row in df.iterrows():
        ch_pos[row['channel']] = [row['x'], row['y'], row['z']]
    return ch_pos
# 头部轮廓
def draw_head(ax, center=(0, 0), radius=1.0, zorder=4):
    """
    绘制头部轮廓、鼻子和耳朵。

    参数:
    - ax : matplotlib Axes 对象
    - center : (x, y) 头中心坐标
    - radius : float, 头半径
    - zorder : 绘制层级
    """

    # 头圆
    head = plt.Circle(center, radius, fill=False, color='k', linewidth=1, zorder=zorder)
    ax.add_artist(head)

    # 鼻子（参考 _make_head_outlines）
    dx = np.exp(np.arccos(np.deg2rad(12)) * 1j)
    dx_real, dx_imag = dx.real, dx.imag
    nose_x = np.array([-dx_real, 0, dx_real]) * radius + center[0]
    nose_y = np.array([dx_imag, 1.15, dx_imag]) * radius + center[1]
    ax.plot(nose_x, nose_y, color='k', linewidth=1, zorder=zorder)

    # 耳朵（参考 _make_head_outlines 手动标定）
    ear_radius = radius * 0.12
    ear_scale = radius * 2  # 根据半径缩放
    theta = np.linspace(np.pi / 2, 3 * np.pi / 2, 30)

    # 左耳
    left_ear_x_array = np.array([0.497, 0.510, 0.518, 0.5299, 0.5419,
                                 0.54, 0.547, 0.532, 0.510, 0.489]) * ear_scale
    left_ear_y_array = np.array([0.0555, 0.0775, 0.0783, 0.0746, 0.0555,
                                 -0.0055, -0.0932, -0.1313, -0.1384, -0.1199]) * ear_scale + center[1]
    ax.plot(center[0] - left_ear_x_array, left_ear_y_array, color='k', linewidth=1, zorder=zorder)

    # 右耳
    right_ear_x_array = np.array([0.497, 0.510, 0.518, 0.5299, 0.5419,
                                  0.54, 0.547, 0.532, 0.510, 0.489]) * ear_scale
    right_ear_y_array = np.array([0.0555, 0.0775, 0.0783, 0.0746, 0.0555,
                                  -0.0055, -0.0932, -0.1313, -0.1384, -0.1199]) * ear_scale + center[1]
    ax.plot(center[0] + right_ear_x_array, right_ear_y_array, color='k', linewidth=1, zorder=zorder)
# 地形图 插值
def rbf_D_interpolate(xy, v, center=(0, 0), radius=1.1, grid_res=300,
                      n_extra=32, rbf_func='multiquadric', smooth=0,
                      border='mean', border_scale=1.0001, n_ngb=4):
    """
    使用 RBF + Delaunay 邻域均值方式生成平滑的 EEG topomap 插值表面。

    参数
    ----
    xy : (N,2) array
        电极二维坐标（与绘图坐标系一致）
    v : (N,) array
        每个电极对应的值（e.g. PSD）
    center : tuple (x0, y0)
        头部圆心（默认 (0,0)）
    radius : float
        头部半径（用于生成边界点与网格范围）
    grid_res : int
        网格分辨率（每轴点数）
    n_extra : int
        边界虚拟点数量
    rbf_func : str
        RBF 内核名称（'multiquadric','thin_plate','gaussian',...）
    smooth : float
        RBF 平滑参数
    border : 'mean' or float
        若 'mean'：边界点用邻近真实通道均值赋值（推荐）
        若 float：边界点赋相同常数值
    border_scale : float
        边界点半径相对 radius 的缩放（略微 >1 用以外推）
    n_ngb : int
        为每个边界点取值时使用的最近真实通道数

    返回
    ----
    zi : (grid_res, grid_res) ndarray
        插值结果（与 grid_x, grid_y 对齐）
    grid_x, grid_y : ndarrays
        meshgrid（由 np.meshgrid 生成）
    """
    xy = np.asarray(xy)
    v = np.asarray(v)
    if xy.ndim != 2 or xy.shape[1] != 2:
        raise ValueError("xy must be shape (n_channels, 2)")

    n_points = xy.shape[0]

    # --- 1. 生成边界虚拟点（圆周） ---
    theta = np.linspace(0.0, 2 * np.pi, n_extra, endpoint=False)
    r_border = radius * border_scale
    border_xy = np.column_stack([center[0] + r_border * np.cos(theta),
                                 center[1] + r_border * np.sin(theta)])

    # --- 2. 用 Delaunay 建图以便找到邻居（对边界点取邻居均值） ---
    # 合并用于三角化的位置（真实点 + 边界点）
    tri_xy = np.vstack([xy, border_xy])
    tri = Delaunay(tri_xy)

    # --- 3. 为边界点赋值 ---
    if isinstance(border, str) and border == 'mean':
        # 使用 Delaunay 的 vertex_neighbor_vertices 索引
        # 注意：tri.vertex_neighbor_vertices 给出 vertices -> neighbor indptr
        indices, indptr = tri.vertex_neighbor_vertices
        v_extra = np.zeros(n_extra)
        used = np.zeros(n_extra, dtype=bool)
        # 边界点在 tri_xy 中的索引范围
        rng = range(n_points, n_points + n_extra)
        for idx, extra_idx in enumerate(rng):
            neigh = indptr[indices[extra_idx]:indices[extra_idx + 1]]
            # 仅保留原始点索引（小于 n_points）
            neigh = neigh[neigh < n_points]
            if neigh.size > 0:
                used[idx] = True
                # 使用最近 n_ngb 个邻居的均值（若邻居多则取最近的 n_ngb）
                if neigh.size > n_ngb:
                    # 计算距离并选取最近 n_ngb
                    d = np.linalg.norm(xy[neigh] - tri_xy[extra_idx], axis=1)
                    order = np.argsort(d)[:n_ngb]
                    sel = neigh[order]
                else:
                    sel = neigh
                v_extra[idx] = v[sel].mean()
        if not used.all() and used.any():
            v_extra[~used] = np.mean(v_extra[used])
        elif not used.any():
            v_extra[:] = np.mean(v)
    else:
        # border 是数值
        v_extra = np.full(n_extra, float(border))

    # --- 4. 合并所有已知点并构建 RBF ---
    all_xy = np.vstack([xy, border_xy])
    all_v = np.concatenate([v, v_extra])

    rbf = Rbf(all_xy[:, 0], all_xy[:, 1], all_v,
              function=rbf_func, smooth=smooth)

    # --- 5. 生成网格（使用 meshgrid，与主函数保持一致） ---
    xmin, xmax = center[0] - radius, center[0] + radius
    ymin, ymax = center[1] - radius, center[1] + radius
    xi = np.linspace(xmin, xmax, grid_res)
    yi = np.linspace(ymin, ymax, grid_res)
    grid_x, grid_y = np.meshgrid(xi, yi)  # meshgrid 与 imshow 对齐

    # --- 6. 评估 RBF，返回与 grid 对齐的 zi ---
    zi = rbf(grid_x, grid_y)

    return zi, grid_x, grid_y
# plv矩阵计算
def calculate_plv(data):
    """
    计算相位锁定值（PLV）矩阵。

    Parameters
    ----------
    data : ndarray, shape (num_channels, num_samples)
        EEG 数据，通道数为 num_channels，样本数为 num_samples。

    Returns
    -------
    plv_matrix : ndarray, shape (num_channels, num_channels)
        计算得到的 PLV 矩阵，表示各通道间的相位同步。
    """
    num_channels, num_samples = data.shape
    plv_matrix = np.zeros((num_channels, num_channels))

    # 计算每个通道的解析信号
    analytic_signals = np.apply_along_axis(hilbert, axis=1, arr=data)

    for i in range(num_channels):
        for j in range(i + 1, num_channels):  # 只计算上三角矩阵，避免重复计算
            # 计算 phase difference
            phase_diff = np.angle(analytic_signals[i] * np.conj(analytic_signals[j]))
            plv = np.abs(np.mean(np.exp(1j * phase_diff)))
            plv_matrix[i, j] = plv
            plv_matrix[j, i] = plv  # 对称矩阵

    return plv_matrix
# 矩阵阈值化
def threshold_proportional(adj, prop=0.2):
    """
    Apply a proportional threshold to retain the top proportion of strongest edges.

    Parameters
    ----------
    adj : ndarray, shape (n_channels, n_channels)
        Adjacency matrix to threshold.
    prop : float
        Proportion of edges to retain (0 < prop <= 1).

    Returns
    -------
    bin_adj : ndarray, shape (n_channels, n_channels)
        Binary adjacency matrix after thresholding.
    """
    n = adj.shape[0]
    triu_idx = np.triu_indices(n, k=1)
    weights = adj[triu_idx]
    k = int(np.floor(len(weights) * prop))

    # Ensure that at least one edge is retained
    k = max(k, 1)

    # Get the threshold value
    thr = np.sort(weights)[-k]

    # Apply the threshold to create a binary adjacency matrix
    bin_adj = np.where(adj >= thr, adj, 0.0)

    return bin_adj
# 单个脑网络
def plot_single_network(ch_names,adj,ax=None,
                        node_size=20, node_color='orange',highlight_nodes=[], show_names=True,
                        edge_color='gray', weighted=True,
                        radius=1.1, figsize=(6, 6),cmap='RdYlBu_r'):
    # 若 ax 未传入，则自己创建
    own_fig = False
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
        own_fig = True
    else:
        fig = ax.figure

    # 坐标归一化
    pos3d = read_ch_pos()
    all_chs_xy = np.array([pos3d[ch][:2] for ch in pos3d.keys()])
    all_chs_xy -= all_chs_xy.mean(axis=0)
    all_chs_xy /= np.sqrt((all_chs_xy ** 2).sum(axis=1)).max()
    xy_dict = dict(zip(pos3d.keys(), all_chs_xy))
    xy = np.array([xy_dict[ch] for ch in ch_names])
    center = xy_dict.get('CZ', np.mean(list(xy_dict.values()), axis=0))

    # ===== 初始化绘图窗口 =====
    ax.set_aspect('equal')
    ax.axis('off')
    # 设置边界（与原类保持一致）
    ear_radius = radius * 0.12
    nose_height = radius * 0.15
    margin_x = radius * 0.12 + 0.05
    ax.set_xlim(center[0] - radius - margin_x, center[0] + radius + margin_x)
    ax.set_ylim(center[1] - radius - ear_radius, center[1] + radius + nose_height + ear_radius)

    # 绘制头部轮廓
    draw_head(ax, center=center, radius=radius)

    # 节点
    for ch in ch_names:
        color = 'red' if ch in highlight_nodes else node_color
        ax.scatter(*xy_dict[ch], s=node_size, color=color, edgecolor='k', zorder=4)
        if show_names:
            ax.text(xy_dict[ch][0], xy_dict[ch][1] + 0.03, ch,
                               ha='center', va='bottom', fontsize=8, zorder=5)

    # colorbar
    norm = mcolors.Normalize(vmin=0, vmax=1)
    color_map = matplotlib.colormaps.get_cmap(cmap)
    # ========= 边 ==========
    N = len(ch_names)
    for i in range(N):
        for j in range(i + 1, N):
            w = adj[i, j]
            if w > 0:
                x = [xy[i, 0], xy[j, 0]]
                y = [xy[i, 1], xy[j, 1]]
                lw = 1.5
                if weighted:
                    ax.plot(x, y,
                        color=color_map(norm(w)),
                        linewidth=lw,
                        alpha=0.7,
                        zorder=3)
                else:
                    ax.plot(x, y,
                        color=edge_color,
                        linewidth=lw,
                        alpha=0.7,
                        zorder=3)

    if own_fig:
        # 不回传 添加颜色条
        sm = cm.ScalarMappable(norm=norm, cmap=color_map)
        cbar = plt.colorbar(sm, ax=ax, fraction=0.035)
        cbar.set_label('Connection Strength', fontsize=10)
        cbar.ax.tick_params(direction='in', labelsize=10)
        plt.show()
        return fig
    else:

        return ax
# 脑网络对比
def plot_multiband_network(ch_names, adj_MI, adj_Rest,cmap='RdYlBu_r'):

    fig, axes = plt.subplots(1, 2, figsize=(8, 4))
    fontsize = 16
    fig.text(0.285, 0.08, 'MI', fontsize=fontsize,  ha='center', va='center', rotation=0)
    fig.text(0.68, 0.08, 'Rest', fontsize=fontsize,  ha='center', va='center', rotation=0)

    im1 = plot_single_network(ch_names,adj_MI,ax=axes[0], show_names=True,cmap=cmap)
    # Rest 行
    im2 = plot_single_network(ch_names,adj_Rest,ax=axes[1],show_names=True,cmap=cmap)

    # --- 合并 colorbar（右侧一个） ---
    norm = mcolors.Normalize(vmin=0, vmax=1)
    color_map = matplotlib.colormaps.get_cmap(cmap)
    sm = cm.ScalarMappable(norm=norm, cmap=color_map)
    cbar = plt.colorbar(sm, ax=axes.ravel().tolist(), fraction=0.02)
    cbar.set_label('Connection Strength', fontsize=10)
    cbar.ax.tick_params(direction='in', labelsize=10)

    # 将图像保存到内存字节流（PNG 格式）
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close(fig)  # 释放内存
    buf.seek(0)
    image_bytes = buf.read()
    buf.close()

    return image_bytes

# 多个频带psd
def compute_band_psd(eeg, fs, bands, labels, trial_idx=0,MI_label=1, Rest_label=2,avg = True):
    """
    eeg: (n_trials, n_channels, n_samples)
    """
    n_trials, n_channels, n_samples = eeg.shape
    band_names = list(bands.keys())
    n_bands = len(band_names)

    psd_MI = np.zeros((n_bands, n_channels))
    psd_Rest = np.zeros((n_bands, n_channels))

    # 先计算所有 trial 的功率谱
    f, Pxx = welch(eeg, fs=fs, axis=-1, nperseg=fs,noverlap = fs // 2)


    for bi, (bname, (f1, f2)) in enumerate(bands.items()):
        idx = np.logical_and(f >= f1, f <= f2)
        band_power = Pxx[:, :, idx].mean(axis=-1)

        band_power_flat = band_power.flatten()
        power_min = band_power_flat.min()
        power_max = band_power_flat.max()
        if power_max - power_min > 1e-12:
            band_power_norm = (band_power - power_min) / (power_max - power_min)
        else:
            band_power_norm = band_power

        if avg:
            psd_MI[bi] = band_power_norm[labels == MI_label].mean(axis=0)
            psd_Rest[bi] = band_power_norm[labels == Rest_label].mean(axis=0)
        else:
            psd_MI[bi] = band_power_norm[labels == MI_label][trial_idx]
            psd_Rest[bi] = band_power_norm[labels == Rest_label][trial_idx]
    return band_names, psd_MI, psd_Rest
# 单个脑地形图
def plot_single_topomap(ch_names, psd_values, cmap='RdYlBu_r', vlim=(0, 1),
                        show_names=True, node_size=3, radius=1.1, grid_res=300,
                        n_contours=None, contour_color='k',
                        ax=None,figsize=(6,6)):
    # 若 ax 未传入，则自己创建
    own_fig = False
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
        own_fig = True
    else:
        fig = ax.figure

    # ===== 初始化绘图窗口 =====
    ax.set_aspect('equal')
    ax.axis('off')
    # ax.set_title("EEG topomap (MNE-like)")

    # 坐标归一化
    pos3d = read_ch_pos()
    all_chs_xy = np.array([pos3d[ch][:2] for ch in pos3d.keys()])
    all_chs_xy -= all_chs_xy.mean(axis=0)
    all_chs_xy /= np.sqrt((all_chs_xy ** 2).sum(axis=1)).max()
    pos2d_dict = dict(zip(pos3d.keys(), all_chs_xy))
    xy = np.array([pos2d_dict[ch] for ch in ch_names])
    center = pos2d_dict.get('CZ', np.mean(list(pos2d_dict.values()), axis=0))

    # 绘制头部轮廓
    draw_head(ax, center=center, radius=radius)
    # 绘制电极
    fontsize = 4
    ax.scatter(xy[:, 0], xy[:, 1], c='k', s=node_size, zorder=5)
    if show_names:
        for i, ch in enumerate(ch_names):
            ax.text(xy[i, 0], xy[i, 1] + 0.03, ch,
                    ha='center', va='bottom', fontsize=fontsize, zorder=6)

    # 数据插值
    zi, grid_x, grid_y = rbf_D_interpolate(
        xy, psd_values, radius=radius,
        grid_res=grid_res
    )
    xmin, xmax = center[0] - radius, center[0] + radius
    ymin, ymax = center[1] - radius, center[1] + radius
    extent = (xmin, xmax, ymin, ymax)
    im = ax.imshow(zi, extent=extent, origin='lower',
                   cmap=cmap, vmin=vlim[0], vmax=vlim[1],
                   interpolation='bicubic', zorder=0)
    # 裁剪路径
    patch_ = Ellipse(center, 2 * radius, 2 * radius, clip_on=True, transform=ax.transData)
    im.set_clip_path(patch_)
    # 初始等高线
    linewidths = 0.5
    if n_contours is None:
        cset = ax.contour(grid_x, grid_y, zi,
                          colors=contour_color, linewidths=linewidths, zorder=2)
    else:
        cset = ax.contour(grid_x, grid_y, zi, levels=n_contours,
                          colors=contour_color, linewidths=linewidths, zorder=2)
    cset.set_clip_path(patch_)



    if own_fig:
        # 不回传 添加颜色条
        plt.colorbar(im, ax=ax, fraction=0.035)
        plt.show()
        return fig
    else:
        # plt.colorbar(im, ax=ax, fraction=0.035)
        return im
# 脑地形图对比
def plot_multiband_topomaps(ch_names, psd_MI, psd_Rest, bands):
    band_names = list(bands.keys())  # 改动 1：新增这行
    n_bands = len(band_names)
    fig, axes = plt.subplots(2, n_bands, figsize=(3*n_bands, 6))

    fontsize = 16

    axes[0, 0].text(-0.1, 0.5, 'MI', transform=axes[0, 0].transAxes, rotation=0, va='center', ha='center', fontsize=fontsize-2)
    axes[1, 0].text(-0.1, 0.5, 'Rest', transform=axes[1, 0].transAxes, rotation=0, va='center', ha='center', fontsize=fontsize-2)

    imgs = []
    for i, bname in enumerate(band_names):
        axes[0, i].set_title(bname, fontsize=fontsize, pad=0)
        # MI 行
        im1 = plot_single_topomap(ch_names,psd_MI[i],ax=axes[0, i], show_names=True)
        # Rest 行
        im2 = plot_single_topomap(ch_names,psd_Rest[i],ax=axes[1, i],show_names=True)
        imgs.append(im1)

    # --- 单个右侧合并 colorbar ---
    cbar = fig.colorbar(imgs[0], ax=axes,fraction=0.02)
    # cbar.set_label("PSD Power",fontsize=fontsize-4)
    cbar.ax.tick_params(direction='in', labelsize=10)

    # 将图像保存到内存字节流（PNG 格式）
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close(fig)  # 释放内存
    buf.seek(0)
    image_bytes = buf.read()
    buf.close()

    return image_bytes

# 小波
def morlet_wavelet(f, fs, n_cycles=7):
    """
    创建 Morlet 小波
    f: 频率
    fs: 采样率
    """
    sigma_t = n_cycles / (2 * np.pi * f)
    t = np.arange(-3*sigma_t, 3*sigma_t, 1/fs)
    wavelet = (np.pi**-0.25) * np.exp(2j*np.pi*f*t) * np.exp(-(t**2)/(2*sigma_t**2))
    return wavelet


# 希尔伯特变换 计算ERDS 效果不佳
def bandpass_filter(data, fs, band, order=4):
    nyq = fs / 2
    b, a = butter(order, [band[0]/nyq, band[1]/nyq], btype='band')
    return filtfilt(b, a, data, axis=-1)
def compute_power_hilbert(filtered_data,is_dB =True):
    analytic = hilbert(filtered_data, axis=-1)
    power = np.abs(analytic) ** 2
    if is_dB:
        power = 10 * np.log10(power)
    return power
def compute_power(data, fs=250,
                  bands={"mu": (8,12), "beta": (13,30)}):
    """
    返回：
        power_dict[band] = (n_trials, n_ch, n_samples)
    """
    power_dict = {}
    for band_name, band_range in bands.items():
        filt = bandpass_filter(data, fs, band_range)
        power = compute_power_hilbert(filt)
        power_dict[band_name] = power

    return power_dict

def compute_erds(power_MI, power_Rest, baseline_period=None):
    """
    计算事件相关去同步/同步 (ERDS)

    Parameters
    ----------
    power_MI, power_Rest: (n_trials, n_ch, n_samples)
        功率数据，单位为 µV² 或 dB（取决于 compute_power_hilbert 的 is_dB 参数）
    baseline_period: tuple (start_idx, end_idx) or None
        基线时间段索引。如果为None，使用 Rest 状态的平均值作为基线

    返回：
        MI_erds_mean, MI_erds_sem
        Rest_erds_mean, Rest_erds_sem
        所有返回值的形状为 (n_ch, n_samples)
    """

    if baseline_period is not None:
        start_idx, end_idx = baseline_period
        baseline = np.concatenate([power_MI[:, :, start_idx:end_idx],
                                  power_Rest[:, :, start_idx:end_idx]], axis=0)
        baseline = baseline.mean(axis=(0, 2), keepdims=True)
    else:
        baseline = power_Rest.mean(axis=(0,2), keepdims=True)

    # === ERDS (%) ===
    MI_erds  = (power_MI  - baseline) / baseline * 100
    Rest_erds = (power_Rest - baseline) / baseline * 100

    return (
        MI_erds.mean(axis=0),  sem(MI_erds, axis=0),
        Rest_erds.mean(axis=0), sem(Rest_erds, axis=0),
    )

def compute_all_erds(MI_power_dict, Rest_power_dict):
    """
    对多个频带同时计算 ERDS。

    输入:
        MI_power_dict[band]   = (n_trials, n_ch, n_samples)
        Rest_power_dict[band] = (n_trials, n_ch, n_samples)

    输出:
        erds_MI[band]   = (mean, sem)
        erds_Rest[band] = (mean, sem)
    """

    erds_MI = {}
    erds_Rest = {}

    for band in MI_power_dict.keys():
        MI_power = MI_power_dict[band]
        Rest_power = Rest_power_dict[band]

        MI_mean, MI_sem, Rest_mean, Rest_sem = compute_erds(MI_power, Rest_power)

        erds_MI[band] = (MI_mean, MI_sem)
        erds_Rest[band] = (Rest_mean, Rest_sem)

    return erds_MI, erds_Rest

def plot_compare_erds(data_MI, data_Rest, mode="power",
                      ch_names = ['FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'CP3', 'CP1', 'CP2', 'CP4', 'P3', 'P1', 'PZ', 'P2', 'P4'],
                      compare_names=['C3', 'CZ', 'C4'], bands=['mu', 'beta'],
                      fs=250, t=None, figsize=(12,6)):

    n_bands = len(bands)
    n_chs = len(compare_names)

    # 自动添加单位
    if mode == "power":
        # y_unit = "Power (µV²)"
        y_unit = "Power (dB)"
    elif mode == "erds":
        y_unit = "ERDS (%)"
    else:
        y_unit = ""

    if t is None:
        n_samples = next(iter(data_MI.values())).shape[-1] \
            if mode=="power" else next(iter(data_MI.values()))[0].shape[-1]
        t = np.arange(n_samples) / fs

    fig, axes = plt.subplots(n_bands, n_chs, figsize=figsize, sharex=True, sharey=True)

    for i, band in enumerate(bands):

        # 选择数据结构
        if mode == "power":
            MI_band = data_MI[band]      # (trials, ch, samples)
            Rest_band = data_Rest[band]

            avg_MI = MI_band.mean(axis=0)
            sem_MI = MI_band.std(axis=0)/np.sqrt(MI_band.shape[0])

            avg_Rest = Rest_band.mean(axis=0)
            sem_Rest = Rest_band.std(axis=0)/np.sqrt(Rest_band.shape[0])

        elif mode == "erds":
            avg_MI, sem_MI = data_MI[band]
            avg_Rest, sem_Rest = data_Rest[band]

        for j, ch in enumerate(compare_names):
            ax = axes[i, j] if n_bands > 1 else axes[j]

            ch_idx = ch_names.index(ch)

            # 绘制 MI
            ax.plot(t, avg_MI[ch_idx], color="C0", label="MI")
            ax.fill_between(t,
                            avg_MI[ch_idx]-sem_MI[ch_idx],
                            avg_MI[ch_idx]+sem_MI[ch_idx],
                            alpha=0.3, color="C0")

            # 绘制 Rest
            ax.plot(t, avg_Rest[ch_idx], color="C1", label="Rest")
            ax.fill_between(t,
                            avg_Rest[ch_idx]-sem_Rest[ch_idx],
                            avg_Rest[ch_idx]+sem_Rest[ch_idx],
                            alpha=0.3, color="C1")

            if i == 0:
                ax.set_title(ch)

            # ← Y 轴加单位
            if j == 0:
                ax.set_ylabel(f"{band}\n{y_unit}")

            if i == n_bands - 1:
                ax.set_xlabel("Time (s)")

            ax.grid(alpha=0.3)

            if i == 0 and j == n_chs - 1:
                ax.legend()

    plt.tight_layout()

    # 将图像保存到内存字节流（PNG 格式）
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close(fig)  # 释放内存
    buf.seek(0)
    image_bytes = buf.read()
    buf.close()

    return image_bytes

# 对比 MI vs Rest 的功率谱密度 PSD
def plot_psd_compare(MI_data, Rest_data, ch_names, compare_names=['C3', 'CZ', 'C4'],
                     fs=250, nperseg=None, average=True, show_sem=True,
                     figsize=(12, 3), save_dir=None, filename="psd.png"):
    """
    对比 MI vs Rest 的功率谱密度 PSD

    MI_data, Rest_data: (n_trials, n_ch, n_samples)
    channels: 需要绘制的通道
    average: 是否对所有试次平均
    show_sem: 是否绘制 SEM 阴影
    """

    n_trials, n_ch, n_samples = MI_data.shape
    n_trials = min(len(MI_data), len(Rest_data))
    # assert Rest_data.shape == MI_data.shape, "MI 和 Rest 数据维度必须一致"

    if nperseg is None:
        nperseg = fs  # 每 1 秒窗长度

    # 计算 MI PSD
    psd_MI_all = []
    for trial in range(n_trials):
        psd_trial = []
        for ch in range(n_ch):
            f, Pxx = welch(MI_data[trial, ch], fs=fs, nperseg=nperseg)
            psd_trial.append(Pxx)
        psd_MI_all.append(psd_trial)
    psd_MI_all = np.array(psd_MI_all)

    # 计算 Rest PSD
    psd_Rest_all = []
    for trial in range(n_trials):
        psd_trial = []
        for ch in range(n_ch):
            _, Pxx = welch(Rest_data[trial, ch], fs=fs, nperseg=nperseg)
            psd_trial.append(Pxx)
        psd_Rest_all.append(psd_trial)
    psd_Rest_all = np.array(psd_Rest_all)

    # ---- Plot ----
    fig, ax = plt.subplots(1, len(compare_names), figsize=figsize)
    if len(compare_names) == 1:
        ax = [ax]

    for i, ch in enumerate(compare_names):
        ch_idx = ch_names.index(ch)
        psd_MI_ch = psd_MI_all[:, ch_idx, :]
        psd_Rest_ch = psd_Rest_all[:, ch_idx, :]

        if average:
            mean_MI = psd_MI_ch.mean(axis=0)
            mean_Rest = psd_Rest_ch.mean(axis=0)

            ax[i].plot(f, mean_MI, color='C0', label='MI')
            ax[i].plot(f, mean_Rest, color='C1', label='Rest')

            if show_sem:
                ax[i].fill_between(f, mean_MI - sem(psd_MI_ch, axis=0),
                                   mean_MI + sem(psd_MI_ch, axis=0), color='C0', alpha=0.3)
                ax[i].fill_between(f, mean_Rest - sem(psd_Rest_ch, axis=0),
                                   mean_Rest + sem(psd_Rest_ch, axis=0), color='C1', alpha=0.3)
        else:
            ax[i].plot(f, psd_MI_ch.T, color='C0', alpha=0.3)
            ax[i].plot(f, psd_Rest_ch.T, color='C1', alpha=0.3)

        ax[i].set_title(ch)
        ax[i].set_xlabel("Frequency (Hz)")
        ax[i].set_ylabel("PSD (μV²/Hz)")
        ax[i].grid(alpha=0.3)
        if i == 0:
            ax[i].legend()

    plt.tight_layout()

    # 将图像保存到内存字节流（PNG 格式）
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=300, bbox_inches='tight')
    plt.close(fig)  # 释放内存
    buf.seek(0)
    image_bytes = buf.read()
    buf.close()

    return image_bytes


def plotMain(
        ch_names = ['FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'CP3', 'CP1', 'CP2', 'CP4', 'P3', 'P1', 'PZ', 'P2', 'P4'],
        compare_names = [ 'C3','CZ','C4'],
        Data = None,labels = None,MI_label = None,Rest_label = None,
        fs = 250):

    trial_idx = 0

    # 数据划分
    if not MI_label:
        label_ = np.unique(labels)
    else:
        label_ = (MI_label,Rest_label)
    MI_data = Data[labels == label_[0]]
    Rest_data = Data[labels == label_[1]]

    # 典型 EEG 频带
    FREQ_BANDS = {
        "Delta (0.8-4Hz)": (0.8, 4),
        "Theta (4-8Hz)": (4, 8),
        "Alpha (8-12Hz)": (8, 12),
        "Beta (12-30Hz)": (12, 30),
        "All (0.8-30Hz)": (0.8, 30)
    }
    # 利用welch估算PSD
    band_names, psd_MI, psd_Rest= compute_band_psd(
        eeg=Data,
        fs=fs,
        bands=FREQ_BANDS,
        labels=labels,
        trial_idx=trial_idx,
        MI_label=MI_label,
        Rest_label=Rest_label,
        avg= True
    )
    # 绘制地形图
    topomaps_imgBytes = plot_multiband_topomaps(
        ch_names=ch_names,
        psd_MI=psd_MI,
        psd_Rest=psd_Rest,
        bands=FREQ_BANDS
    )

    # 绘制脑网络
    mi_plv_matrix = calculate_plv(MI_data[trial_idx])
    mi_BI_matrix = threshold_proportional(mi_plv_matrix, prop=0.3)
    rest_plv_matrix = calculate_plv(Rest_data[trial_idx])
    rest_BI_matrix = threshold_proportional(rest_plv_matrix, prop=0.3)
    network_imgBytes = plot_multiband_network(ch_names, mi_BI_matrix, rest_BI_matrix)

    # ERDS 先计算erds，后平均
    MI_power = compute_power(MI_data)
    Rest_power = compute_power(Rest_data)
    erds_dict_MI, erds_dict_Rest = compute_all_erds(MI_power, Rest_power)
    erds_imgBytes = plot_compare_erds(erds_dict_MI, erds_dict_Rest, ch_names=ch_names,
                      compare_names=compare_names, bands=['mu', 'beta'],
                      fs=fs, mode="erds")

    # 绘制PSD
    psd_imgBytes = plot_psd_compare(MI_data, Rest_data, ch_names = ch_names, compare_names=compare_names,
                     fs=fs, nperseg=None, average=True, show_sem=True,
                     figsize=(12, 3))
    return {'topomaps_imgBytes':base64.b64encode(topomaps_imgBytes).decode(),'network_imgBytes':base64.b64encode(network_imgBytes).decode(),
            'erds_imgBytes':base64.b64encode(erds_imgBytes).decode(),'psd_imgBytes':base64.b64encode(psd_imgBytes).decode()}

if __name__ == '__main__':
    allData = np.random.uniform(-50,50,size=(80,21,1000))
    allLabel = np.random.randint(1,3,size=(80,))
    allData = allData[:len(allLabel)]
    ch_names = ['FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'CP3', 'CP1',
                'CP2', 'CP4', 'P3', 'P1', 'PZ', 'P2', 'P4']
    compare_names = ['C3', 'CZ', 'C4']
    ret = plotMain(ch_names=ch_names, compare_names=compare_names, Data=allData, labels=allLabel, MI_label=1, Rest_label=2,
             fs=250)
    print('计算完成，开始发送')
    from Zmq.zmqClient import zmqClient

    zmqClient = zmqClient('192.168.76.101', 8088)
    zmqClient.connect()
    zmqClient.send_to_all('miReport', ret)
