from IPython.parallel import Client
rc = Client()

rc.ids

%autopx

from mpi4py import MPI

MPI.COMM_WORLD.Get_rank()

import yt
yt.enable_parallelism()
import numpy as np
import os
import shutil
import h5py
import gc
import glob

os.chdir('/home/goldbaum/trove')

from yt.utilities.physical_constants import mh, G

def _total_dynamical_time(field, data):
    tdyn = np.sqrt(3*np.pi/(32*G*(data['density'] + data['dark_matter_density']))).in_units('s')
    return tdyn

def _sfr(field, data):
    sfr = np.where(data['density']*data.ds.parameters['Mu']/mh > 50, 
                   0.01*data['density']/data['total_dynamical_time'],
                   0)*yt.units.g/yt.units.cm**3/yt.units.s
    return sfr

from yt.data_objects.particle_filters import add_particle_filter

def star(pfilter, data):
    filter = data["io", "particle_type"] == 2
    return filter

add_particle_filter("stars", function=star, filtered_type='io', requires=["particle_type"])

def young_star(pfilter, data):
    filter = (data.ds.current_time - data["io", "creation_time"]).in_units('Myr') < 10
    return filter

add_particle_filter("young_stars", function=young_star, filtered_type='io', requires=["creation_time"])

def formed_star(pfilter, data):
    filter = data["io", "creation_time"] > 0
    return filter

add_particle_filter("formed_stars", function=formed_star, filtered_type="io", requires=["creation_time"])

base_dir = "./nofeedback_hgf/"
glob_pattern = "DD????/DD????"

output_dir = base_dir+"covering_grids/"

if yt.is_root():
    if not os.path.isdir(output_dir):
        os.mkdir(output_dir)

filenames = sorted(glob.glob(base_dir+glob_pattern))

ts = yt.load(filenames)

# Level of refinement
level = 10

# Dimensions of the base grid, assuming a cubic grid.
n_base = 64

galaxy_center = [0.5, 0.5, 0.5]

gas_fields = (('gas', 'density'), ('gas', 'thermal_energy'), 
              ('gas', 'velocity_x'), ('gas', 'velocity_y'), ('gas', 'velocity_z'),)
gas_units = ('g/cm**3', 'erg/g', 'cm/s', 'cm/s', 'cm/s')

sfr_fields = (('gas', 'sfr_density'),)
sfr_units = ('g/cm**3/s',)

star_fields = ((('stars', 'particle_mass'),) +
               tuple(('stars', 'particle_position_%s' % d) for d in 'xyz') + 
               tuple(('stars', 'particle_velocity_%s' % d) for d in 'xyz')) 
star_units = ('g', 'cm', 'cm', 'cm', 'cm/s', 'cm/s', 'cm/s', )

young_star_fields = ((('young_stars', 'particle_mass'),) +
                     tuple(('young_stars', 'particle_position_%s' % d) for d in 'xyz') + 
                     tuple(('young_stars', 'particle_velocity_%s' % d) for d in 'xyz')) 
young_star_units = ('g', 'cm', 'cm', 'cm', 'cm/s', 'cm/s', 'cm/s', )

formed_star_fields = ((('formed_stars', 'particle_mass'),) +
                      tuple(('formed_stars', 'particle_position_%s' % d) for d in 'xyz') +
                      tuple(('formed_stars', 'particle_velocity_%s' % d) for d in 'xyz'))
formed_star_units = ('g', 'cm', 'cm', 'cm', 'cm/s', 'cm/s', 'cm/s', )

dims = np.array([1024, 1024, 256])

def process_object(dd, filename, field_names, field_units):
    with h5py.File(filename) as data_file:
        for field, unit in zip(field_names, field_units):
            print "Creating covering grid for %s" % str(field)
            data = dd[field].in_units(unit)
            data_file.create_dataset(field[1], data=data, compression='lzf')

class CacheHandler(object):
    def __init__(self, use_species, species_name, output_directory, ds):
        self.use_species = use_species
        self._species_name = species_name
        self._ds_name = str(ds)
        self._output_dir = output_directory + '/{}/'.format(species_name)
        if not os.path.exists(self._output_dir):
            os.makedirs(self._output_dir)
        if os.path.exists(self.filename):
            self.use_species=False

    @property
    def filename(self):
        return self._output_dir + self._ds_name + '_' + self._species_name.replace('_','-') + '_covering_grid.h5'

def process_dataset(ds, use_gas=True, use_sfr=True, use_stars=True, 
                    use_young_stars=True, use_formed_stars=True):
    gas = CacheHandler(use_gas, 'gas', output_dir, ds)
    sfr = CacheHandler(use_sfr, 'sfr', output_dir, ds)
    stars = CacheHandler(use_stars, 'stars', output_dir, ds)
    young_stars = CacheHandler(use_young_stars, 'young_stars', output_dir, ds)
    formed_stars = CacheHandler(use_formed_stars, 'formed_stars', output_dir, ds)

    handlers = (gas, sfr, stars, young_stars, formed_stars)
    do_species = (handler.use_species for handler in handlers)

    if any(do_species):
        pass
    else:
        return

    ds.add_field('total_dynamical_time', _total_dynamical_time, units='s')
    ds.add_field('sfr_density', _sfr, units='g/cm**3/s')
    ds.add_particle_filter('stars')
    ds.add_particle_filter('young_stars')
    ds.add_particle_filter('formed_stars')

    center = ds.arr(galaxy_center, 'code_length')

    left_edge = center - ds.length_unit*dims/(2.*n_base*2**level)
    right_edge = center + ds.length_unit*dims/(2.*n_base*2**level)

    cg = ds.smoothed_covering_grid(level, left_edge=left_edge, dims=dims)
    ad = ds.all_data()

    if gas.use_species is True:
        process_object(cg, gas.filename, gas_fields, gas_units)

    if sfr.use_species is True:
        process_object(cg, sfr.filename, sfr_fields, sfr_units)

    if stars.use_species is True:
        process_object(ad, stars.filename, star_fields, star_units)
        
    if young_stars.use_species is True:
        process_object(ad, young_stars.filename, young_star_fields, young_star_units)

    if formed_stars.use_species is True:
        process_object(ad, formed_stars.filename, formed_star_fields, formed_star_units)

for ds in ts.piter():
    process_dataset(ds)
    print "finished {}".format(str(ds))

%autopx
%matplotlib inline
from matplotlib import pyplot as plt
from matplotlib.colors import LogNorm

fig, ax = plt.subplots()

fig.set_size_inches(10, 8)

with h5py.File('./nofeedback/covering_grids/DD0099_covering_grid.h5') as f:
    im = ax.imshow(np.sum(f['gas/sfr_density'][:], axis=-1), norm=LogNorm())
    cb = fig.colorbar(im)