import mne
import os
import os.path as op
import numpy as np
from mne.datasets import sample, fetch_fsaverage
from mne.minimum_norm import apply_inverse, read_inverse_operator, make_inverse_operator

data_path = mne.datasets.sample.data_path()
subject = 'sample'
data_dir = op.join(data_path, 'MEG', subject)
subjects_dir = op.join(data_path, 'subjects')
bem_dir = op.join(subjects_dir, subject, 'bem')

fname_mixed_src = op.join(bem_dir, '%s-oct-6-mixed-src.fif' % subject)
fname_aseg = op.join(subjects_dir, subject, 'mri', 'aseg.mgz')
fname_model = op.join(bem_dir, '%s-5120-bem.fif' % subject)
fname_bem = op.join(bem_dir, '%s-5120-bem-sol.fif' % subject)
fname_evoked = data_dir + '/sample_audvis-ave.fif'
fname_trans = data_dir + '/sample_audvis_raw-trans.fif'
fname_fwd = data_dir + '/sample_audvis-meg-oct-6-mixed-fwd.fif'
fname_cov = data_dir + '/sample_audvis-shrunk-cov.fif'

#create a mixed source space
labels_vol = ['Left-Cerebellum-Cortex',
              'Right-Cerebellum-Cortex']
src = mne.setup_source_space(subject, spacing='oct5',
                             add_dist=False, subjects_dir=subjects_dir)

vol_src = mne.setup_volume_source_space(
    subject, mri=fname_aseg, pos=10.0, bem=fname_model,
    volume_label=labels_vol, subjects_dir=subjects_dir,
    add_interpolator=False,  # just for speed, usually this should be True
    verbose=True)

#Generate the mixed source space
src += vol_src

fwd = mne.make_forward_solution(
    fname_evoked, fname_trans, src, fname_bem,
    mindist=5.0,  # ignore sources<=5mm from innerskull
    meg=True, eeg=False, n_jobs=1)

leadfield = fwd['sol']['data']
src_fwd = fwd['src']
n = sum(src_fwd[i]['nuse'] for i in range(len(src_fwd)))

condition = 'Left Auditory'
evoked = mne.read_evokeds(fname_evoked, condition=condition,
                          baseline=(None, 0))
noise_cov = mne.read_cov(fname_cov)

snr = 3.0            # use smaller SNR for raw data
inv_method = 'MNE'  # sLORETA, MNE, dSPM
#parc = 'aparc'       # the parcellation to use, e.g., 'aparc' 'aparc.a2009s'
lambda2 = 1.0 / snr ** 2
inverse_operator = make_inverse_operator(evoked.info, fwd, noise_cov,
                                         depth=None, fixed=False)

stc = apply_inverse(evoked, inverse_operator, lambda2, inv_method,
                    pick_ori=None)

print(src)
print(stc.vertices)

#morphing
#recuperating volume source estimate from mixed source estimate stc
nb_sources_cortex = stc.vertices[0].shape[0] + stc.vertices[1].shape[0]
vertices_vol = np.concatenate((stc.vertices[2], stc.vertices[3]))
stc_vol = mne.VolSourceEstimate(stc.data[nb_sources_cortex:], vertices_vol, 
                                    tmin=stc.tmin, tstep=stc.tstep,
                                    subject = stc.subject)
#Maybe not the cleanest way to isolate a volume source estimate from a mixed source estimate?

#recuperating the source space I want to morph to
src_fsaverage_vol = subjects_dir + '/fsaverage/bem/fsaverage-vol-5-src.fif'
src_vol  = mne.read_source_spaces(src_fsaverage_vol)

#recuperating the source space to morph from
cereb_used = mne.SourceSpaces(inverse_operator['src'][2:])
#the source space (inverse_op['src']) is made of 4 source spaces: 2 surface source spaces and 2 volume source spaces
#- here I recuperate the 2 volume source spaces

morph = mne.compute_source_morph(src = cereb_used, 
                                    subject_from=stc_vol.subject,
                                    subject_to = 'fsaverage', 
                                    subjects_dir=subjects_dir,
                                    niter_affine=[10, 10, 5], niter_sdr=[10, 10, 5],  # just for speed
                                    src_to=src_vol, verbose=True)

stc_vol = morph.apply(stc_vol)

ValueError: vertices do not match between morph (77) and stc (33) for the volume:
[ 619  620  634  635  636  651  860  874  875  876  888  903  906  918
  935 1129 1131 1143 1144 1145 1146 1158 1159 1160 1174 1188 1189 1204
 1416 1430 1431 1445 1716  608  609  623  624  625  638  639  640  652
  653  654  668  863  864  877  879  880  892  896  907  911  922  923
  924  925 1133 1147 1148 1149 1150 1162 1164 1165 1177 1180 1195 1209
 1417 1433 1434 1448 1449 1464 1717]
[ 619  620  634  635  636  651  860  874  875  876  888  903  906  918
  935 1129 1131 1143 1144 1145 1146 1158 1159 1160 1174 1188 1189 1204
 1416 1430 1431 1445 1716]