from yt.mods import *
from yt.analysis_modules.halo_profiler.api import *
from yt.analysis_modules.halo_finding.api import *

output = 15
#hp = HaloProfiler('output_%04i'%output,output_dir="Halos",
#     halo_list_file='DD0108.out',n_profile_bins=100)
hp = HaloProfiler('RedshiftOutput%04i'%output,output_dir="Halos",
     halo_list_file='RD%04i.out'%output,n_profile_bins=100)
#     ,recenter='Max_Gas_Density' )


hp.add_halo_filter(VirialFilter, must_be_virialized=True,
                   overdensity_field='ActualOverdensity',
                   virial_overdensity=200,
                   virial_filters=[['TotalMassMsun','>=','1e6']],
                   virial_quantities=['TotalMassMsun','RadiusMpc'],
                   use_log=True)

hp.make_profiles(njobs=-1, prefilters=["halo['mass'] > 1e6"],
                 filename='VirialQuantities.h5')

def my_analysis(halo, sphere):
    #total_mass = sphere.quantities['TotalMass']()
    id = halo['id']
    if id < 0: return
    f=open('./halo_metallicity_%04i.txt'%output,'a')
    enriched_mass=sphere["CellMass"][sphere["TotalMetallicity"]>1e-4].sum()
    total_mass =sphere["CellMass"].sum()
    #print enriched_mass,total_mass
    f.write(str(id)+" "+str(enriched_mass)+" "+str(total_mass)+"\n")
    f.close()

#hp.analyze_halo_spheres(my_analysis_virial, halo_list='all',    # 'all'  "filtered",
#                            analysis_output_dir='special_analysis')

#hp.make_profiles(njobs=-1, prefilters=["halo['mass'] > 1e7"],
#                 filename='VirialQuantities.h5')