from netpyne import specs, sim
import os

# Information
hoc_file_EC = 'ECtemplate.hoc'
cell_name_EC = 'ECcell'
cell_label_EC = 'Pyr_EC'
cell_type_EC = 'PYR'
cell_model_EC = 'PYR'

hoc_file_CA3 = 'CA3_Cell.hoc'  # New HOC file
cell_name_CA3 = 'CA3cell'      # New cell name
cell_label_CA3 = 'CA3'
cell_type_CA3 = 'PYR'     
cell_model_CA3 = 'PYR'   

output_dir = 'output'
if not os.path.exists(output_dir):
    os.mkdir(output_dir)

###############################################################################
netParams = specs.NetParams()

pop_label_EC = cell_label_EC + '_pop'
netParams.popParams[pop_label_EC] = {'cellType': cell_type_EC,
                                     'cellModel': cell_model_EC,
                                     'numCells': 20}  

importedCellParams_EC = netParams.importCellParams(label=cell_label_EC,
                                                   conds={'cellType': cell_type_EC, 'cellModel': cell_model_EC},
                                                   fileName=hoc_file_EC,
                                                   cellName=cell_name_EC)

pop_label_CA3 = cell_label_CA3 + '_pop'
netParams.popParams[pop_label_CA3] = {'cellType': cell_type_CA3,
                                      'cellModel': cell_model_CA3,
                                      'numCells': 30} 

# Import CA3 cells into cellParams
importedCellParams_CA3 = netParams.importCellParams(label=cell_label_CA3,
                                                    conds={'cellType': cell_type_CA3, 'cellModel': cell_model_CA3},
                                                    fileName=hoc_file_CA3,
                                                    cellName=cell_name_CA3)

###############################################################################
# NetPyNE config
simConfig = specs.SimConfig()					           
simConfig.duration = 200 						          
simConfig.dt = 0.01								               
simConfig.verbose = True							              
simConfig.recordTraces = {'V_soma':{'sec':'soma','loc':0.5,'var':'v'}}  
simConfig.recordStep = 0.01
simConfig.filename = os.path.join(output_dir, cell_name_EC) 			
simConfig.saveJson = False
simConfig.analysis['plotTraces'] = {'include': ['all'], 'saveFig': True}  # Plot recorded traces for this list of cells
simConfig.hParams['v_init'] = -70

### Run simulation ###
sim.createSimulateAnalyze(netParams=netParams, simConfig=simConfig)

### Plot morphology for EC cells ###
sim.analysis.plotShape(showSyns=True, dist=0.6, saveFig=True)

# Update simConfig for CA3 cells
simConfig.filename = os.path.join(output_dir, cell_name_CA3)  # Set file output name for CA3 cells

### Run simulation for CA3 cells ###
sim.createSimulateAnalyze(netParams=netParams, simConfig=simConfig)

### Plot morphology for CA3 cells ###
sim.analysis.plotShape(showSyns=True, dist=0.6, saveFig=True)