# put some system-specific variables in the namespace ('le' dict)
%run ~/set_localenv_vars.py

nb_name = 'replicating_spiegler2016'


# this is where I look for things
outdir = '%s/notebooks/%s' %(le['data_dir'],nb_name)

# this is where things actually live
outdir_loc = '%s/%s' %('/mnt/titania-hecuba/mcintosh_lab/john/Data/notebooks', nb_name)
#outdir_loc = '%s/%s' %('/hecuba/mcintosh_lab/john/Data/notebooks', nb_name)
import os; 
if not os.path.isdir(outdir_loc): os.makedirs(outdir_loc)
if not os.path.isdir(outdir): os.system('ln -s %s %s' %(outdir_loc,outdir))

    # root path to tvb-library, tvb-data, etc. libraries. 
tvb_folder = le['code_dir'] + '/libraries_of_others/github'

tvb_dat_dir = tvb_folder + '/tvb-data'

ctx_file = tvb_dat_dir + '/tvb_data/surfaceData/cortex_16384.zip'
rm_file = tvb_dat_dir  + '/tvb_data/regionMapping/regionMapping_16k_76.bz2'
conn_file = tvb_dat_dir + '/tvb_data/connectivity/connectivity_76.zip'


# stuff for workdocs-cloudfiles

aws_key = 'drowssaperucesyreva'
aws_secret = '?teytidettopsuoyevah'  

# Generic imports

import os,sys,glob,h5py,itertools,multiprocessing,\
       numpy as np,pandas as pd
from datetime import datetime
from time import time
from scipy import optimize
from scipy.signal import hilbert
from scipy.sparse import bsr_matrix
from scipy.sparse.linalg import eigs as sp_eigs


# Visualization stuff

%matplotlib inline
from matplotlib import pyplot as plt
from matplotlib.pyplot import subplot
import matplotlib as mpl
import matplotlib.pyplot as pyplot
import matplotlib.colors
import matplotlib.ticker as ticker
import matplotlib.colors as colors
from matplotlib.tri import Triangulation
from matplotlib import cm
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from IPython.display import Image,display as d,clear_output
import seaborn as sns

from numpy import pi,cos,sin


# aws api and workdocs-cloudfiles stuff

sys.path.append(le['ipynb_workdocs_dir'])
from new_utils import nb_fig,cloudfiles_nb


# TVB stuff

nbso,nbse = sys.stdout,sys.stderr # hack part 1/2 to keep output printing properly
sys.path.append(tvb_folder +'/tvb-library')
sys.path.append(tvb_folder + '/tvb-data')
from tvb.simulator.lab import *
sys.stdout,sys.stderr = nbso,nbse  # ...hack part 2/2

from tvb.datatypes.cortex import Cortex
from tvb.datatypes.region_mapping import RegionMapping
from tvb.datatypes.projections import ProjectionMatrix


from copy import copy # seems that this needs to be done last otherwise it gets 
                      # replaced with numpy copy function

#Visualization things - see here https://johngriffiths.github.io/LabNotebook/mpl-surface-viz.html


