import nest
import random
import matplotlib.pyplot as plt

nest.ResetKernel()

extent = 10
num_neurons = 200

mask_obj = nest.CreateMask('doughnut',
                           {'inner_radius': 0.75,
                            'outer_radius': 3.75})

pos = nest.spatial.free(
                          pos=nest.random.uniform(min=-extent / 2,
                                                  max=extent / 2),
                          num_dimensions=2,
                          edge_wrap=False)

neurons = nest.Create('iaf_psc_alpha', num_neurons, positions=pos)


ctr = nest.FindCenterElement(neurons)
ctr_position = nest.GetPosition(ctr)

#############################
# Problem 1:
# The NodeCollection created by the nest.SelectNodesByMask() function

nodes_in_mask = nest.SelectNodesByMask(neurons, ctr_position, mask_obj)

# Try to get positions of new NodeCollection
# -> Error #1  nest.lib.hl_api_exceptions.LayerExpected: LayerExpected in SLI function GetPosition_g:
nest.GetPosition(nodes_in_mask)

# Try to plot new NodeCollection
# -> Error #2   TypeError: 'NoneType' object is not subscriptable
fig = nest.PlotLayer(nodes_in_mask)

plt.show()

##########################
# Problem 2
# Create a NodeCollection "new_collection" containing 3 random nodes from "neurons"

node_ids = []

neurons_list = neurons.tolist()

# get ids of 3 random nodes from NodeCollection
for x in range(0, 3):
    center = random.choice(neurons_list)
    node_ids.append(center)

# sort because ids need to be sorted to create NodeCollection
node_ids.sort()

# create NodeCollection from drawn node ids
new_collection = nest.NodeCollection(node_ids)

# Try to get positions of new NodeCollection
# -> Error #1  nest.lib.hl_api_exceptions.LayerExpected: LayerExpected in SLI function GetPosition_g:
nest.GetPosition(new_collection)

# Try to plot new NodeCollection
# -> Error #2  TypeError: 'NoneType' object is not subscriptable
fig = nest.PlotLayer(new_collection)


plt.show()