# The use case:
#     1. Create a spatial 2D-Population of Neurons
#     2. "randomly" select some of the Neurons as source neurons
#     3. for each source neuron
#            establish connections from the source neuron into the neurons population using a mask

# Problem occurs when a mask is used inside the connection-dict given to the nest.Connect(...) function
# If the mask is removed, no error is produced (however the result is not as desired)

import nest
import random
from mpi4py import MPI

nest.ResetKernel()
nest.set_verbosity('M_ERROR')
num_threads = 1 # number of threads makes no difference to the error
nest.local_num_threads = num_threads
num_processes = nest.num_processes
num_vp = nest.total_num_virtual_procs
random.seed(42)

'''
####################################################################
# Parameters (and explanation)


# The error occurs if the mask is used inside the connection dict of nest.Connect (use_mask = True)
# If no mask is used in the cdict (use_mask = False) no error is produced
use_mask = True

######
# The following values influence whether the error occurs or not, with masks activated.

# has a big influence
neurons_num = 20  # Total number of neurons

extent = 2.
edge_wrap = True  # problem occurs with edge_wrap = True  and with edge_wrap = False, makes no difference

# value makes no real difference, just that the error occurs at different settings (neurons_num, extent,...)
mask_radius = 0.5

# Number of randomly selected source neurons, connections are established only from these neurons.
# The denominator influences at which settings of num_neurons, extent, mask_radius and source_neurons_num the error occurs,
#   but it still occurs
# Lower denominator value -> fewer settings run successfully
source_neurons_num = int(neurons_num / 3)
'''
############################################################
## Setting examples
# For each setting there are multiple examples for alternative values which lead to the same outcome
#   Only one value should be changed at a time


# Setting 1: creates error with 3 and 4 MPI processes
# neurons_num = 20
# extent = 2.                     #also with x \in {1., 1.5, 1.8, 2., 2.5} default 2.
# mask_radius = 0.5                  #also with x \in {0.2, 0.5, 0.6, 0.7, 1.}   default 0.5
# source_neurons_num = int(neurons_num/2)      # also with neurons_num/x, x \in {1, 2,  3, 4, 5, 6, 10, 20, 1/2, 1/3}  default 2
# use_mask = True
# edge_wrap = True

# Setting 1.1: works with 3 and 4 MPI processes
neurons_num = 20
extent = 2.                   #also with x \in {1., 1.5, 1.8, 2., 2.5} default 2.
mask_radius = 0.5                   #also with x \in {0.2, 0.5, 0.6, 0.7, 1.}  default 0.5
source_neurons_num = int(neurons_num/2)     # also with neurons_num/x, x \in {1, 2, 3, 4, 5, 6, 10, 20, 1/2, 1/3} default 2
use_mask = False           # difference to setting 1
edge_wrap = True


# Setting 2: works with 3 and 4 MPI processes
# neurons_num = 15           # difference to setting 1 which didn't work
# extent = 2.                    # also with x \in {1., 1.5, 1.8, 2., 2.5, 5., 10., 15.} default 2.
# mask_radius = 0.5                    # also with x \in {0.2,0.5, 0.6, 0.7, 1.}   default 0.5
# source_neurons_num = int(neurons_num/2)     # also with neurons_num/x, x \in {1, 2, 3, 4, 5, 6, 10, 20, 1/2, 1/3} default 2
# use_mask = True
# edge_wrap = True


# Setting 3: creates error with 3 and 4 MPI processes
# neurons_num = 20  # also x \in {20, 23, 50} default 20
# extent = 1.
# mask_radius = 0.2
# source_neurons_num = int(neurons_num/2)   #also with neurons_num/x, x\in{1, 2}
# use_mask = True
# edge_wrap = True

# Setting 3.1: works with 3 and 4 MPI processes
# neurons_num = 20               # also x \in {20, 23, 50} default 20
# extent = 1.
# mask_radius = 0.2
# source_neurons_num = int(neurons_num/2)        #also with neurons_num/x, x \in {1, 2} default 2
# use_mask = False               # difference to setting 3
# edge_wrap = True


# Setting 4: works with 4 MPI processes but creates error with 3 MPI processes
# neurons_num = 10    # also with x \in {10, 22, 25}  default 10
# extent = 1.
# mask_radius = 0.2
# source_neurons_num = int(neurons_num/1)
# use_mask = True
# edge_wrap = True


# Setting 5: works with 3 and 4 MPI processes
# neurons_num = 24       # also x \in {12, 21, 24, 48, 51} default 24
# extent = 1.
# mask_radius = 0.2
# source_neurons_num = int(neurons_num/1)
# use_mask = True
# edge_wrap = True


###################################################################
# Mask

mask = nest.CreateMask(masktype='circular', specs={'radius': mask_radius})


#####################################################################
# Create Neurons

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

neurons = nest.Create(model='iaf_psc_alpha', n=neurons_num, positions=neurons_pos)
neurons_list = neurons.tolist()   # used for implementational reasons


##################################################################
# print info

if nest.Rank() == 0:
    print("-----------------------------------------------------")
    print("num threads :", num_threads)
    print("num of MPI processes:", num_processes)
    print("total num virtual processes :", num_vp)
    print("extent:", extent)
    print("neurons_num:", neurons_num)
    print("edge_wrap:", edge_wrap)
    print('mask_radius:', mask_radius)
    print("-----------------------------------------------------")


######################################################################
# Create connections

# taking every x-th neuron as source neurons (same behaviour as randomly selecting a list of source neurons)
source_node_ids = neurons_list[::3]


if use_mask:  # cdict with mask
    cdict = {'rule': 'pairwise_bernoulli',
             'p': nest.random.uniform(),
             'allow_autapses': False,
             'allow_multapses': False,
             'mask': mask   # de-/activating this line shows the problem
             }
else:  # cdict without mask line
    cdict = {'rule': 'pairwise_bernoulli',
             'p': nest.random.uniform(),
             'allow_autapses': False,
             'allow_multapses': False
             }


# connect the source neurons to neurons by iterating over the global_ids of nest
for node_id in source_node_ids:
    # The following is implemented this way because of my use-case where,
    #   the spatial information of the source neuron is required in order to connect.
    # Therefore, the NodeCollection needs to be sliced to create a one element NodeCollection containing
    #     the spatial information.
    # It is not possible to obtain the spacial information of a neuron from a NodeCollection
    #   created by nest.SelectNodesByMask or by slicing the NodeCollection with an interval (nodecollection[::i]).
    #   In both cases the metadata containing the spatial information is lost.
    source_node_index = neurons.index(node_id)
    source_node = neurons[source_node_index]
    nest.Connect(source_node, neurons, cdict)  # in this line the problem occurs if the mask is used in the cdict


# no difference if this barrier is removed
# serves only to beautify the output
MPI.COMM_WORLD.Barrier()

if nest.Rank() == 0:
    print("-----------------------------------------------------")
    print("num threads :", num_threads)
    print("num of MPI processes:", num_processes)
    print("total num virtual processes :", num_vp)
    print("extent:", extent)
    print("neurons_num:", neurons_num)
    print("edge_wrap:", edge_wrap)
    print('mask_radius:', mask_radius)
    print("-----------------------------------------------------")