def plot_surface_mpl(vtx,tri,data=None,rm=None,reorient='tvb',view='superior',
                     shaded=False,ax=None,figsize=(6,4), title=None,
                     lthr=None,uthr=None, nz_thr = 1E-20,
                     shade_kwargs = {'edgecolors': 'k', 'linewidth': 0.1,
                                     'alpha': None, 'cmap': 'coolwarm',
                                     'vmin': None, 'vmax': None}):
                        
  r"""Plot surfaces, surface patterns, and region patterns with matplotlib
    
  This is a general-use function for neuroimaging surface-based data, and 
  does not necessarily require construction of or interaction with tvb 
  datatypes. 

  See also:  plot_surface_mpl_mv



  Parameters
  ----------
  
  vtx           : N vertices x 3 array of surface vertex xyz coordinates 

  tri           : N faces x 3 array of surface faces

  data          : array of numbers to colour surface with. Can be either 
                  a pattern across surface vertices (N vertices x 1 array),
                  or a pattern across the surface's region mapping 
                  (N regions x 1 array), in which case the region mapping 
                  bust also be given as an argument. 
                  
  rm            : region mapping - N vertices x 1 array with (up to) N 
                  regions unique values; each element specifies which 
                  region the corresponding surface vertex is mapped to 

  reorient      : modify the vertex coordinate frame and/or orientation 
                  so that the same default rotations can subsequently be 
                  used for image views. The standard coordinate frame is 
                  xyz; i.e. first,second,third axis = left-right, 
                  front-back, and up-down, respectively. The standard 
                  starting orientation is axial view; i.e. looking down on
                  the brain in the x-y plane.
                  
                  Options: 

                    tvb (default)   : swaps the first 2 axes and applies a rotation
                                              
                    fs              : for the standard freesurfer (RAS) orientation; 
                                      e.g. fsaverage lh.orig. 
                                      No transformations needed for this; so is 
                                      gives same result as reorient=None

  view          : specify viewing angle. 
  
                  This can be done in one of two ways: by specifying a string 
                  corresponding to a standard viewing angle, or by providing 
                  a tuple or list of tuples detailing exact rotations to apply 
                  around each axis. 
                  
                  Standard view options are:
    
                  lh_lat / lh_med / rh_lat / rh_med / 
                  superior / inferior / posterior / anterior

                  (Note: if the surface contains both hemispheres, then medial 
                   surfaces will not be visible, so e.g. 'rh_med' will look the 
                   same as 'lh_lat')
                   
                  Arbitrary rotations can be specied by a tuple or a list of 
                  tuples, each with two elements, the first defining the axis 
                  to rotate around [0,1,2], the second specifying the angle in 
                  degrees. When a list is given the rotations are applied 
                  sequentially in the order given. 
                  
                  Example: rotations = [(0,45),(1,-45)] applies 45 degrees 
                  rotation around the first axis, followed by 45 degrees rotate 
                  around the second axis. 

  lthr/uthr     : lower/upper thresholds - set to zero any datapoints below / 
                  above these values
  
  nz_thr        : near-zero threshold - set to zero all datapoints with absolute 
                  values smaller than this number. Default is a very small 
                  number (1E-20), which unless your data has very small numbers, 
                  will only mask out actual zeros. 

  shade_kwargs  : dictionary specifiying shading options

                  Most relevant options (see matplotlib 'tripcolor' for full details):
                  
                    - 'shading'        (either 'gourand' or omit; 
                                        default is 'flat')
                    - 'edgecolors'     'k' = black is probably best
                    - 'linewidth'      0.1 works well; note that the visual 
                                       effect of this will depend on both the 
                                       surface density and the figure size 
                    - 'cmap'           colormap
                    - 'vmin'/'vmax'    scale colormap to these values
                    - 'alpha'          surface opacity
                  
  ax            : figure axis
  
  figsize       : figure size (ignore if ax provided)
  
  title         : text string to place above figure
  
  
  
                  
  Usage
  -----
       

  Basic freesurfer example:

  import nibabel as nib
  vtx,tri = nib.freesurfer.read_geometry('subjects/fsaverage/surf/lh.orig')
  plot_surface_mpl(vtx,tri,view='lh_lat',reorient='fs')



  Basic tvb example:
  
  ctx = cortex.Cortex.from_file(source_file = ctx_file,
                                region_mapping_file =rm_file)
  vtx,tri,rm = ctx.vertices,ctx.triangles,ctx.region_mapping
  conn = connectivity.Connectivity.from_file(conn_file); conn.configure()
  isrh_reg = conn.is_right_hemisphere(range(conn.number_of_regions))
  isrh_vtx = np.array([isrh_reg[r] for r in rm])
  dat = conn.tract_lengths[:,5]

  plot_surface_mpl(vtx=vtx,tri=tri,rm=rm,data=dat,view='inferior',title='inferior')

  fig, ax = plt.subplots()
  plot_surface_mpl(vtx=vtx,tri=tri,rm=rm,data=dat, view=[(0,-90),(1,55)],ax=ax,
                   title='lh angle',shade_kwargs={'shading': 'gouraud', 'cmap': 'rainbow'})

   
  """
    
  # Copy things to make sure we don't modify things 
  # in the namespace inadvertently. 
    
  vtx,tri = vtx.copy(),tri.copy()
  if data is not None: data = data.copy()

  # 1. Set the viewing angle 
  
  if reorient == 'tvb':
    # The tvb default brain has coordinates in the order 
    # yxz for some reason. So first change that:   
    vtx = np.array([vtx[:,1],vtx[:,0],vtx[:,2]]).T.copy()
    
    # Also need to reflect in the x axis
    vtx[:,0]*=-1

  # (reorient == 'fs' is same as reorient=None; so not strictly needed
  #  but is included for clarity)
   


  # ...get rotations for standard view options
    
  if   view == 'lh_lat'    : rots =  [(0,-90),(1,90)  ]
  elif view == 'lh_med'    : rots =  [(0,-90),(1,-90) ] 
  elif view == 'rh_lat'    : rots =  [(0,-90),(1,-90) ]
  elif view == 'rh_med'    : rots =  [(0,-90),(1,90)  ]
  elif view == 'superior'  : rots =   None
  elif view == 'inferior'  : rots =   (1,180)
  elif view == 'anterior'  : rots =   (0,-90)
  elif view == 'posterior' : rots =  [(0, -90),(1,180)]
  elif (type(view) == tuple) or (type(view) == list): rots = view 

  # (rh_lat is the default 'view' argument because no rotations are 
  #  for that one; so if no view is specified when the function is called, 
  #  the 'rh_lat' option is chose here and the surface is shown 'as is' 
                            
                            
  # ...apply rotations                          
     
  if rots is None: rotmat = np.eye(3)
  else:            rotmat = get_combined_rotation_matrix(rots)
  vtx = np.dot(vtx,rotmat)

                                    
      
  # 2. Sort out the data
                                    
                                    
  # ...if no data is given, plot a vector of 1s. 
  #    if using region data, create corresponding surface vector 
  if data is None: 
    data = np.ones(vtx.shape[0]) 
  elif data.shape[0] != vtx.shape[0]: 
    data = np.array([data[r] for r in rm])
    
  # ...apply thresholds
  if uthr: data *= (data < uthr)
  if lthr: data *= (data > lthr)
  data *= (np.abs(data) > nz_thr)

                                    
  # 3. Create the surface triangulation object 
  
  x,y,z = vtx.T
  tx,ty,tz = vtx[tri].mean(axis=1).T
  tr = Triangulation(x,y,tri[np.argsort(tz)])
                
  # 4. Make the figure 

  if ax is None: fig, ax = plt.subplots(figsize=figsize)  
  
  #if shade = 'gouraud': shade_opts['shade'] = 
  tc = ax.tripcolor(tr, np.squeeze(data), **shade_kwargs)
                        
  ax.set_aspect('equal')
  ax.axis('off')
    
  if title is not None: ax.set_title(title)
        
        
        
        
