from Model_Network import Neurons, Network, Stimulation, Simulation, SaveNetworkPara
from Model_Network import SaveData, Learning, RunStaticNetwork, RunPlasticNetwork
from Plot_NetResults import HeatmapTest, Plot_Currents2PC, Plot_PopulationRate
import warnings
warnings.filterwarnings("ignore")
dtype = np.float32
# %% Figure 1 - Supplement Figure 3
folder = 'Fig_1_Supp_3'
### Define neuron and network parameters
NeuPar = Neurons(FlagDendRec=np.int32(0))
VE_scale = 0.5
wPP, wPE, wDS = dtype(-0.5), dtype(2.5), dtype(-5.0)
NetPar = Network(NeuPar, wPP=wPP, wPE=wPE, wDS=wDS)
### Define simulation parameters & recording options
SimPar_test = Simulation()
SimPar = Simulation(dt=0.2)
SavePar = SaveData()
### Define plasticity parameter
LearnPar = Learning(NeuPar)
#####################
# Before plasticity #
VE, VP, MP = 1, 1, 1
wSV, wVS, wPP, wEP = -0.6, -0.5, -0.1, -40.0
elif Panel==2: # d
VE, VP, MP = 1, 0, 1
wSV, wVS, wPP, wEP = -0.6, -0.5, -0.1, -40.0
elif Panel==3: # e
VE, VP, MP = 0, 1, 1
wSV, wVS, wPP, wEP = -0.6, -0.5, -1.5, -40.0
elif Panel==4: # f
VE, VP, MP = 0, 0, 1
wSV, wVS, wPP, wEP = -0.6, -0.5, -1.5, -40.0
wPV = -(VP + abs(wSV)*MP - (1-wPP)/(-0.07*wEP) * VE) # gain = 0.07
wPS = -(abs(wVS)*abs(wPV) + (1-wSV*wVS)*MP)
NeuPar = Neurons()
NetPar = Network(NeuPar, wPP=wPP, wPS=wPS, wPV=wPV, wEP=wEP, flag_hetero=0)
### Define input parameters
stim_max, SD = 50.0, 10.0
r0 = np.array([1,2,2,4])
StimPar = Stimulation(NeuPar, NetPar, SD, None, stim_max, r0 = r0, VE=VE, VP=VP, MP=MP, VS=VS, VV=VV)
### Define simulation parameters
SimPar = Simulation()
### Run simulation
RunStaticNetwork(NeuPar, NetPar, StimPar, SimPar, folder, fln)
SaveNetworkPara(NeuPar, NetPar, StimPar, None, folder, fln)