import nest
import sys, os
import numpy as np
from mpi4py import MPI
from time import time

nest.set_verbosity("M_ERROR")
nest.ResetKernel()

comm = MPI.COMM_WORLD
RNK = comm.rank
MOTHER = 0



TNVP = int(sys.argv[1])
GRNG = int(sys.argv[2])

for ii in range(10):
    # print(f"\n\n====={ii}=====\n\n")
    nest.SetKernelStatus({
        "grng_seed":GRNG,
        "rng_seeds":range(GRNG+TNVP+1, GRNG+2*TNVP+1) 
        })



    n  = nest.Create("iaf_psc_delta", 100)
    PG = nest.Create("poisson_generator", 1, {"rate":5e5})
    SD = nest.Create("spike_detector")

    nest.Connect(n, n, {"rule":"pairwise_bernoulli", "p":0.25})
    nest.Connect(PG, n)
    nest.Connect(n, SD)

    t0=time()
    nest.Simulate(1e4)
    t1=time()


    senders = nest.GetStatus(SD, "events")[0]["senders"]




    rec_buff = comm.gather(senders, root=MOTHER)
    if RNK==MOTHER:
        print(type(rec_buff))
        # print(f"{ii}:\tMOTHER received {rec_buff}")

        pass

    if RNK==MOTHER:
        np.save(f"sender_{ii}_{nest.Rank()}.npy", np.ravel(rec_buff))
    comm.Barrier()
    print(f"{RNK}@{ii} says: EXIT 0\t|time taken :{t1-t0}")
    if RNK==MOTHER:
        print("\n")

    nest.SetStatus(SD, {"n_events":0})