def plot_surface_mpl_mv(vtx=None,tri=None,data=None,rm=None,hemi=None,   # Option 1
                        vtx_lh=None,tri_lh=None,data_lh=None,rm_lh=None, # Option 2
                        vtx_rh=None,tri_rh=None,data_rh=None,rm_rh=None,
                        title=None,**kwargs):

  r"""Convenience wrapper on plot_surface_mpl for multiple views 
   
  This function calls plot_surface_mpl five times to give a complete 
  picture of a surface- or region-based spatial pattern. 

  As with plot_surface_mpl, this function is written so as to be 
  generally usable with neuroimaging surface-based data, and does not 
  require construction of of interaction with tvb datatype objects. 

  In order for the medial surfaces to be displayed properly, it is 
  necessary to separate the left and right hemispheres. This can be 
  done in one of two ways: 

  1. Provide single arrays for vertices, faces, data, and 
     region mappings, and addition provide arrays of indices for 
     each of these (vtx_inds,tr_inds,rm_inds) with 0/False 
     indicating left hemisphere vertices/faces/regions, and 1/True 
     indicating right hemisphere. 

     Note: this requires that 

  2. Provide separate vertices,faces,data,and region mappings for 
     each hemisphere (vtx_lh,tri_lh; vtx_rh,tri_rh,etc...)


 
  Parameters
  ----------

  (see also plot_surface_mpl parameters info for more details)

  (Option 1)

  vtx               :  surface vertices
 
  tri               : surface faces

  data              : spatial pattern to plot

  rm                : surface vertex to region mapping

  hemi              : hemisphere labels for each vertex
                      (1/True = right, 0/False = left) - 
      

  OR

  (Option 2)

  vtx_lh            : left hemisphere surface_vertices
  vtx_rh            : right ``      ``    ``     ``
  
  tri_lh            : left hemisphere surface faces 
  tri_rh            : right ``      ``    ``     ``

  data_lh          : left hemisphere surface_vertices
  data_rh          : right ``      ``    ``     ``

  rm_lh            : left hemisphere region_mapping
  rm_rh            : right ``      ``    ``     ``


  title            : title to show above middle plot
 
  kwargs           : additional tripcolor kwargs; see plot_surface_mpl

 

  Examples
  ----------

  # TVB default data

  # Plot one column of the region-based tract lengths 
  # connectivity matrix. The corresponding region is 
  # right auditory cortex ('rA1')

  ctx = cortex.Cortex.from_file(source_file = ctx_file,
                                region_mapping_file =rm_file)
  vtx,tri,rm = ctx.vertices,ctx.triangles,ctx.region_mapping
  conn = connectivity.Connectivity.from_file(conn_file); conn.configure()
  isrh_reg = conn.is_right_hemisphere(range(conn.number_of_regions))
  isrh_vtx = np.array([isrh_reg[r] for r in rm])
  dat = conn.tract_lengths[:,5]

  plot_surface_mpl_mv(vtx=vtx,tri=tri,rm=rm,data=dat,
                      hemi=isrh_vtx,title=u'rA1 \ntract length')

  plot_surface_mpl_mv(vtx=vtx,tri=tri,rm=rm,data=dat,
                    hemi=isrh_vtx,title=u'rA1 \ntract length',
                    shade_kwargs = {'shading': 'gouraud',
                                    'cmap': 'rainbow'}) 


  """
   

 
  if vtx is not None:                                    # Option 1
    tri_hemi = hemi[tri].any(axis=1)
    tri_lh,tri_rh = tri[tri_hemi==0],tri[tri_hemi==1]
  elif vtx_lh is not None:                               # Option 2
    vtx = np.vstack([vtx_lh,vtx_rh])
    tri = np.vstack([tri_lh,tri_rh+tri_lh.max()+1])

  if data_lh is not None:                                # Option 2
    data = np.hstack([data_lh,data_rh])
    
  if rm_lh is not None:                                  # Option 2 
    rm = np.hstack([rm_lh,rm_rh + rm_lh.max() + 1])
    
 

  # 2. Now do the plots for each view

  # (Note: for the single hemispheres we only need lh/rh arrays for the 
  #  faces (tri); the full vertices, region mapping, and data arrays
  #  can be given as arguments, they just won't be shown if they aren't 
  #  connected by the faces in tri )
  
  # LH lateral
  plot_surface_mpl(vtx,tri_lh,data=data,rm=rm,view='lh_lat',
                   ax=subplot(2,3,1),**kwargs)
    
  # LH medial
  plot_surface_mpl(vtx,tri_lh, data=data,rm=rm,view='lh_med',
                   ax=subplot(2,3,4),**kwargs)
    
  # RH lateral
  plot_surface_mpl(vtx,tri_rh, data=data,rm=rm,view='rh_lat',
                   ax=subplot(2,3,3),**kwargs)
    
  # RH medial
  plot_surface_mpl(vtx,tri_rh, data=data,rm=rm,view='rh_med',
                   ax=subplot(2,3,6),**kwargs)
    
  # Both superior
  plot_surface_mpl(vtx,tri, data=data,rm=rm,view='superior',
                   ax=subplot(1,3,2),title=title,**kwargs)
    
  plt.subplots_adjust(left=0.0, right=1.0, bottom=0.0,
                      top=1.0, wspace=0, hspace=0)    
    
    
    
    

