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

pf = load("RD0028/RedshiftOutput0028")

# Find all the haloes, and include star particles.
#sv = pf.h.region([0.5,0.5,0.5],[0.4,0.4,0.4],[0.6,0.6,0.6])
#haloes = parallelHF(pf, dm_only=False, subvolume = sv)
#vl=pf["RefineRegionLeftEdge"]
#vr=pf["RefineRegionRightEdge"]
#sv = pf.h.region((vl+vr)/2, vl, vr)

haloes = LoadHaloes(pf, "RD0028/RD0028")

# Set up the spectrum builder.
spec = SpectrumBuilder(pf, bcdir="/u/sciteam/haox/yt_scripts/BC03", model="salpeter")

# Iterate over the haloes.
i=0
for halo in haloes:
    # Get the pertinent arrays.
    if i <= 53: 
       i=i+1
       continue
    i=i+1
    ct = halo["creation_time"]
    sm = halo["ParticleMassMsun"]
    metal = halo["metallicity_fraction"]
    pt = halo["particle_type"]
    # Select just the stars.
    stars = (ct > 0)  #np.logical_and(ct > 0, pt == 7)
    ct = ct[stars]
    sm = sm[stars]
    metal = metal[stars]
    # Calculate the spectrum.
    spec.calculate_spectrum(star_mass=sm, star_creation_time=ct,
    star_metallicity_fraction=metal)
    
    spec.write_out(name="RD0028/SED/halo%05d.out" % halo.id)
    # Write out the SED using the default flux normalization.
    #spec.write_out_SED(name="RD0066/halo%05d.out" % halo.id)