#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Sat Jan 14 23:05:17 2017
    
   
@author: daisy
"""

#%% 
import mne
import os.path as op
from mne.beamformer import lcmv
import numpy as np

subject = 'lanmengying'
data_path= '/Users/daisy/Desktop/MEG_data/s38/run2_mooney_normal'
subjects_dir ='/Applications/freesurfer/subjects'
con1='mooney'
con2='normal'

#%% Make forward solution

#create BEM model
ico = 4 
conductivity = (0.3,) 
model = mne.make_bem_model(subject=subject, ico=ico,
                           conductivity=conductivity,
                           subjects_dir=subjects_dir)

#create BEM solution
bem = mne.make_bem_solution(model)
fname_bem = op.join(subjects_dir, subject,'bem', subject+'-bem.fif')
mne.write_bem_solution(fname_bem, bem)

#Compute source space on surface
spacing='ico5'
src = mne.setup_source_space(subject, spacing=spacing,
                             subjects_dir=subjects_dir,
                             add_dist=False)

fname_src = subjects_dir +'/'+subject+'/bem/'+subject+'-'+spacing[0:3]+'-'+ spacing[3]+'-src.fif'
src.save(fname_src,overwrite=True)
#src = mne.read_source_spaces(fname_src)

#compute forward solution
fname_trans = data_path +'/'+ subject + '-trans.fif'
fname_raw =data_path +'/'+ subject+'-raw.fif'
fname_fwd = op.join(subjects_dir, subject, 'bem', subject+'-fwd.fif')
forward = mne.make_forward_solution(fname_raw,trans=fname_trans,src=src,bem=fname_bem)
forward = mne.convert_forward_solution(forward, surf_ori=True) 


#%% lcmv

# read files
epochs_fname=data_path + '/'+ subject+ '-epo.fif'
epochs = mne.read_epochs(epochs_fname, preload=True)

# calaulate evoke
evoked1= epochs[con1].average() 
evoked2= epochs[con2].average() 

# Estimate covarariance matrix
noise_cov = mne.compute_covariance(epochs, tmin=-0.25,tmax=0, method='shrunk')
data_cov=mne.compute_covariance(epochs,tmin=0.1, tmax=0.35, method='shrunk')

# ignore compensation grade
evoked1.info['comps']=[]
evoked2.info['comps']=[]

# Calculate inverse solution
stc1=lcmv(evoked1,forward,noise_cov,data_cov,reg=0.01,pick_ori='max-power',weight_norm='unit-noise-gain')
stc2=lcmv(evoked2,forward,noise_cov,data_cov,reg=0.01,pick_ori='max-power',weight_norm='unit-noise-gain')

# Save surface data
fname_stc1=data_path+ '/'+subject+'_'+con1+'_maxpower'
fname_stc2=data_path+ '/'+subject+'_'+con2+'_maxpower'

stc1.save(fname_stc1)
stc2.save(fname_stc2)