def get_combined_rotation_matrix(rotations):
  '''Return a combined rotation matrix from a dictionary of rotations around 
     the x,y,or z axes'''
  rotmat = np.eye(3)
    
  if type(rotations) is tuple: rotations = [rotations] 
  for r in rotations:
    newrot = get_rotation_matrix(r[0],r[1])
    rotmat = np.dot(rotmat,newrot)
  return rotmat



def get_rotation_matrix(rotation_axis, deg):

  '''Return rotation matrix in the x,y,or z plane'''



  # (note make deg minus to change from anticlockwise to clockwise rotation)
  th = -deg * (pi/180) # convert degrees to radians

  if rotation_axis == 0:
    return np.array( [[    1,         0,         0    ],
                      [    0,      cos(th),   -sin(th)],
                      [    0,      sin(th),    cos(th)]])
  elif rotation_axis ==1:
    return np.array( [[   cos(th),    0,        sin(th)],
                      [    0,         1,          0    ],
                      [  -sin(th),    0,        cos(th)]])
  elif rotation_axis ==2:
    return np.array([[   cos(th),  -sin(th),     0    ],
                     [    sin(th),   cos(th),     0   ],
                     [     0,         0,          1   ]])

%%javascript
IPython.load_extensions('calico-spell-check', 
                        'calico-document-tools', 'calico-cell-tools');

