import mayavi.mlab as mlab
import matplotlib.pyplot as plt

import os.path as op
import numpy as np

from mne.evoked import read_evokeds
from mne.cov import read_cov
from mne.source_space import read_source_spaces, setup_volume_source_space
from mne.bem import read_bem_solution

from mne.forward import make_forward_solution
from mne.minimum_norm import make_inverse_operator, apply_inverse
from mne import (pick_channels, pick_types_forward, read_labels_from_annot,
                 stc_to_label, fit_dipole, read_trans)

from mne.io.constants import FIFF
from mne.transforms import invert_transform, apply_trans

subjects_dir = '/mnt/isilon/meg_lab/bloyl/varITI_variability/fs_dir/'
subject = 'A040'

evoked_fname = op.join(subjects_dir, subject, 'meg_data',
                       '%s-ave.fif' % subject)
bem_fname = op.join(subjects_dir, subject, 'bem', 'A040-bem.sol')
trans_fname = op.join(subjects_dir, subject, 'meg_data', '%s-trans.fif'
                      % subject)
cov_fname = op.join(subjects_dir, subject, 'meg_data', '%s-cov.fif' % subject)
src_fname = op.join(subjects_dir, subject, 'meg_data', 'cortical-src.fif')
dip_src_fname = op.join(subjects_dir, subject, 'meg_data', 'anat-src.fif')

# read inputs
evoked = read_evokeds(evoked_fname)[0]
noise_cov = read_cov(cov_fname)
bem = read_bem_solution(bem_fname)
src = read_source_spaces(src_fname)

head_mri_trans = read_trans(trans_fname)
mri_head_trans = invert_transform(head_mri_trans)

# conform transforms
assert(mri_head_trans['to'] == FIFF.FIFFV_COORD_HEAD)
assert(mri_head_trans['from'] == FIFF.FIFFV_COORD_MRI)

# mask for just right sensors...
lt_chans = [k['ch_name'] for k in evoked.info['chs']
            if k['loc'][0] <= 0 or k['kind'] == FIFF.FIFFV_REF_MEG_CH]
rt_chans = [k['ch_name'] for k in evoked.info['chs']
            if k['loc'][0] >= 0 or k['kind'] == FIFF.FIFFV_REF_MEG_CH]

lt_picks = pick_channels(evoked.info['ch_names'], lt_chans)

evoked_plot = evoked.copy()
evoked_plot.data[lt_picks, :] = 0

# select the rt_channels
evoked.pick_channels(rt_chans)

##############################################################################
# fit cortical sheet source locations

# make forward operator
fwd = make_forward_solution(evoked.info, trans=trans_fname, src=src,
                            bem=bem, meg=True, eeg=False, n_jobs=1)

fwd = pick_types_forward(fwd, meg=True, eeg=False)

# make inverse operator
inv_op = make_inverse_operator(evoked.info, fwd, noise_cov, loose=1)

stc = apply_inverse(evoked, inv_op, 1 / 9., method='dSPM',
                    pick_ori=None)

# extract functional label
aparc_label_names = ['bankssts-rh', 'transversetemporal-rh',
                     'superiortemporal-rh']

labels = [read_labels_from_annot(subject, parc='aparc',
                                 subjects_dir=subjects_dir,
                                 regexp=lname)[0]
          for lname in aparc_label_names]

anat_label = labels[0]
for l in labels[1:]:
    anat_label += l

stc_mean = stc.copy().crop(0, 0.350).mean()

stc_mean_label = stc_mean.in_label(anat_label)
data = np.abs(stc_mean_label.data)
stc_mean_label.data[data < 0.6 * np.max(data)] = 0.

lh_labels, rh_labels = stc_to_label(stc_mean_label, src=src, smooth=True,
                                    subjects_dir=subjects_dir, connected=True)

func_label = rh_labels[0]

# exxtract timecourse
mean = stc.extract_label_time_course(func_label, src, mode='mean')
stc_func_label = stc.in_label(func_label)

#############################################################################
# fit a signel location at the center of the functional label..
# rt_aud_src_sRAS = func_label.pos.mean(axis=0)
rt_aud_src_sRAS = np.array([0.0452731, -0.01446342, -0.00180349])
rt_aud_src_head = apply_trans(mri_head_trans, rt_aud_src_sRAS)

###
# Dipole fit.
dip_fixed = fit_dipole(evoked, noise_cov, bem, trans_fname,
                       pos=rt_aud_src_head)[0]
##############################################################################
# fit anatomical source locations
pos = dict()
pos['rr'] = np.array([rt_aud_src_sRAS])
pos['nn'] = np.array([[1., 0., 0.] for k in pos['rr']])
anat_src = setup_volume_source_space(subject=None, pos=pos, bem=bem_fname,
                                     mindist=5)
# anat_src = read_source_spaces(dip_src_fname)
# make forward operator
fwd_anat = make_forward_solution(evoked.info, trans=trans_fname, src=anat_src,
                                 bem=bem, meg=True, eeg=False, n_jobs=1)
if fwd_anat['nsource'] != len(anat_src[0]['rr']):
    raise RuntimeError('Some Sources were excluded!!!!')
fwd_anat = pick_types_forward(fwd_anat, meg=True, eeg=False)

# make inverse operator
inv_op_anat = make_inverse_operator(evoked.info, fwd_anat, noise_cov, loose=1)

stc_anat = apply_inverse(evoked, inv_op_anat, 0, method='dSPM',
                         pick_ori=None)

##############################################################################

plt.figure()
h0, = plt.plot(1e3 * stc.times, mean.T / mean.max(), 'r', linewidth=3,
               label='Surface dSPM (mean func Label)')
h1, = plt.plot(1e3 * dip_fixed.times,
               dip_fixed.amplitude / dip_fixed.amplitude.max(), 'g',
               linewidth=3, label='Fixed dipole')
h2, = plt.plot(1e3 * stc_anat.times, stc_anat.data[0] / stc_anat.data[0].max(),
               'b', linewidth=3, label='Single Loc dSPM')
plt.legend()
plt.show()

# plot source locations
rt_loc = rt_aud_src_sRAS
mlab.points3d(rt_loc[0], rt_loc[1], rt_loc[2], scale_factor=0.002,
              color=(1, 0, 0))

# mlab.points3d(src[1]['rr'][:, 0], src[1]['rr'][:, 1], src[1]['rr'][:, 2],
#              scale_factor=0.001, color=(1,1,1), opacity=0.2)

mlab.points3d(func_label.pos[:, 0], func_label.pos[:, 1], func_label.pos[:, 2],
              scale_factor=0.001, color=(0, 1, 0))

mlab.show()
plt.show()