from mpi4py import MPI
import numpy as np
import cProfile

import bifurcation

def main():
  comm = MPI.COMM_WORLD
  rank = comm.Get_rank()
  size = comm.Get_size()


  num_rs = 1000
  iters = 10000
  tail_len = 100
  range_start = 0
  range_end = 4
  range_len = (range_end - range_start) / size

  chunk = num_rs//size
  start = rank * range_len
  end = (rank + 1) * range_len

  rs = np.linspace(start, end, num=chunk, endpoint=False)

  histories = bifurcation.get_bifurcation_map(rs, iters, tail_len)
  histories = np.array(histories)

  recvbuf = None
  if rank == 0:
    recvbuf = np.empty((num_rs, tail_len), 'd')
    # recvbuf = np.empty(size, 'int64')

  comm.Gather(histories, recvbuf, root=0)
  if rank == 0:
    print('Done')
    # bifurcation.plot_bifurcations(np.linspace(range_start, range_end, num_rs), recvbuf)



if __name__ == '__main__':
  main()