#!/usr/bin/env python
# coding: utf-8

# In[ ]:


get_ipython().run_line_magic('pylab', 'nbagg')
from tvb.simulator.lab import *


# # Simulate cortical surface with s/M/EEG forward solutions
# 
# Cortical surface with subcortical regions, sEEG, EEG & MEG, using a stochastic
# integration.

# In[ ]:


from tvb.datatypes.cortex import Cortex
from tvb.datatypes.region_mapping import RegionMapping
from tvb.datatypes.projections import ProjectionMatrix, ProjectionSurfaceEEG
from tvb.datatypes.sensors import SensorsEEG

oscillator = models.Generic2dOscillator(a=numpy.array([0.1]), tau=numpy.array([2.0]))

white_matter = connectivity.Connectivity.from_file("connectivity_76.zip") # TVB-3018: "connectivity_192.zip"
white_matter.speed = numpy.array([4.0])
white_matter_coupling = coupling.Difference(a=numpy.array([0.014]))

rm_f_name = 'regionMapping_16k_76.txt' # TVB-3018: 'regionMapping_16k_192.txt'
rm = RegionMapping.from_file(rm_f_name)
sensorsEEG = SensorsEEG.from_file('eeg_unitvector_62.txt.bz2')
prEEG = ProjectionSurfaceEEG.from_file('projection_eeg_62_surface_16k.mat', matlab_data_name="ProjectionMatrix")

heunint = integrators.HeunStochastic(
    dt=2**-4,
    noise=noise.Additive(nsig=numpy.array([2 ** -5, ]))
)


# ## Monitors
# See shown here 3 different ways of configuring monitors.
# These methods are available for all projection monitors

# In[ ]:


fsamp = 1e3/1024.0 # 1024 Hz

monitor_MEG=monitors.MEG.from_file(rm_f_name=rm_f_name)
monitor_MEG.period=fsamp
mons = (
    monitors.EEG(sensors=sensorsEEG, projection=prEEG, region_mapping=rm, period=fsamp),
    monitor_MEG,
    monitors.iEEG.from_file('seeg_588.txt', 'projection_seeg_588_surface_16k.npy', rm_f_name=rm_f_name, period=fsamp),
    monitors.ProgressLogger(period=10.0),
)

local_coupling_strength = numpy.array([2 ** -10])
default_cortex = Cortex.from_file(region_mapping_file=rm_f_name)
default_cortex.region_mapping_data.connectivity = white_matter
default_cortex.coupling_strength = local_coupling_strength


# Bellow, adjust the simulation length to match your machine and needs

# In[ ]:


sim = simulator.Simulator(
    model=oscillator,
    connectivity=white_matter,
    coupling=white_matter_coupling,
    integrator=heunint,
    monitors=mons,
    surface=default_cortex,
    simulation_length=200.0
)
sim.configure()


# In[ ]:


eeg, meg, seeg, _ = sim.run()


# In[ ]:


figure()

for i, mon in enumerate((eeg, meg, seeg)):
    subplot(3, 1, i + 1)
    time, data = mon
    plot(time, data[:, 0, :, 0], 'k', alpha=0.1)
    ylabel(['EEG', 'MEG', 'sEEG'][i])

tight_layout()


# In[ ]: