"""
netParams.py 

NetPyNE high-level specifications for cortical network model to study epilepsy 

Contributors: salvadordura@gmail.com
"""

from netpyne import specs
import pickle, json
#import math

netParams = specs.NetParams()   # object of class NetParams to store the network parameters

try:
	from __main__ import cfg  # import SimConfig object with params from parent module
except:
	from cfg import cfg


#------------------------------------------------------------------------------
# VERSION 
#------------------------------------------------------------------------------
netParams.version = 1

#------------------------------------------------------------------------------
#
# NETWORK PARAMETERS
#
#------------------------------------------------------------------------------

#------------------------------------------------------------------------------
# General network parameters
#------------------------------------------------------------------------------

netParams.scale = cfg.scale # Scale factor for number of cells # NOT DEFINED YET! 3/11/19 # How is this different than scaleDensity? 
netParams.sizeX = cfg.sizeX # x-dimension (horizontal length) size in um
netParams.sizeY = cfg.sizeY # y-dimension (vertical height or cortical depth) size in um
netParams.sizeZ = cfg.sizeZ # z-dimension (horizontal depth) size in um
netParams.shape = 'cylinder'  # cylindrical (column-like) volume

netParams.defaultDelay = [DT] # 2.0 # default conn delay (ms)
netParams.propVelocity = 5e6 # 500.0 # propagation velocity (um/ms)
netParams.defaultThreshold = 7.0 

#------------------------------------------------------------------------------
# Cell parameters
#------------------------------------------------------------------------------

## Import Pyramidal cell params from py file 
netParams.importCellParams(label='Pyr_simple', conds={'cellType': 'Pyr'}, fileName='PY.hoc', cellName='PY', cellInstance = True)

## Import PV cell params from hoc file 
netParams.importCellParams(label='FS_simple', conds={'cellType': 'FS'}, fileName='FS.hoc', cellName='FS', cellInstance = True)


#------------------------------------------------------------------------------
# Population parameters
#------------------------------------------------------------------------------

## Set layer boundaries and densities
layer = {'2': [0.12, 0.38], '5':[0.38,1.0]}

density = {'L2': 1.0e4, 'L4': 1.0e4, 'L5': 1.0e4} # cells per mm2

# define E and I pops to use later in conn
Epops = ['E2', 'E5']  
# Epops = ['E2']  
Ipops = ['I2']  

   

# specify population parameter 
## L2
netParams.popParams['E2'] = {'cellType': 'Pyr', 'ynormRange': layer['2'],  'numCells':65}     
netParams.popParams['I2'] = {'cellType': 'FS',  'ynormRange': layer['2'],  'numCells':32}      
 
## L5
netParams.popParams['E5'] = {'cellType': 'Pyr', 'ynormRange': layer['5'],   'numCells':206}      



# option to create populations with a single cell (for debugging)
if cfg.singleCellPops:
    for pop in netParams.popParams.values(): pop['numCells'] = 1


#------------------------------------------------------------------------------
# Synaptic mechanism parameters
#------------------------------------------------------------------------------

netParams.synMechParams['AMPA'] = {'mod': 'Exp2Syn', 'tau1': 0.02, 'tau2': 3.0, 'e': 0}
netParams.synMechParams['GABA'] = {'mod': 'Exp2Syn', 'tau1': 0.02, 'tau2': 6.0, 'e': -73}




#------------------------------------------------------------------------------
# Local connectivity parameters
#------------------------------------------------------------------------------
if cfg.addConn:

# EE
    netParams.connParams['E2->E2'] ={
    'preConds':{'pop':'E2'},
    'postConds':{'pop':'E2'},
    'probability':'0.15',
    'weight':1*cfg.EEGain,    
		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['AMPA'],
    'sec':'soma'
    }

    netParams.connParams['E2->E5'] ={
    'preConds':{'pop':'E2'},
    'postConds':{'pop':'E5'},
    'probability':'0.15', 
    'weight':1*cfg.EEGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['AMPA'],
    'sec':'soma'
    }

    netParams.connParams['E5->E2'] ={
    'preConds':{'pop':'E5'},
    'postConds':{'pop':'E2'},
    'probability':'0.15', 
    'weight':1*cfg.EEGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['AMPA'],
    'sec':'soma'
    }

    netParams.connParams['E5->E5'] ={
    'preConds':{'pop':'E5'},
    'postConds':{'pop':'E5'},
    'probability':'0.15', 
    'weight':1*cfg.EEGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['AMPA'],
    'sec':'soma'
    }

