#Original from Britton Smith, shared on 22. September 2016:
#Edited by Carla Bernhardt, 2019
import sys,yt
yt.enable_parallelism()
from yt.data_objects.particle_filters import particle_filter
from yt.analysis_modules.halo_finding.rockstar.api import RockstarHaloFinder

def get_max_static_level(ds):
        max_level = -1
        for par, val in ds.parameters.items():
                if "StaticRefineRegionLevel" in par:
                        max_level = max(max_level, val)
                if "MustRefineParticlesRefineToLevel" in par:
                        must_refine_level = val
        final_level = max((max_level + 1), must_refine_level)
        #print final_level
        return final_level

def calculate_particle_mass(ds, nested_levels=0):
        "Calculate particle mass from box size and number of particles."

        if "CosmologySimulationOmegaCDMNow" in ds.parameters:
                omega = ds.parameters["CosmologySimulationOmegaCDMNow"]
        else:
                omega = ds.omega_matter
        refine_by = ds.parameters["RefineBy"]

        cell_volume = (ds.domain_right_edge - ds.domain_left_edge).prod() / ds.domain_dimensions.prod()
        rho_bar = omega * ds.cosmology.critical_density(0.0) * (1 + ds.current_redshift)**3
        particle_mass = rho_bar * cell_volume * (1. / refine_by**nested_levels)**3
        return particle_mass.in_units("Msun")

@particle_filter("enzo_2_dark_matter", requires=["creation_time"], filtered_type="io")
def enzo_2_dm(pfilter, data):
        return data["creation_time"] <= 0

# @particle_filter("enzo_2_dark_matter_nested", requires=["particle_mass"],
#                  filtered_type="enzo_2_dark_matter")
@particle_filter("enzo_2_dark_matter_nested", requires=["particle_mass"], filtered_type="all")
def _enzo_2_dm_nested(pfilter, data):
        max_nested = get_max_static_level(data.ds)
        #yt.mylog.info("Getting dark matter particles for nested level %d." % max_nested)
        return data["particle_mass"] <= 1.1 * calculate_particle_mass(data.ds, nested_levels=max_nested)

def setup_ds(ds):
        for pf in ["enzo_2_dark_matter_nested"]:
                rv = ds.add_particle_filter(pf)
                if not rv:
                        raise RuntimeError("Adding filter %s failed!" % pf)

es = yt.simulation(sys.argv[1], "Enzo", find_outputs=True)
#es = yt.simulation(sys.argv[1], "Enzo")

# unigrid
# es.get_time_series()
# rh = RockstarHaloFinder(es, num_readers=2, num_writers=4)

# nested
es.get_time_series(setup_function=setup_ds)
rh = RockstarHaloFinder(es, num_readers=2, num_writers=4, particle_type="enzo_2_dark_matter_nested")

#rh.run()
rh.run(restart=True)