# plot EEG topograpy with mne 
# https://mne.tools/stable/index.html

import mne
import numpy as np
import scipy.io
import matplotlib.pyplot as plt
from matplotlib import mlab as mlab
# from preprocess import import_data

data = 0
label = 0
# get the data and label 
# data - (samples, channels, trials)
# label -  (label, 1)

data = np.transpose(data, (2, 1, 0))
label = np.squeeze(np.transpose(label))
idx = np.where(label == 1)
data_draw = data[idx]

mean_trial = np.mean(data_draw, axis=0)  # mean trial
# use standardization or normalization to adjust
mean_trial = (mean_trial - np.mean(mean_trial)) / np.std(mean_trial)


mean_ch = np.mean(mean_trial, axis=1)  # mean samples with channel dimension left


# Draw topography
biosemi_montage = mne.channels.make_standard_montage('biosemi64')  # set a montage, see mne document
index = [37, 9, 10, 46, 45, 44, 13, 12, 11, 47, 48, 49, 50, 17, 18, 31, 55, 54, 19, 30, 56, 29]  # correspond channel
biosemi_montage.ch_names = [biosemi_montage.ch_names[i] for i in index]
biosemi_montage.dig = [biosemi_montage.dig[i+3] for i in index]
info = mne.create_info(ch_names=biosemi_montage.ch_names, sfreq=250., ch_types='eeg')  # sample rate

evoked1 = mne.EvokedArray(mean_trial, info)
evoked1.set_montage(biosemi_montage)
plt.figure(1)
# im, cn = mne.viz.plot_topomap(np.mean(mean_trial, axis=1), evoked1.info, show=False)
im, cn = mne.viz.plot_topomap(mean_ch, evoked1.info, show=False)
plt.colorbar(im)

plt.savefig('./topo/test.png')
print('the end')