cnb = cloudfiles_nb('aws', [aws_key,aws_secret])
cnb.initialize_folder(nb_name)

os.chdir(outdir)

# ('_sp' refers to the variable symbol as used in Spiegler et al; 
#   so as to avoid ambiguity)


gamma_sp = 1.21
epsilon_sp =  12.3083

# is this the correct? 
# - units in Spiegler are S^-1. 
# - default value for g2do d is 0.02
# - so this gives 0.07674
eta_sp = (1/1000.) * 76.74    #eta_sp = 76.74 # 1. # 76.74 ##1. # 76.74 # 1.



sp_g2do_params = dict(d = eta_sp,
                      tau = 1.,
                      f = 1.,
                      e = 0., 
                      g = -gamma_sp,
                      alpha = 1., 
                      gamma = 1., 
                      c = 0.,
                      b= -epsilon_sp, # should not be negative? LOOKS LIKE IT AL COMES DOWN TO THIS PARAM
                      beta = 0,
                      a = 0.)

# <!-- collapse=False-->


# Run sims

# Specify model with Spiegler parameters
sp_g2do = models.Generic2dOscillator(**sp_g2do_params)
#sp_g2do.d = (1/1000.)*eta_sp

# Evolve system for 1000 sytems, step duration 0.1ms
time,dat = sp_g2do.stationary_trajectory(n_step=5000,dt=0.1)

# Put sim results into a pandas dataframe
V_dat = np.squeeze(dat[:,0,:,:])
W_dat = np.squeeze(dat[:,1,:,:])
df_VW = pd.DataFrame([V_dat,W_dat],
                     index=['V', 'W'], 
                     columns=time).T



# Make 2 quick plots:

f = outdir + '/pics/single_node_phase_plane_dynamics.png'

fig, ax = plt.subplots(ncols=2, figsize=(12,3))

# V and W vs. time
df_VW.plot(ax=ax[0])

