"""
This is the second bit of code for lab1.  
Here we're going to explore voltage a bit more. 
Step 5, run this printPklSpikes, and look at the png file.
"""

#import the cod for the pyNN nest interface.
import pyNN.nest as sim
from pyNN.utility import normalized_filename

#set up some constants in all caps for my style.
RUN_DURATION = 100
TIME_STEP = 1.0
#TIME_STEP = 0.1

numberNeurons = 2

def init():
    sim.setup(timestep=TIME_STEP,min_delay=TIME_STEP, max_delay=TIME_STEP)

def makeNeurons():
    #cells=sim.Population(numberNeurons,sim.IF_cond_exp,  cellparams={})
    #cells=sim.Population(numberNeurons,sim.IF_cond_alpha,  cellparams={})
    #step 7
    cells=sim.Population(numberNeurons,sim.IF_cond_exp,cellparams={'v_reset': -70.0})
    cells.record({'spikes','v'})
    return (cells)

#make a synapse from a spike source to a neuron,
#leaving the other neuron alone.
def makeSynapses(cells):
    #Make a spike source
    inputSpikeTimes = [10.0,20.0,40.0,70]
    spikeArray = {'spike_times': [inputSpikeTimes]}
    spikeGen=sim.Population(1,sim.SpikeSourceArray,spikeArray,
                            label='inputSpikes_0')

    #make a synapse
    firingWt = 0.2
    connector = []
    connector = connector + [(0,0,firingWt,TIME_STEP)]
    fromListConnector = sim.FromListConnector(connector)
    sim.Projection(spikeGen, cells,  fromListConnector, receptor_type='excitatory')

def plotIt(population):
    from pyNN.utility.plotting import Figure, Panel
    figure_filename = normalized_filename(".", "lab1a voltage",
                                          "png", "nest")
    panels = []
    for variable in ('v'):
        panels.append(
            Panel(population.get_data().segments[0].filter(name=variable)[0],
                  data_labels=[population.label], yticks=True),
            )
    # add ylabel to top panel in each group
    panels[0].options.update(ylabel='Membrane potential (mV)')
    # add xticks and xlabel to final panel
    panels[-1].options.update(xticks=True, xlabel="Time (ms)")

    Figure(*panels,
           title="Example of IF_cond_exp neuron",
           annotations="Simulated with NEST" 
    ).save(figure_filename)
    print(figure_filename)
    

#----main---
init()

cells = makeNeurons() 

makeSynapses(cells)

sim.run(RUN_DURATION)

cells.write_data("lab1aSpikes.pkl",'spikes')
cells.write_data("lab1aVoltage.pkl",'v')

plotIt(cells)