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

output = 221

fn = 'output_%04i'%(output)

pf = load(fn)
vl=pf["RefineRegionLeftEdge"]
vr=pf["RefineRegionRightEdge"]
sv = pf.h.region((vl+vr)/2, vl, vr)

sm = sv["ParticleMassMsun"]
ct = sv["creation_time"]
pt = sv["particle_type"]
px = sv["particle_position_x"]
py = sv["particle_position_y"]
pz = sv["particle_position_z"]
# First pick out only PopIII stars.
stars = np.logical_and(ct > 0, pt == 1)
ct = ct[stars]
sm = sm[stars]
px = px[stars]
py = py[stars]
pz = pz[stars]

stars = np.logical_and(sm<300, sm>5)
ct = ct[stars]
sm = sm[stars]
px = px[stars]
py = py[stars]
pz = pz[stars]


f=open('./PopIII_list_%04i_type1.txt'%output,'a')
f.write('#  Mass CreationTime Position (x,y,z) ')
for i in xrange(ct.size):
    f.write(str(i)+"  "+str(sm[i])+"  "+str(ct[i])+"  "+str(px[i])+"  "+str(py[i])+"  "+str(pz[i])+"\n")

f.close()