# V vs. W trajectory
ymax = df_VW.W.max() + df_VW.W.std()*1.5
ymin = df_VW.W.min() - df_VW.W.std()*1.5
xmax = df_VW.V.max() + df_VW.V.std()*1.5
xmin = df_VW.V.min() - df_VW.V.std()*1.5
df_VW.plot(x='V', y='W', legend=False,ax=ax[1],
           xlim=[xmin,xmax],ylim=[ymin,ymax],alpha=0.5);



plt.savefig(f, bbox_inches='tight')
plt.close()


# Embed in LN
cap = ''; label = 'Single node phase plane dynamics'
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

f = outdir + '/pics/Spiegler_fig2a_replication.png'


# First run the model for the parameter values and 
# put into a pandas dataframe

bif_params = [-2 , -0.5 ,1.]# 1.]

res = {}

for b in bif_params:

  mod = models.Generic2dOscillator(**sp_g2do_params)
  mod.g = b

  initstim = np.array([2.,0.])[:,np.newaxis]
    
  time,dat = mod.stationary_trajectory(n_step=5000,dt=0.1,initial_conditions=initstim)
    
    
  V_dat = np.squeeze(dat[:,0,:,:])
  W_dat = np.squeeze(dat[:,1,:,:])
  df_VW = pd.DataFrame([V_dat,W_dat],
                       index=['V', 'W'], 
                       columns=time).T
  
  res[b] = df_VW
    

df_VWb = pd.concat(res);
df_VWb.index.names = ['g', 'time']


# Now make the 3D plot

# (note that matplotlib's 3d plotting doesn't get the order of the elements quite right
# (maybe use plotly or something to get a better version of that))

mpl.rcParams['legend.fontsize'] = 15;

fig = plt.figure();
ax = fig.gca(projection='3d');

for b_it,b in enumerate(bif_params):

  z = df_VWb.ix[b].W.values
  y = df_VWb.ix[b].V.values
  x = np.ones(y.shape[0])*(0.2*b_it)  

  ax.plot(x,y,z, label='$\gamma = %s $' %b);#,alpha=0.5,linewidth=0.5)

    
# Draw at thick black line through the first two
# ...and  a dashed black line through the third one
ax.plot([-0.1,0.3],[0,0],[0,0],c='k', alpha=0.3);
ax.plot([0.3,0.6],[0,0],[0,0],c='k', alpha=0.3, linestyle='--');


ax.view_init(30, 250); # angle)

#ax.axis('off')
ax.legend();
plt.show();


fig.savefig(f, bbox_inches='tight');
plt.close();
plt.clf();
clear_output();

# <!-- collapse=False-->

# Embed in LN
cap = ''; label = 'Spielger Fig 2a replication'
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500));
d(im)

# <!-- collapse=False-->


# Make fig
f = outdir + '/pics/Spiegler_fig2b_replication.png'

fig, ax = plt.subplots(ncols=3, figsize=(12,3))
for b_it, b in enumerate(bif_params):  df_VWb.ix[b].ix[:150].plot(ax=ax[b_it],ylim=[-5.5,5.5])    
plt.tight_layout() 

plt.savefig(f, bbox_inches='tight')
plt.close()


# Embed in LN
cap = ''; label = 'Spielger Fig 2b replication'
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

conn = connectivity.Connectivity(load_default=True)

stim_node_names = ['rPFCCL']
stim_node_inds = []
for s in stim_node_names:
  for r_it,r in enumerate(conn.region_labels): 
    if s == r: 
      stim_node_inds.append(r_it)
print stim_node_inds

# <!-- collapse=False-->
f = outdir + '/pics/stimulus_pattern.png'

eqn_t = equations.PulseTrain()
eqn_t.parameters['onset'] = 1.5e3
eqn_t.parameters['T'] = 500.0
eqn_t.parameters['tau'] = 50.0

# configure stimulus spatial pattern
stim_weight = 0.1
weighting = np.zeros((76, ))
weighting[stim_node_inds] = stim_weight

