""" Usefull tools """

import os

__all__ = ['dump_gcr']

def dump_gcr(g, outfile, particles="all", mode='w-', **kwargs):
    """Dump full cosmic ray distribution.

    Parameters
    ----------
    g : galpropy.Galprop
        The Galprop instance.
    outfile : str
        Output file name including the path. The directory will be created if
        does not exist.
    particles : list[str] or set[str], or "all", optional
        Specify the Particles to dump. (default "all")
    mode : {'w-', 'x', 'w', 'a'}, optional
        w- or x  Create file, fail if exists  (default)
        w        Create file, truncate if exists
        a        Read/write if exists, create otherwise
    kwargs : dict, optional
        Extra arguments to `h5py.File`.

    """
    try:
        import h5py
        import numpy as np
    except ImportError:
        print("To store GCR please install h5py and numpy, nothing dumped")
        return
    outfile = os.path.abspath(outfile)
    os.makedirs(os.path.dirname(outfile), exist_ok=True)
    if particles == "all":
        particles = {p.name for p in g.gcr}

    with h5py.File(outfile, mode, **kwargs) as f:
        dset = f.create_dataset('r', data=g.gcr[0].r())
        dset.attrs['unit'] = 'kpc'
        dset = f.create_dataset('z', data=g.gcr[0].z())
        dset.attrs['unit'] = 'kpc'
        dset = f.create_dataset('Ekin', data=g.GetEkin())
        dset.attrs['unit'] = 'GeV/n'
        for gcr in g.gcr:
            if gcr.name in particles:
                group = f.create_group(gcr.name)
                crdist = gcr.cr_density
                dist = gcr.A_pos() * 1e3 * np.array(crdist.dist).reshape(
                        (crdist.n_rgrid, crdist.n_zgrid, crdist.n_pgrid))
                group.create_dataset('flux', data=dist)
                group.attrs['Z'] = gcr.Z()
                group.attrs['A'] = gcr.A()
                group.attrs['unit'] = r'cm^-2 s^-1 sr^-1 GeV^-1'