## EI

    netParams.connParams['E2->I2'] ={
    'preConds':{'pop':'E2'},
    'postConds':{'pop':'I2'},
    'probability':'0.15', 
    'weight':1*cfg.EIGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
#'delay':defaultDelay,
    'synMech':['AMPA'],
    'sec':'soma'
    }

    netParams.connParams['E5->I2'] ={
    'preConds':{'pop':'E5'},
    'postConds':{'pop':'I2'},
    'probability':'0.15', 
    'weight':1*cfg.EIGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
#'delay':defaultDelay,
    'synMech':['AMPA'],
    'sec':'soma'
    }

### II
    netParams.connParams['I2->I2'] ={
    'preConds':{'pop':'I2'},
    'postConds':{'pop':'I2'},
    'probability':'0.40', 
       'weight':1*cfg.IIGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
#'delay':defaultDelay,
    'synMech':['GABA'],
    'sec':'soma',
    # 'synsPerConn':1
    }



##### IE
    netParams.connParams['I2->E2'] ={
    'preConds':{'pop':'I2'},
    'postConds':{'pop':'E2'},
    'probability':'0.35', 
    #'divergence':500,
    #'weight':0.018*cfg.IEGain,
    'weight':1*cfg.IEGain,
		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['GABA'],
    'sec':'soma',
    # 'synsPerConn':1

    }

    netParams.connParams['I2->E5'] ={
    'preConds':{'pop':'I2'},
    'postConds':{'pop':'E5'},
    'probability':'0.35', 
    #'divergence':500,
    #'weight':0.018*cfg.IEGain,
    'weight':1*cfg.IEGain,

		'delay':'defaultDelay+dist_3D/propVelocity',
    'synMech':['GABA'],
    'sec':'soma',
    # 'synsPerConn':1
    }


  #------------------------------------------------------------------------------
# Reaction-diffusion (RxD) parameters
#------------------------------------------------------------------------------  

# RxD params

### constants
constants = {}
netParams.rxdParams['constants'] = constants



### regions
regions = {}
regions['cyt'] = {'cells': 'all', 'secs': 'all', 'nrn_region':'i'} 
regions['fhs_i'] = {'cells': ['I2'], 'secs': 'all', 'geometry':{'class':'Shell','args':{'lo':1,'hi':1.00727}},'nrn_region':'o'} # interstitial space
regions['fhs_e'] = {'cells': ['E2','E5'], 'secs': 'all', 'geometry':{'class':'Shell','args':{'lo':1,'hi':1.00889}},'nrn_region':'o'} # interstitial space


regions['cyt_fhs_i_membrane'] = {'cells': ['I2'], 'secs': 'all', 'geometry':{'class':'membrane','args':{}}}
regions['cyt_fhs_e_membrane'] = {'cells': ['E2','E5'], 'secs': 'all', 'geometry':{'class':'membrane','args':{}}}

regions['fhs_i_ecs_membrane'] = {'cells': ['I2'], 'secs': 'all', 'geometry': {'class': 'ScalableBorder', 'args': {'diam_scale': 1.00727, 'on_cell_surface': True}}}
regions['fhs_e_ecs_membrane'] = {'cells': ['E2','E5'], 'secs': 'all', 'geometry': {'class': 'ScalableBorder', 'args': {'diam_scale': 1.00889, 'on_cell_surface': True}}}




fraction = 1.0
margin = 10  # extracellular volume additional margin 
x = [0 - margin, cfg.sizeX*fraction + margin]
y = [-cfg.sizeY*fraction - margin, 0 + margin]
z = [0 - margin, cfg.sizeZ*fraction + margin]

alpha = 0.1
dx = 50 #um
regions['ecs'] = {'extracellular': True, 'xlo': x[0], 'ylo': y[0], 'zlo': z[0], 'xhi': x[1], 'yhi': y[1], 'zhi': z[1], 'dx': dx, 'volume_fraction': alpha, 'tortuosity': 1.5}  
netParams.rxdParams['regions'] = regions


### species 
species = {}
diff_const = [MCRATE] #1.0 um^2/ms
# species['k'] = {'regions': ['fhs','cyt'], 'd': 0.0, 'charge': 1, 'initial': 'lambda nd: 140.0*mM if nd.region == cyt else 3.0*mM'}
# species['k'] = {'regions': ['fhs','cyt'], 'd': 0.0, 'charge': 1, 'initial': lambda nd: 140.0 if nd.region.name == 'cyt' else 3.0}
species['k'] = {'regions': ['fhs_i','fhs_e','cyt'], 'd': 0.0, 'charge': 1, 'initial': lambda nd: 138 if nd.region.name == 'cyt' else [KOINIT]}



species['buf_i'] = {'regions':['fhs_i'], 'd':0.0, 'charge':0, 'initial': [BINITI]}
species['kbuf_i'] = {'regions':['fhs_i'], 'd':0.0, 'charge':0, 'initial': 500-[BINITI]}