stimulus = patterns.StimuliRegion(
    temporal=eqn_t,
    connectivity=conn,
    weight=weighting)


#Configure space and time
stimulus.configure_space()
stimulus.configure_time(np.arange(0., 2e3, 2**-4))

#And take a look
plot_pattern(stimulus)


plt.savefig(f, bbox_inches='tight')
plt.close()



# Embed figure
cap = ''; label = 'stimulus pattern'
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

def run_sim(g2do_params,conn=None,
            sim_len = 2e3,dt = 0.1,
            tau = 2,lca=0.01,
            tavg_per = 1.0,
            stim_node_inds=[17],
            stim_eqn_onset =1.5e3,
            stim_eqn_T = 500.0,
            stim_eqn_tau = 50.0,
            stim_weight = 0.1):

  if conn == None:
    conn = connectivity.Connectivity(load_default=True)
    
  model = models.Generic2dOscillator(**g2do_params)
  cpl=coupling.Linear(a=lca)
  solver=integrators.HeunDeterministic(dt=dt)
  mons=(monitors.TemporalAverage(period=tavg_per),)
    
  eqn_t = equations.PulseTrain()
  eqn_t.parameters['onset'] = stim_eqn_onset  
  eqn_t.parameters['T'] = stim_eqn_T
  eqn_t.parameters['tau'] = stim_eqn_tau

  weighting = np.zeros((76, ))
  weighting[stim_node_inds] = stim_weight

  stimulus = patterns.StimuliRegion(temporal=eqn_t,
                                    connectivity=conn,
                                    weight=weighting)
  stimulus.configure_space()
  stimulus.configure_time(np.arange(0., 3e3, 2**-4))
    
  sim = simulator.Simulator(model = model,connectivity=conn,coupling=cpl,
                            integrator=solver, 
                            monitors=mons,
                            stimulus=stimulus)
  sim.configure()

  (tavg_time, tavg_data),  = sim.run(simulation_length=sim_len) 
    
  df_tavg = pd.DataFrame(np.squeeze(tavg_data),index=tavg_time)
  df_tavg.index.names = ['t']
  df_tavg.columns.names= ['node']
  df_tavg.columns = conn.region_labels    
    
  return df_tavg

# <!-- collapse=False-->


# Run simulation
df_tavg_sp1 = run_sim(sp_g2do_params,lca=0.1,sim_len=2e3)


# Make the figure
f = outdir + '/pics/df_tavg_sp1_ts.png'
fig, ax = plt.subplots(nrows=2, figsize=(12,6))
df_tavg_sp1.plot(legend=False,ax=ax[0])
df_tavg_sp1.ix[1450:1650].plot(legend=False,ax=ax[1])
plt.savefig(f, bbox_inches='tight')
plt.close()


# Embed in LN
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# <!-- collapse=False-->


# Run sim
ps = copy(sp_g2do_params) # copy cos we don't want to change beta in this dict
ps['beta'] = 1
df_tavg_sp2 = run_sim(ps,lca=0.1,sim_len=2e3)


# Make figure
f = outdir + '/pics/df_tavg_sp2_ts.png'
fig, ax = plt.subplots(nrows=2, figsize=(12,6))
df_tavg_sp2.plot(legend=False,ax=ax[0])
df_tavg_sp2.ix[1450:1650].plot(legend=False,ax=ax[1])
plt.savefig(f, bbox_inches='tight')
plt.close()

# Embed in LN
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# <!-- collapse=False-->
f = outdir + '/pics/df_tavg_sp3_ts.png'


# Run simulation
sp = copy(sp_g2do_params)
sp['e'] = 3.
df_tavg_sp3 = run_sim(ps,lca=0.1,sim_len=2e3)


# Make figure
f = outdir + '/pics/df_tavg_sp3_ts.png'
fig, ax = plt.subplots(nrows=2, figsize=(12,6))
df_tavg_sp3.plot(legend=False,ax=ax[0])
df_tavg_sp3.ix[1450:1650].plot(legend=False,ax=ax[1])
plt.savefig(f, bbox_inches='tight')
plt.close()

