from mpi4py import MPI
import ptreelib

from pNbody import *

usempi = 1

# open disk and create tree
nb = Nbody('../../examples/disk.dat',ftype='gadget')

if usempi:
  nb.Tree = ptreelib.Tree(npart=nb.npart,pos=nb.pos,vel=nb.vel,mass=nb.mass,num=nb.num,tpe=nb.tpe)
  nb.ExchangeParticles()
  nb.Tree.SetParticles(nb.npart,nb.pos,nb.vel,nb.mass,nb.num,nb.tpe) 
  nb.Tree.BuildTree()
  #print nb.Tree.Potential(array([[-20.4709, -0.840655, -19.8087]],float32),0.1)
else:  
  nb.Tree = treelib.Tree(npart=nb.npart,pos=nb.pos,vel=nb.vel,mass=nb.mass)
  #print nb.Tree.Potential(array([[-20.4709, -0.840655, -19.8087]],float32),0.1)


#print nb.Tree.Potential(array([[-10,-10,0]],float32),0.1)



# create the grid
nx = 128
ny = 128
xmin = -100
xmax= 100
ymin = -100
ymax =  100
pos = libgrid.get_Points_On_Carthesian_2d_Grid(nx,ny,xmin,xmax,ymin,ymax,offx=0,offy=0)
gr = Nbody(pos=pos,status='new',ftype='gadget')		# this should be //
# rotate it
gr.rotate(axis="x",angle=pi/2)
gr.rename('grid.dat')
gr.write()

pot = nb.Tree.Potential(gr.pos,1.0)
pot.shape = (nx,ny)
#acc = nb.Tree.Acceleration(gr.pos,1.0)
#acc = copy.deepcopy(acc[:,0])
#acc.shape = (nx,ny)

mat = pot

if mpi.mpi_IsMaster():
  libutil.mplot(mat,save='map.fits')