species['buf_e'] = {'regions':['fhs_e'], 'd':0.0, 'charge':0, 'initial': [BINITE]}
species['kbuf_e'] = {'regions':['fhs_e'], 'd':0.0, 'charge':0, 'initial': 500-[BINITE]}

species['garbage_i'] = {'regions':['fhs_i'], 'd':0.0, 'charge':0, 'initial': 0}
species['garbage_e'] = {'regions':['fhs_e'], 'd':0.0, 'charge':0, 'initial': 0}

species['extP'] = {'regions': ['ecs'], 'd': diff_const, 'charge': 1, 'initial': [KOINIT]}


netParams.rxdParams['species'] = species


### reactions
reactions = {}
reactions['buffering_i'] = {'reactant':'k+buf_i','product':'kbuf_i','rate_f': '0.0008*k*buf_i / (1 + rxd.rxdmath.exp(((k - [KOTHI]) / [SLOPEI])))','rate_b':'0.0008*(500.0-buf_i)', 'custom_dynamics':True}
reactions['buffering_e'] = {'reactant':'k+buf_e','product':'kbuf_e','rate_f': '0.0008*k*buf_e / (1 + rxd.rxdmath.exp(((k - [KOTHE]) / [SLOPEE])))','rate_b':'0.0008*(500.0-buf_e)', 'custom_dynamics':True}
reactions['potassium_garbage_i'] = {'reactant':'k','product':'garbage_i','rate_f': '0.0909*0.0008*(500.0-buf_i)','rate_b': 0, 'custom_dynamics':True}
reactions['potassium_garbage_e'] = {'reactant':'k','product':'garbage_e','rate_f': '0.0909*0.0008*(500.0-buf_e)','rate_b':0, 'custom_dynamics':True}

netParams.rxdParams['reactions'] = reactions

mcReactions = {}
mM_to_mol_per_um = 6.0221409e+23 * 1e-18

mcReactions['exchangeK_i'] = {'reactant': 'k[fhs_i]', 'product': 'extP[ecs]', 'rate_f':mM_to_mol_per_um*diff_const*alpha/(0.5*dx) , 'rate_b':mM_to_mol_per_um*diff_const*alpha/(0.5*dx)  ,'membrane': 'fhs_i_ecs_membrane'}
mcReactions['exchangeK_e'] = {'reactant': 'k[fhs_e]', 'product': 'extP[ecs]', 'rate_f':mM_to_mol_per_um*diff_const*alpha/(0.5*dx) , 'rate_b':mM_to_mol_per_um*diff_const*alpha/(0.5*dx)  ,'membrane': 'fhs_e_ecs_membrane'}




netParams.rxdParams['multicompartmentReactions'] = mcReactions


# Rates

rates = {}
rates['buffer_extra_contri_i'] = {'species': 'k[fhs_i]', 'rate': '-10*0.0909*0.0008*(500.0-buf_i[fhs_i])'}
rates['buffer_extra_contri_e'] = {'species': 'k[fhs_e]', 'rate': '-10*0.0909*0.0008*(500.0-buf_e[fhs_e])'}
#netParams.rxdParams['rates'] = rates # rates don't work with netpyne

# Stimulation parameters


netParams.stimSourceParams['bkg_e'] = {'type': 'NetStim', 'rate': [BCKRATEE], 'noise': [BCKNOISEE], 'seed':[EXSEED]} 
netParams.stimSourceParams['bkg_i'] = {'type': 'NetStim', 'rate': [BCKRATEI], 'noise': [BCKNOISEI], 'seed':[INSEED]}


netParams.stimTargetParams['bkg_e->Pyr'] = {'source': 'bkg_e', 'conds': {'cellType': 'Pyr'}, 'weight': [BCKWTEE], 'delay': 0, 'synMech': 'AMPA'}
netParams.stimTargetParams['bkg_e->FS'] = {'source': 'bkg_e', 'conds': {'cellType': 'FS'}, 'weight': [BCKWTEI], 'delay': 0, 'synMech': 'AMPA'}

netParams.stimTargetParams['bkg_i->Pyr'] = {'source': 'bkg_i', 'conds': {'cellType': 'Pyr'}, 'weight': [BCKWTIE], 'delay': 0, 'synMech': 'GABA'}
netParams.stimTargetParams['bkg_i->FS'] = {'source': 'bkg_i', 'conds': {'cellType': 'FS'}, 'weight': [BCKWTII], 'delay': 0, 'synMech': 'GABA'}

#------------------------------------------------------------------------------
# Description
#------------------------------------------------------------------------------

netParams.description = """
v1 - Created simple 4-layer cortical model based Allen Brain V1 params
"""