# Create config file
with open('config.py', 'w') as fp:
    fp.write(
'''import os
base_path = os.path.abspath(".")
data_path = os.path.abspath("simulations")
jobscript_template = "python {base_path}/run_simulation.py {label}"
submit_cmd = "bash -c"
''')

%matplotlib inline
import numpy as np
import os
import nest
import json

from multiarea_model import MultiAreaModel
from config import base_path, data_path

import sys
sys.path.append('./figures/MAM2EBRAINS')

!pip install nested_dict dicthash

# Jupyter notebook display format setting
from IPython.display import display, HTML
style = """
<style>
table {float:left}
</style>
"""
display(HTML(style))

# Ignore and don't display warning messages
import warnings
warnings.filterwarnings("ignore")

# Downscaling factor
# Value range/options: (0, 1.]
# Value assgined: 0.005
scale_down_to = 0.005 # Change it to 1. for running the fullscale network

# Scaling factor for cortico-cortical connections (chi) 
# Value range/options: [1., 2.5]
# Value assgined: 1.0
cc_weights_factor = 1.0

# Cortical areas included in the simulation
# Value range/options: any sublist of complete_ares_list
# Value assgined: complete_ares_list
areas_simulated = ['V1', 'V2', 'VP', 'V3', 'V3A', 'MT', 'V4t', 'V4', 'VOT', 'MSTd', 'PIP', 'PO', 'DP', 'MIP', 'MDP', 'VIP', 'LIP', 'PITv', 'PITd', 'MSTl', 'CITv', 'CITd', 'FEF', 'TF', 'AITv', 'FST', '7a', 'STPp', 'STPa', '46', 'AITd', 'TH']

# Firing rates used to replace the non-simulated areas
# Value range/options: None, 'hom_poisson_stat', 'het_poisson_stat', 'het_current_nonstat'
# Value assgined: 'het_poisson_stat'
replace_non_simulated_areas = 'het_poisson_stat'

# Connection parameters
conn_params = {
    'replace_non_simulated_areas': 'het_poisson_stat', # Whether to replace non-simulated areas by Poisson sources with the same global rate, by default: None
    'g': -11., # It sets the relative inhibitory synaptic strength, by default: -16.
    'K_stable': 'K_stable.npy', # Whether to apply the stabilization method of Schuecker, Schmidt et al. (2017), by default: None
    'fac_nu_ext_TH': 1.2, # Increase the external input to 2/3E and 5E in area TH
    'fac_nu_ext_5E': 1.125, # Increase the external Poisson indegree onto 5E
    'fac_nu_ext_6E': 1.41666667, # Increase the external Poisson indegree onto 6E
    'av_indegree_V1': 3950. # Adjust the average indegree in V1 based on monkey data
}

# Input parameters
input_params = {
    'rate_ext': 10. # Rate of the Poissonian spike generator (in spikes/s)
} 

# Neuron parameters
neuron_params = {
    'V0_mean': -150., # Mean for the distribution of initial membrane potentials, by default: -100.
    'V0_sd': 50.} # Standard deviation for the distribution of initial membrane potentials, by default: 50.

# Network parameters
network_params = {
    'N_scaling': scale_down_to, # Scaling of population sizes, by default: 1.
    'K_scaling': scale_down_to, # Scaling of indegrees, by default: 1.
    'fullscale_rates': 'tests/fullscale_rates.json', # Absolute path to the file holding full-scale rates for scaling synaptic weights, by default: None
    'input_params': input_params, # Input parameters
    'connection_params': conn_params, # Connection parameters
    'neuron_params': neuron_params # Neuron parameters
} 

# Simulation parameters
sim_params = {
    'areas_simulated': areas_simulated,
    't_sim': 2000., # Simulated time (in ms), by default: 10.0
    # 't_sim': 1500., # Simulated time (in ms), by default: 10.0
    'num_processes': 1, # The number of MPI processes, by default: 1
    'local_num_threads': 1, # The number of threads per MPI process, by default: 1
    'recording_dict': {'record_vm': False},
    'rng_seed': 1  # global random seed
}

# Theory paramters (theory_params)
theory_params = {
    'dt': 0.1 # The time step of the mean-field theory integration, by default: 0.01
} 

M = MultiAreaModel(network_params, 
                   simulation=True,
                   sim_spec=sim_params,
                   theory=True,
                   theory_spec=theory_params)

p, r = M.theory.integrate_siegert()
print("Mean-field theory predicts an average "
      "firing rate of {0:.3f} spikes/s across all populations.".format(np.mean(r[:, -1])))

# Indegrees
# Dictionary of nodes indegrees organized as:
# {<source_area>: {<source_pop>: {<target_area>: {<target_pop>: indegree_values}}}}
# M.K

# Synapses
# Dictionary of synapses that target neurons receive, it is organized as:
# {<source_area>: {<source_pop>: {<target_area>: {<target_pop>: number_of_synapses}}}}
# M.synapses

# run the simulation, depending on the model parameter and downscale ratio, the running time varies largely.
M.simulation.simulate()

label_spikes = M.simulation.label
label = M.simulation.label

from MAM2EBRAINS_LOAD_DATA import load_data
A, tsteps, firing_rate = load_data(M)

from MAM2EBRAINS_VISUALIZATION import plot_instan_mean_firing_rate
plot_instan_mean_firing_rate(tsteps, firing_rate, sim_params)

from MAM2EBRAINS_VISUALIZATION import plot_resting_state
plot_resting_state(M, A, label_spikes, data_path, sim_params)

# area_list = ['V1', 'V2', 'VP', 'V3', 'V3A', 'MT', 'V4t', 'V4', 'VOT', 'MSTd', 'PIP', 'PO', 'DP', 'MIP', 'MDP', 'VIP', 'LIP', 'PITv', 'PITd', 'MSTl', 'CITv', 'CITd', 'FEF', 'TF', 'AITv', 'FST', '7a', 'STPp', 'STPa', '46', 'AITd', 'TH']
# output = {'pdf', 'png', 'eps'}, optional
A.show_rates()