from neuromllite import Network, Cell, InputSource, Population, Synapse
from neuromllite import Projection, RandomConnectivity, Input, Simulation
import sys

################################################################################
###   Build new network

net = Network(id="HHTest")
net.notes = "Example HH cell"
net.parameters = {"N": 1}


cell = Cell(id="hhcell", neuroml2_source_file="../hhcell.cell.nml")
cell.parameters = {}

params = {
}


for p in params:
    cell.parameters[p] = p
    net.parameters[p] = params[p]

net.cells.append(cell)


pop = Population(
    id="hhPop", size="1", component=cell.id, properties={"color": ".7 0 0"}
)
net.populations.append(pop)

net.parameters["delay"] = "5ms"
net.parameters["stim_amp"] = "0.05nA"
net.parameters["duration"] = "25ms"
input_source = InputSource(
    id="iclamp_0",
    neuroml2_input="pulseGenerator",
    parameters={"amplitude": "stim_amp", "delay": "delay", "duration": "duration"},
)
net.input_sources.append(input_source)

net.inputs.append(
    Input(
        id="stim",
        input_source=input_source.id,
        population=pop.id,
        percentage=100,
        weight=1,
    )
)

print(net)
print(net.to_json())
net_yaml_file = net.to_yaml_file("%s.nmllite.yaml" % net.id)
net_json_file = net.to_json_file("%s.nmllite.json" % net.id)


################################################################################
###   Build Simulation object & save as JSON

record_variables = {"v": {"all": "*"}}

sim = Simulation(
    id="Sim%s" % net.id,
    network=net_yaml_file,
    duration="50",
    dt="0.025",
    record_variables=record_variables,
)

sim.to_yaml_file("%s.yaml" % sim.id)
sim.network = net_json_file
sim.to_json_file("%s.json" % sim.id)


################################################################################
###   Run in some simulators

from neuromllite.NetworkGenerator import check_to_generate_or_run
import sys

check_to_generate_or_run(sys.argv, sim)