from mpi4py import MPI
import ptreelib

from pNbody import *

usempi = 1

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

if usempi:
  nb.Tree = ptreelib.Tree(npart=nb.npart,pos=nb.pos,vel=nb.vel,mass=nb.mass,num=nb.num,tpe=nb.tpe,params={"OutputInfo":1})
  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)
  # --> donne -406.29 !!!
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)
  # --> donne -0.19275261



# create the grid
nx = 256
ny = 256
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()

# potential
#pot = nb.Tree.Potential(gr.pos,1.0)
#pot.shape = (nx,ny)

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

# density
#hsml = gr.mass/gr.mass*1
#if usempi:
#  den,hsml = nb.Tree.Density(gr.pos,hsml)
#else:  
#  den,hsml = nb.Tree.Density(gr.pos,hsml,33,3)
#den.shape = (nx,ny)
#hsml.shape = (nx,ny)

# observable
hsml = gr.mass/gr.mass*1
observable = nb.pos[:,0]
if usempi:
  den,hsml = nb.Tree.Density(gr.pos,hsml)
  observable = nb.Tree.SphEvaluate(gr.pos,hsml,nb.rho,observable)
else:  
  den,hsml = nb.Tree.Density(gr.pos,hsml,33,3)
  observable = nb.Tree.SphEvaluate(gr.pos,hsml,nb.rho,observable,33,3)

observable.shape = (nx,ny)


mat = observable

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

