import sys
from yt.mods import *
import numpy as np
from yt.analysis_modules.halo_finding.api import *
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

output = 2

n=2500
#fn = 'output_%04i'%(output)
fn = 'RedshiftOutput%04i'%output

pf = load(fn)
haloes = LoadHaloes(pf, "RD%04i"%output)
#vl=pf["RefineRegionLeftEdge"]
#vr=pf["RefineRegionRightEdge"]
#sv = pf.h.region((vl+vr)/2, vl, vr)
f=open('./PopII_haloes_%04i_%04i.txt'%(output,rank),'a')

i=0
for halo in haloes:
    i = i+1 
    if i <= n*rank: continue
    if i > n*(rank+1): break
    if halo.total_mass() <= 1e6: break
    hm = halo.total_mass()
    hr = halo.maximum_radius()
    hc = halo.center_of_mass()
    sp = halo.get_sphere()
    sm = sp["ParticleMassMsun"]
    ct = sp["creation_time"]
    pt = sp["particle_type"]
    pn = sp["particle_index"]
    px = sp["particle_position_x"]
    py = sp["particle_position_y"]
    pz = sp["particle_position_z"]
    # First pick out only PopIII stars.
    stars = np.logical_and(ct > 0, pt == 7)
    ct = ct[stars]
    sm = sm[stars]
    pn = pn[stars]
    px = px[stars]
    py = py[stars]
    pz = pz[stars]
    
    f.write(str(i-1)+" "+str(hm)+" "+str(hr)+" "+str(hc[0])+" "+str(hc[1])+" "+str(hc[2])+":\n")
    #f.write("   "+str(ct.size)+"   "+str(smp.mean())+"\n")
    f.write("   "+str(ct.size)+"   "+str(sm.mean())+"\n")
    for j in xrange(ct.size):
       f.write("   "+str(j)+" "+str(sm[j])+"  "+str(ct[j])+"  "+str(pn[j])+"  "+str(px[j])+"  "+str(py[j])+"  "+str(pz[j])+"\n")
    f.write("\n\n")
    #i=i+1

f.close()