import logging
import os
from itertools import combinations

import matplotlib.pyplot as plt
import mne
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from mne_connectivity.spectral import spectral_connectivity_epochs
from mne_connectivity.viz import plot_connectivity_circle

__all__ = ['plot_connectivity_circle']

# ========================= 计算功能连接 ========================= #

# Define electrode coordinates (based on standard 10-20 system)
ELECTRODE_COORDS = {
    'P3': (-52.2, -60.0, 30.0),
    'C3': (-52.2, -16.4, 57.8),
    'F3': (-35.5, 49.4, 32.4),
    'Fz': (0.0, 50.0, 40.0),
    'F4': (40.2, 47.6, 32.1),
    'C4': (54.1, -18, 57.5),
    'P4': (52.2, -60.0, 30.0),
    'Cz': (0.0, 0.0, 60.0),
    'A1': (-80.0, -20.0, -20.0),
    'Fp1': (-30.0, 60.0, 10.0),
    'Fp2': (30.0, 60.0, 10.0),
    'T3': (-70.2, -21.3, -10.7),
    'T5': (-70.0, -50.0, -10.0),
    'O1': (-26.8, -100.2, 12.8),
    'O2': (24.1, -100.5, 12.8),
    'F7': (-60.0, 30.0, 10.0),
    'F8': (60.0, 30.0, 10.0),
    'A2': (80.0, -20.0, -20.0),
    'T6': (70.0, -50.0, -10.0),
    'T4': (71.9, -25.2, -8.2),
    'Pz': (0.0, -60.0, 30.0),
}


def calculate_connectivity(epochs, fmin, fmax):
    """Calculate PLI connectivity for given epochs."""
    con = spectral_connectivity_epochs(
        epochs,
        method='pli',
        mode='multitaper',
        sfreq=epochs.info['sfreq'],
        fmin=fmin,
        fmax=fmax,
        faverage=True,
        verbose=False,
        indices=None,
    )
    con_data = con.get_data()

    n_channels = len(epochs.info['ch_names'])
    con_data = reshape_to_full_matrix(con_data, n_channels)
    return con_data


def reshape_to_full_matrix(data, n_channels):
    """Reshape connectivity data to a full matrix."""
    full_matrix = np.zeros((n_channels, n_channels))
    triu_indices = np.array(list(combinations(range(n_channels), 2)))

    if data.ndim == 2:
        data = data[:, 0]

    if len(data) > len(triu_indices):
        data = data[: len(triu_indices)]
    elif len(data) < len(triu_indices):
        raise ValueError(
            f'Mismatch between data length ({len(data)}) and expected pairs ({len(triu_indices)}).'
        )

    for i, (row, col) in enumerate(triu_indices):
        full_matrix[row, col] = data[i]
        full_matrix[col, row] = data[i]

    return full_matrix


def plot_nilearn_connectivity(
    connectivity_matrix, ch_names, title='Connectivity Matrix', num_top_connections=10
):
    """
    Visualize connectivity matrix using nilearn's plot_connectome.

    Args:
        connectivity_matrix: The connectivity matrix to visualize
        ch_names: List of channel names
        title: Title for the plot
        num_top_connections: Number of top connections to show
    """
    from nilearn import plotting  # pyright: ignore[reportMissingImports]

    # Get coordinates for available channels
    node_coords = []
    valid_indices = []
    valid_names = []

    for idx, name in enumerate(ch_names):
        if name in ELECTRODE_COORDS:
            node_coords.append(ELECTRODE_COORDS[name])
            valid_indices.append(idx)
            valid_names.append(name)

    node_coords = np.array(node_coords)

    # Filter connectivity matrix for available electrodes
    filtered_matrix = connectivity_matrix[np.ix_(valid_indices, valid_indices)]

    # Create custom colormap
    custom_cmap = LinearSegmentedColormap.from_list(
        'custom_cmap', ['darkblue', 'white', 'darkred']
    )

    # Filter top connections
    abs_connectivity = np.abs(filtered_matrix)
    sorted_indices = np.argsort(abs_connectivity, axis=None)[::-1]
    if 0 < num_top_connections < len(sorted_indices):
        top_indices = sorted_indices[:num_top_connections]
    else:
        top_indices = sorted_indices

    # Create matrix with only top connections
    top_matrix = np.zeros_like(filtered_matrix)
    for idx in top_indices:
        i, j = np.unravel_index(idx, filtered_matrix.shape)
        top_matrix[i, j] = filtered_matrix[i, j]

    # Plot connectome
    plotting.plot_connectome(
        top_matrix,
        node_coords,
        node_color='red',
        node_size=50,
        edge_cmap=custom_cmap,
        edge_vmin=-0.1,
        edge_vmax=0.1,
        edge_kwargs={'linewidth': 2.0},
        edge_threshold=None,
        title=f'{title}\nTop {num_top_connections} Connections',
        display_mode='ortho',
        colorbar=True,
    )

    return plt.gcf()
