from mpi4py import MPI
import ptreelib

from pNbody import *

nb = Nbody('../../examples/disk.dat',ftype='gadget')
Tree = ptreelib.Tree(nb.npart,nb.pos,nb.vel,nb.mass,nb.num,nb.tpe)
#Tree.info()


# here, we should exchange particles (to do in pNbody, hard work ?)
num,procs = Tree.GetExchanges()


if mpi.mpi_ThisTask() == 0:
  io.write_array('qq',procs)


H =  histogram(procs,arange(mpi.mpi_NTask()))  
T = mpi.mpi_AllgatherAndConcatArray(H) 
T.shape = (mpi.mpi_NTask(),mpi.mpi_NTask())


if mpi.mpi_IsMaster():  
  pass
  #print T
  #Task = 2
  #ToTask = 0
  #print T[mpi.mpi_NTask()-1-Task,ToTask]
  
# now, we have to send / recv  
SendPart = T[mpi.mpi_NTask()-1-mpi.mpi_ThisTask(),:]
RecvPart = T[:,mpi.mpi_ThisTask()]

if mpi.mpi_IsMaster():  
  pass
  #print SendPart
  #print RecvPart


NTask = mpi.mpi_NTask()
ThisTask = mpi.mpi_ThisTask()


new_procs = array([])
new_pos   = array([])
new_pos.shape = (0,3)

  

# send/recv (1)
for i in range(NTask):  
  
  if   i > ThisTask:    
     # here, we send to i 
     
     N = T[NTask-1-ThisTask,i]
     mpi.mpi_Send(N,i)
     #print "%d send to %d %d"%(ThisTask,i,N)
     
     if N>0:
       to_procs  = compress((procs==i),procs)
       to_num	 = compress((procs==i),num)
       to_pos	 = libutil.compress_from_lst(nb.pos,nb.num,to_num)
       #mpi.mpi_Send(to_procs,i)
       mpi.mpi_Send(to_pos,i)
       
     nb.pos = libutil.compress_from_lst(nb.pos,nb.num,to_num,reject=True)
     nb.num = libutil.compress_from_lst(nb.num,nb.num,to_num,reject=True)
    
       
  elif i < ThisTask:
     N = mpi.mpi_Recv(i)
     #print "%d recv fr %d %d"%(ThisTask,i,N) 
     if N > 0:
       #new_procs = concatenate((new_procs,mpi.mpi_Recv(i))) 
       new_pos   = concatenate((new_pos,mpi.mpi_Recv(i)))


# send/recv (1)
for i in range(NTask):  
  
  if   i < ThisTask:    
     # here, we send to i 
     
     N = T[NTask-1-ThisTask,i]
     mpi.mpi_Send(N,i)
     #print "%d send to %d %d"%(ThisTask,i,N)
     
     if N>0:
       to_procs  = compress((procs==i),procs)
       to_num	 = compress((procs==i),num)
       to_pos	 = libutil.compress_from_lst(nb.pos,nb.num,to_num)
       #mpi.mpi_Send(to_procs,i)
       mpi.mpi_Send(to_pos,i)

     nb.pos = libutil.compress_from_lst(nb.pos,nb.num,to_num,reject=True)
     nb.num = libutil.compress_from_lst(nb.num,nb.num,to_num,reject=True)
       
  elif i > ThisTask:
     N = mpi.mpi_Recv(i)
     #print "%d recv fr %d %d"%(ThisTask,i,N) 
     if N > 0:
       #new_procs = concatenate((new_procs,mpi.mpi_Recv(i))) 
       new_pos   = concatenate((new_pos,mpi.mpi_Recv(i)))
     
    
     
# check
c = (new_procs!=ThisTask).astype(Int)
if sum(c)>0:
  print "here we are in trouble"

nb.pos = concatenate((nb.pos,new_pos))

nb.vel = None
nb.mass= None
nb.num = None
nb.u   = None
nb.rho = None
nb.rsp = None
nb.opt = None 
nb.npart = array([len(nb.pos),0,0,0,0,0])
nb.npart_tot = None
nb.init()

#nb.info()

nb.set_pio('yes')
nb.write()


#f = open("o%d.dat"%(ThisTask),'w')
#for i in range(len(nb.pos)):
#  f.write("%8.5f\n"%nb.pos[i,0])
#  
#f.close()  


######################################

print T
print "-----------------------"

Tree = ptreelib.Tree(nb.npart,nb.pos,nb.vel,nb.mass,nb.num)
#Tree.info()


# here, we should exchange particles (to do in pNbody, hard work ?)
num,procs = Tree.GetExchanges()


if mpi.mpi_ThisTask() == 0:
  io.write_array('qq',procs)


H =  histogram(procs,arange(mpi.mpi_NTask()))  
T = mpi.mpi_AllgatherAndConcatArray(H) 
T.shape = (mpi.mpi_NTask(),mpi.mpi_NTask())


print T