# Embed in LN
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# <!-- collapse=False-->


# Run sim

ps = copy(sp_g2do_params)
ps['beta'] = 1.
ps['e'] = 3.
df_tavg_sp4 = run_sim(ps,lca=0.1,sim_len=2e3)


#Make fig
f = outdir + '/pics/df_tavg_sp4_ts.png'
fig, ax = plt.subplots(nrows=2, figsize=(12,6))
df_tavg_sp4.plot(legend=False,ax=ax[0])
df_tavg_sp4.ix[1450:1650].plot(legend=False,ax=ax[1])
plt.savefig(f, bbox_inches='tight')
plt.close()

# Embed in LN`
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# <!-- collapse=False-->

#Make fig
f = outdir + '/pics/df_tavg_sp2_heatmap.png'
fig, ax = plt.subplots(figsize=(12,12))
sns.heatmap(df_tavg_sp2.ix[1490:1650].T, ax=ax,
            vmin=-0.01,vmax=0.01,xticklabels='')
plt.savefig(f, bbox_inches='tight')
plt.close()

# Embed in LN`
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# These are some visualization things necessary for the brain surface plots
ctx = cortex.Cortex.from_file(source_file = ctx_file,
                              region_mapping_file =rm_file)
vtx,tri,rm = ctx.vertices,ctx.triangles,ctx.region_mapping
conn = connectivity.Connectivity.from_file(conn_file); conn.configure()
isrh_reg = conn.is_right_hemisphere(range(conn.number_of_regions))
isrh_vtx = np.array([isrh_reg[r] for r in rm])
dat = conn.tract_lengths[:,0] # 'dat' is one column of the region-based tract lengths 
                              # connectivity matrix. The corresponding region is 
                              # right auditory cortex ('rA1')

# <!-- collapse=False-->
f = outdir + '/pics/df_tavg_sp2_brainsurf.png'


# Make fig
fig, ax = plt.subplots(ncols=10, nrows=2,figsize=(15,3))
cmap = cm.Reds
cmap.set_under(color='w')

kws = {'edgecolors': 'k', 'vmin': 0., 'cmap': cmap, 
       'vmax': 0.02, 'alpha': None, 'linewidth': 0.01}

ts = [1500.5, 1520.5,1540.5,1560.5,1580.5,1600.5,1650.5,1700.5,1750.5,1800.5]

for t_it,t in enumerate(ts):

  dat = df_tavg_sp2.ix[t].abs().values
    
  plot_surface_mpl(vtx=vtx,tri=tri,data=dat,rm=rm,ax=ax[0][t_it],
                   shade_kwargs=kws,
                   view='rh_lat')
    
  plot_surface_mpl(vtx=vtx,tri=tri,data=dat,rm=rm,ax=ax[1][t_it],
                   shade_kwargs=kws,
                   view='superior')
    
    
  ax[0][t_it].set_title('t=%1.1fms' %t)
    
        
plt.savefig(f, bbox_inches='tight')
plt.close()


# Embed in LN`
cap = ''; label = ''
im = nb_fig(f,label,cap,cnb,show_fignum=False,size=(800,500))
d(im)

# Run sims
roi_names = conn.region_labels
roi_inds = range(roi_names.shape[0])

stim_roisims = {}
sp = copy(sp_g2do_params)
sp['beta'] = 1.

for roi_name,roi_ind in zip(roi_names,roi_inds):
    
  df = run_sim(sp,stim_node_inds = [roi_ind],
               lca=0.1,sim_len=2000)
  stim_roisims[roi_name] = df
    
    
# Collect into a single dataframe
df_stimroisims = pd.concat(stim_roisims)
df_stimroisims.index.names = ['stim_roi', 't']
df_stimroisims.columns.names = ['roi']

# Same to file
df_stimroisims.to_pickle(outdir + '/df_stimroisims.pkl')