# coding: utf-8
import pandas as pd
import pspec_new
from baseprocessor import *

mass_func = lambda dset: dset["rho"] * dset["dx"] ** 3  # Mass function
vol_func = lambda dset: dset["dx"] ** 3  # Volume function
getter_T = lambda dset: dset["P"] / dset["rho"]  # Temperature


def getter_B_int(dset):
    B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
    return B_norm


def getter_rho(dset):
    return dset["rho"]


class PostProcessor(HDF5Container):
    """
    This class enable to compute and save derived quantities from the raw output
    """

    # Axes information
    _ax_nb = {"x": 0, "y": 1, "z": 2}  # Number of each axes
    _axes_h = {"x": "y", "y": "x", "z": "x"}  # Associated horizontal axe
    _axes_v = {"x": "z", "y": "z", "z": "y"}  # Associated vertical axe

    G = 1.0  # Gravitational constant

    cells_loaded = False

    def __init__(self, path=None, num=None, path_out=None, pp_params=None, tag=None):
        """
        Creates the basic structures needed for the outputs
        """

        super(PostProcessor, self).__init__(path, path_out, pp_params, tag)

        # Open outfile
        if not self.pp_params.out.tag == "":
            tag_name = self.pp_params.out.tag + "_"
        else:
            tag_name = ""

        self.filename = (
            self.path_out + "/postproc_" + tag_name + format(num, "05") + ".h5"
        )

        self.cells_filename = (
            self.path_out + "/cells_" + tag_name + format(num, "05") + ".h5"
        )

        if not os.path.exists(self.path_out):
            os.makedirs(self.path_out)

        self.open()

        # Ramses Output
        self.path = path
        self.run = os.path.basename(path)
        self.num = num
        self._ro = pymses.RamsesOutput(
            path,
            num,
            order=self.pp_params.pymses.order,
            verbose=self.pp_params.pymses.verbose,
        )
        self._amr = self._ro.amr_source(self.pp_params.pymses.variables)
        self.info = self._ro.info.copy()

        # Density operator
        self._rho_op = ScalarOperator(
            lambda dset: dset["rho"], self._ro.info["unit_density"]
        )

        # Density ray tracer
        if self.pp_params.pymses.fft:
            self._rt = splatting.SplatterProcessor(
                self._amr, self._ro.info, self._rho_op
            )
        else:
            self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op)

        # Set the extend of the image
        self._radius = 0.5 / self.pp_params.pymses.zoom
        self.lbox = self.info["boxlen"]
        center = self.pp_params.pymses.center
        im_extent = [
            (-self._radius + center[0]),
            (self._radius + center[0]),
            (-self._radius + center[1]),
            (self._radius + center[1]),
        ]

        # Get time
        time = self._ro.info["time"]  # time in codeunits

        # Set post processing attributes
        self.save.root._v_attrs.dir = os.path.dirname(path)
        self.save.root._v_attrs.run = os.path.basename(path)
        self.save.root._v_attrs.num = num
        self.save.root._v_attrs.lbox = self.lbox
        self.save.root._v_attrs.unit_length = self.info["unit_length"]
        self.save.root._v_attrs.time = time

        if not "/maps" in self.save:
            self.save.create_group("/", "maps", "2D maps")
        self.save.root.maps._v_attrs.center = center
        self.save.root.maps._v_attrs.radius = self._radius
        self.save.root.maps._v_attrs.im_extent = im_extent

        # Initialize cameras
        self._cam = {}
        for ax_los in self._ax_nb:  # los = line of sight
            ax_h = self._axes_h[ax_los]
            ax_v = self._axes_v[ax_los]

            self._cam[ax_los] = Camera(
                center=self.pp_params.pymses.center,
                line_of_sight_axis=ax_los,
                region_size=[2.0 * self._radius, 2.0 * self._radius],
                distance=self._radius,
                far_cut_depth=self._radius,
                up_vector=ax_v,
                map_max_size=self.pp_params.pymses.map_size,
            )

        self.close()

        self.log_id = "[{}, {}] ".format(self.run, self.num)

        self.def_rules()

    def load_cells(self):
        """
        Load all cells from the source file in the memory.
        Cells will be accessible trough self.cells
        (/!\ Long and memory heavy)
        """
        if not self.cells_loaded:
            if os.path.exists(self.cells_filename):
                cells_hdf5 = tables.open_file(self.cells_filename, mode="r")
                try:
                    node = cells_hdf5.get_node("/cells")
                    self.cells = {}
                    for key in node._v_children:
                        self.cells[key] = cells_hdf5.get_node("/cells/" + key).read()
                finally:
                    cells_hdf5.close()
            else:
                cell_source = CellsToPoints(self._amr)
                cells_pymses = cell_source.flatten()
                self.cells = {}
                for key in cells_pymses.fields:
                    self.cells[key] = cells_pymses[key]
                self.cells["dx"] = cells_pymses.get_sizes()
                self.cells["pos"] = cells_pymses.points

                if self.pp_params.process.save_cells:
                    cells_hdf5 = tables.open_file(self.cells_filename, mode="w")
                    try:
                        for key in self.cells:
                            cells_hdf5.create_array(
                                "/cells", key, self.cells[key], "", createparents=True
                            )
                    finally:
                        cells_hdf5.close()

        self.cells_loaded = True

    def unload_cells(self):
        """
        Free space in the memory by telling the garbage collectors that
        self.cells is not needed
        """
        if self.cells_loaded:
            del self.cells
            self.cells_loaded = False

    def _slice(self, getter, ax_los="z", z=0, unit=cst.none):
        """
        Slice process function.
        Return a slice of the source box.

        Parameters
        ----------
        getter : callable
             A callable that extract the wanted data from a pymses dataset

        ax_los : string
             The axis perpendicular to the slice plane

        z : float
             Coordinate of the slice on the ax_los axis

        unit : cst.Unit
             Unit of the resulting dataset

        Returns
        -------
        A numpy array containing the slice
        """
        op = ScalarOperator(getter, unit)
        datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z)
        return datamap.map.T

    def _ax_avg(
        self, getter, ax_los, unit=cst.none, mass_weighted=True, surf_qty=False
    ):
        """
        Map of the average of a quantity (given by getter) along an axis (ax_los)
        Return 2D array
        """
        if mass_weighted:

            def num(cells):
                value = getter(cells)
                mass = mass_func(cells)
                # Transpose (.T) is for vectorial values
                return (mass * value.T).T

            op = FractionOperator(num, mass_function, unit)
        else:
            op = ScalarOperator(getter, unit)

        if self.pp_params.pymses.fft:
            rt = splatting.SplatterProcessor(self._amr, self._ro.info, op)
        else:
            rt = raytracing.RayTracer(self._amr, self._ro.info, op)

        datamap = rt.process(self._cam[ax_los], surf_qty=surf_qty)
        return datamap.map.T

    def _get_axis(self, axis):

        if isinstance(axis, str):
            axis = self._ax_nb[axis]

        self.load_cells()
        return np.sort(np.unique(self.cells["pos"][:, axis]))

    def _plane_avg_uniform(self, getter, axis, unit=cst.none, mass_weighted=False):
        """
        Profile of the average of a quantity (given by getter) perpendicular to an axis
        WARNING : This version only works on an uniform grid, need of a box version for AMR
        """
        self.load_cells()
        if isinstance(axis, str):
            axis = self._ax_nb[axis]
        axis_data = self.cells["pos"][:, axis]
        value = getter(self.cells)

        df = pd.DataFrame({"axis": axis_data})
        if mass_weighted:
            mass = mass_func(self.cells)
            tot_mass = np.sum(mass)
            df["value"] = value * mass / tot_mass
        else:
            df["value"] = value

        if self.pp_params.process.unload_cells:
            self.unload_cells()

        df.sort_values("axis", inplace=True)

        return df.groupby("axis").mean().values[:, 0]

    def _vol_avg(self, getter, mass_weighted=True):
        self.load_cells()
        value = getter(self.cells)
        if mass_weighted:
            mass = mass_func(self.cells)
            # Transpose (.T) is for vectorial values
            data = np.sum((mass * value.T).T, axis=0) / np.sum(mass)
        else:
            data = np.sum(value, axis=0)

        if self.pp_params.process.unload_cells:
            self.unload_cells()
        return data

    def _vol_pdf(self, getter, bins=100, logbins=False, weight_func=vol_func):
        self.load_cells()
        data = getter(self.cells)
        if logbins:
            data = np.log10(data)
        weights = weight_func(self.cells)
        if self.pp_params.process.unload_cells:
            self.unload_cells()

        values, edges = np.histogram(data, bins, weights=weights)
        centers = 0.5 * (edges[1:] + edges[:-1])
        return (np.stack([values, centers]), {"logbins": logbins})

    def _Brho(self, bins=100, logbins=True):
        """
        Mean of B in rho bins
        """
        self.load_cells()
        B = getter_B_int(self.cells)
        rho = getter_rho(self.cells)

        if logbins:
            rho_bins = np.logspace(
                np.log10(np.min(rho)), np.log10(np.max(rho)), bins, base=10
            )
        else:
            rho_bins = np.linspace(np.min(rho), np.max(rho), bins)

        weights = mass_func(self.cells)

        # For each cell, bin_number contains the number of the bins it belongs to
        bin_number = np.zeros(len(B))

        # Go through the min value of rho of each bin
        for rho_min in rho_bins[:-1]:
            bin_number = bin_number + (rho > rho_min).astype(int)

        # Compute the mean in each bin
        B_mean = np.zeros(len(rho_bins) - 1)
        for i in range(len(B_mean)):
            B_mean[i] = np.mean(B[bin_number == i])

        # Get the center of each bin
        if logbins:
            centers = 10 ** (0.5 * (np.log10(rho_bins[1:]) + np.log10(rho_bins[:-1])))
        else:
            centers = 0.5 * (rho_bins[1:] + rho_bins[:-1])

        if self.pp_params.process.unload_cells:
            self.unload_cells()
        return ({"rho": centers, "B": B_mean}, {"logbins": logbins})

    def cos_vfluct_B(self):

        mean_speed = self.save.get_node("/globals/mwa_speed").read()

        def getter_cos_vfluct_B(dset):
            vel_fluct = dset["vel"] - mean_speed
            B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
            v_norm = np.sqrt(np.sum(vel_fluct ** 2, axis=1))
            # Compute the dot product in each cell
            dot_prod = np.einsum("ij,ij->i", vel_fluct, dset["Br"])
            return np.abs(dot_prod) / (v_norm * B_norm)

        return self._vol_pdf(getter_cos_vfluct_B)

    def _mwa_sigma(self, axes=["x", "y", "z"]):
        mw_speed = self.save.get_node("/globals/mwa_speed").read()

        if axes == ["x", "y", "z"]:

            def getter(dset):
                return np.sum((dset["vel"] - mw_speed) ** 2, axis=1)

        else:

            def getter(dset):
                sigma_squared = 0.0
                for ax in axes:
                    ax_nb = self._ax_nb[ax]
                    sigma_sq_ax = (dset["vel"][:, ax_nb] - mw_speed[ax_nb]) ** 2
                    sigma_squared = sigma_squared + sigma_sq_ax
                return sigma_squared

        return np.sqrt(self._vol_avg(getter, mass_weighted=True))

    def _coldens(self, ax_los):
        datamap = self._rt.process(self._cam[ax_los], surf_qty=True)
        return datamap.map.T

    def _rho(self, ax_los, z=0.0):
        datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=z)
        return (datamap_rho.map).T

    def _vector_h(self, name, unit, ax_los, z=0.0):
        h_op = ScalarOperator(
            lambda dset: dset[name][:, self._ax_nb[self._axes_h[ax_los]]],
            self._ro.info[unit],
        )
        dmap_h = slicing.SliceMap(self._amr, self._cam[ax_los], h_op, z=z).map.T
        return dmap_h

    def _vector_v(self, name, unit, ax_los, z=0.0):
        v_op = ScalarOperator(
            lambda dset: dset[name][:, self._ax_nb[self._axes_v[ax_los]]],
            self._ro.info[unit],
        )
        dmap_v = slicing.SliceMap(self._amr, self._cam[ax_los], v_op, z=z).map.T
        return dmap_v

    def _speed_h(self, ax_los, z=0.0):
        return self._vector_h("vel", "unit_velocity", ax_los, z)

    def _speed_v(self, ax_los, z=0.0):
        return self._vector_v("vel", "unit_velocity", ax_los, z)

    def _B_h(self, ax_los, z=0.0):
        return self._vector_h("Br", "unit_mag", ax_los, z)

    def _B_v(self, ax_los, z=0.0):
        return self._vector_v("Br", "unit_mag", ax_los, z)

    def _B_int(self, ax_los, z=0.0):
        """
        Slice ont the intensity of the magnétic field
        """
        B_op = ScalarOperator(
            lambda dset: np.sqrt(np.sum(dset["Br"] ** 2, axis=1)),
            self._ro.info["unit_mag"],
        )
        dmap_B = (slicing.SliceMap(self._amr, self._cam[ax_los], B_op, z=z)).map.T
        dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read()
        return dmap_B

    def _temperature(self, ax_los, z=0.0):
        P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"])
        dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=z)).map.T
        dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read()
        return dmap_P / dmap_rho

    def _levels(self, ax_los):
        self._amr.set_read_levelmax(self.pp_params.pymses.levelmax)
        level_op = MaxLevelOperator()
        rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op)
        datamap = rt_level.process(self._cam[ax_los], surf_qty=True)
        return datamap.map.T

    def _jeans(self, ax_los):
        dmap_T = self.save.get_node("/maps/T_" + ax_los).read()
        dmap_rho = self.save.get_node("/maps/rho_" + ax_los).read()
        dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho)
        return dmap_jeans

    def _jeans_ratio(self, ax_los):
        dmap_jeans = self.save.get_node("/maps/jeans_" + ax_los).read()
        dmap_levels = self.save.get_node("/maps/levels_" + ax_los).read()
        dmap_jeans_ratio = dmap_jeans * 2 ** (dmap_levels)
        return dmap_jeans_ratio

    def _toomreQ_disk(self, ax_los):
        """
        Compute the Toomre Q parameter in a Keplerian disk
        """

        # Operator to compute the angular speed times rho
        def omega_rho_func(dset):
            pos = dset.get_cell_centers()
            pos = pos - (self.pp_params.disk.pos_star / self.lbox)
            xx = pos[:, :, 0]
            yy = pos[:, :, 1]
            rc = np.sqrt(xx ** 2 + yy ** 2)  # cylindrical radius
            vx = dset["vel"][:, :, 0]
            vy = dset["vel"][:, :, 1]
            omega_rho = 1.0 / rc ** 2
            omega_rho = omega_rho * dset["rho"]
            vyx = vy * xx
            vxy = vx * yy
            omega_rho = omega_rho * (vyx - vxy)
            return omega_rho

        # Operator to compute the angular speed
        omega_op = FractionOperator(
            omega_rho_func, lambda dset: dset["rho"], 1.0 / self._ro.info["unit_time"]
        )

        # Operator to compute the sound speed
        cs_op = FractionOperator(
            lambda dset: dset["P"],
            lambda dset: dset["rho"],
            self._ro.info["unit_velocity"],
        )

        # Ray tracer for the angular speed
        rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op)

        # Ray tracer for the sound speed
        if self.pp_params.pymses.fft:
            rt_cs = splatting.SplatterProcessor(
                self._amr, self._ro.info, cs_op, surf_qty=False
            )
        else:
            rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op)

            dmap_omega = rt_omega.process(self._cam[ax_los])
            dmap_cs = rt_cs.process(self._cam[ax_los])
            dmap_col = self.save.root.maps.coldens_z.read()
            map_Q = (
                (self.lbox * dmap_cs.map.T)
                * dmap_omega.map.T
                / (np.pi * self.G * dmap_col)
            )
        return map_Q

    def _radial_bins(self, _):
        """
        Computes radial bins (for disk)
        """
        pos_star = self.pp_params.disk.pos_star
        im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox

        # radius of the corner of the box plus a margin
        rad_of_box = (
            np.sqrt(
                (im_extent[1] - pos_star[0]) ** 2 + (im_extent[3] - pos_star[1]) ** 2
            )
            + 0.1
        )

        bin_in = self.pp_params.disk.bin_in
        bin_out = self.pp_params.disk.bin_out
        nb_bin = self.pp_params.disk.nb_bin

        # radial bins
        if self.pp_params.disk.binning == "log":
            lrad_in = np.log10(bin_in)
            lrad_ext = np.log10(bin_out)
            rad_bins = np.logspace(lrad_in, lrad_ext, num=nb_bin)
        elif self.pp_params.disk.binning == "lin":
            rad_bins = np.linspace(bin_in, bin_out, num=nb_bin)

        # Add boundaries
        rad_bins = np.concatenate(([0.0], rad_bins, [rad_of_box]))
        return rad_bins

    def _rr(self, _):
        """
        Computes the radius from the center
        """
        im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox
        map_size = self.pp_params.pymses.map_size
        pos_star = self.pp_params.disk.pos_star

        x = np.linspace(im_extent[0], im_extent[1], map_size)
        y = np.linspace(im_extent[2], im_extent[3], map_size)
        xx, yy = np.meshgrid(x, y)
        rr = np.sqrt((xx - pos_star[0]) ** 2 + (yy - pos_star[1]) ** 2)
        return rr

    def _bins_on_map(self, ax_los):
        rad_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
        rr = self.save.get_node("/maps/rr_" + ax_los).read()

        # Find appropriate bin for each coordinate set
        bins = np.zeros(rr.shape, dtype=int)
        for r in rad_bins[1:]:
            bins = bins + (rr >= r).astype(int)
        return bins

    def _rad_avg(self, name, ax_los):
        radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
        bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read()
        dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read()

        # mean of all the cells in the bin
        mean_bin = np.zeros(len(radial_bins) - 1)
        for j in range(len(radial_bins) - 1):
            mean_bin[j] = np.mean(dmap[bins_on_map == j])
        return mean_bin

    def _rad_avg_map(self, name, ax_los):

        radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
        bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read()
        rr = self.save.get_node("/maps/rr_" + ax_los).read()
        mean_bin = self.save.get_node("/radial/rad_avg_" + name + "_" + ax_los).read()

        # Add value for border
        mean_bin = np.concatenate(([mean_bin[0]], mean_bin))

        rr_flat = rr.flatten()
        bins_on_map_flat = bins_on_map.flatten()

        # Compute the map azimuthally averaged
        # use linear interpolation to improve accuracy
        avg_flat = (radial_bins[bins_on_map_flat + 1] - rr_flat) * mean_bin[
            bins_on_map_flat
        ]
        avg_flat = (
            avg_flat
            + (rr_flat - radial_bins[bins_on_map_flat]) * mean_bin[bins_on_map_flat + 1]
        )
        avg_flat = avg_flat / (
            radial_bins[bins_on_map_flat + 1] - radial_bins[bins_on_map_flat]
        )
        avg_map = np.reshape(avg_flat, rr.shape)

        return avg_map

    def _fluct_map(self, name, ax_los):

        dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read()
        avg_map = self.save.get_node("/maps/avg_map_" + name + "_" + ax_los).read()

        return dmap / avg_map

    def _rad_fluct_pdf(self, name, ax_los):
        fluct_map = self.save.get_node("/maps/fluct_" + name + "_" + ax_los).read()
        rr = self.save.get_node("/maps/rr_" + ax_los).read()

        mask_pdf = (rr > self.pp_params.disk.rmin_pdf) & (
            rr < self.pp_params.disk.rmax_pdf
        )

        nb_cells = np.sum(mask_pdf.flatten())
        values, edges = np.histogram(
            np.log10(fluct_map[mask_pdf].flatten()),
            self.pp_params.pdf.nb_bin,
            weights=np.ones(nb_cells) / nb_cells,
        )
        centers = 0.5 * (edges[1:] + edges[:-1])
        return np.stack([values, centers])

    def _fit_pdf(self, name, ax_los):
        pdf = self.save.get_node("/hist/pdf_" + name + "_" + ax_los)
        values, centers = pdf.read()
        mask_fit = (
            (centers > self.pp_params.pdf.xmin_fit)
            & (centers < self.pp_params.pdf.xmax_fit)
            & (values > 0)
        )
        (slope, origin, correlation, _, stderr) = linregress(
            centers[mask_fit], np.log10(values[mask_fit])
        )

        pdf.attrs.slope = slope
        pdf.attrs.origin = origin
        pdf.attrs.correlation = correlation
        pdf.attrs.stderr = stderr
        pdf.attrs.var = np.var
        return True

    def _sinks(self):
        csv_name = (
            self.path
            + "/output_"
            + str(self.num).zfill(5)
            + "/sink_"
            + str(self.num).zfill(5)
            + ".csv"
        )
        sinks = np.loadtxt(csv_name, delimiter=",")
        header = [
            "Id",
            "M",
            "dmf",
            "x",
            "y",
            "z",
            "vx",
            "vy",
            "vz",
            "rot_period",
            "lx",
            "ly",
            "lz",
            "acc_rate",
            "acc_lum",
            "age",
            "int_lum",
            "Teff",
        ]
        if len(sinks) == 0:
            sinks = np.zeros(len(header))

        sinks_dict = {}
        for key, a in zip(header, sinks):
            sinks_dict[key] = a

        return sinks_dict

    def _pspec(self):
        outfile = self.path_out + "/pspec.h5"
        pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile)
        return outfile

    def def_rules(self):

        self.rules = {
            # Base rules
            "coldens": Rule(
                self,
                self._coldens,
                "Column density",
                "/maps",
                unit=self.info["unit_density"] * self.info["unit_length"],
            ),
            "rho": Rule(
                self,
                self._rho,
                "Density slice",
                "/maps",
                unit=self.info["unit_density"],
            ),
            "speed_h": Rule(
                self,
                self._speed_h,
                "Horizontal speed slice wrt the line of sight",
                "/maps",
                unit=self.info["unit_velocity"],
            ),
            "speed_v": Rule(
                self,
                self._speed_v,
                "Vertical speed slice wrt the line of sight",
                "/maps",
                unit=self.info["unit_velocity"],
            ),
            "B_h": Rule(
                self,
                self._speed_h,
                "Horizontal slice of the magnetic field wrt the line of sight",
                "/maps",
                unit=self.info["unit_mag"],
            ),
            "B_v": Rule(
                self,
                self._speed_v,
                "Vertical  slice  of the magnetic field wrt the line of sight",
                "/maps",
                unit=self.info["unit_mag"],
            ),
            "T": Rule(
                self,
                self._temperature,
                "Temperature slice",
                "/maps",
                dependencies=["rho"],
                unit=self.info["unit_temperature"],
            ),
            "levels": Rule(
                self, self._levels, "Max level within line of sight", "/maps"
            ),
            "jeans": Rule(
                self,
                self._jeans,
                "Jeans lenght slice",
                "/maps",
                dependencies=["rho", "T"],
            ),
            "jeans_ratio": Rule(
                self,
                self._jeans_ratio,
                "Jeans' lenght divided by the max resolution",
                "/maps",
                dependencies=["jeans", "levels"],
            ),
            "Q": Rule(
                self,
                self._toomreQ_disk,
                "Toomre Q parameter for a Keplerian disk",
                "/maps",
                dependencies=["coldens"],
                is_valid=lambda _: self.pp_params.disk.on,
            ),
            "sinks": Rule(
                self,
                self._sinks,
                group="/datasets",
                unit={
                    "Id": cst.none,
                    "M": cst.Msun,
                    "dmf": cst.Msun,
                    "x": "",
                    "y": "",
                    "z": "",
                    "vx": "",
                    "vy": "",
                    "vz": "",
                    "rot_period": "[y]",
                    "lx": "|l|",
                    "ly": "|l|",
                    "lz": "|l|",
                    "acc_rate": "[Msol/y]",
                    "acc_lum": "[Lsol]",
                    "age": cst.year,
                    "int_lum": "[Lsol]",
                    "Teff": cst.K,
                },
            ),
            "pspec": Rule(self, self._pspec, "Power spectrum", "/hdf5"),
            # Helpers
            "radial_bins": Rule(self, self._radial_bins, "Radial bins", "/radial"),
            "rr": Rule(self, self._rr, "Coordinate map", "/maps"),
            "bins_on_map": Rule(
                self,
                self._bins_on_map,
                "Convert map coordinates to bins",
                "/maps",
                dependencies=["radial_bins", "rr"],
            ),
            "B_int": Rule(
                self,
                self._B_int,
                "Magnetic intensity slice",
                "/maps",
                dependencies=["rho"],
                unit=self.info["unit_mag"],
            ),
            # PDF
            "rho_pdf": Rule(
                self,
                partial(self._vol_pdf, partial(simple_getter, "rho"), logbins=True),
                "Global rho-PDF",
                "/hist",
                unit=self.info["unit_density"],
            ),
            "T_pdf": Rule(
                self,
                partial(self._vol_pdf, getter_T, logbins=True),
                "Global T-PDF",
                "/hist",
                unit=self.info["unit_temperature"],
            ),
            "P_pdf": Rule(
                self,
                partial(self._vol_pdf, getter_T, logbins=True),
                "Global P-PDF",
                "/hist",
                unit=self.info["unit_pressure"],
            ),
            "cos_pdf": Rule(
                self,
                partial(self.cos_vfluct_B),
                "Global cos fluctuation-PDF",
                "/hist",
                dependencies=["mwa_speed"],
                unit=cst.none,
            ),
            "Brho": Rule(
                self,
                self._Brho,
                "Average of B as a function of rho",
                "/datasets",
                unit={"rho": self.info["unit_density"], "B": self.info["unit_mag"]},
            ),
            # Profiles
            "axis": Rule(
                self,
                partial(self._get_axis),
                "Axis",
                "/profile",
                unit=self.info["unit_length"],
            ),
            "rho_prof": Rule(
                self,
                partial(self._plane_avg_uniform, partial(simple_getter, "rho")),
                "Rho profile",
                "/profile",
                unit=self.info["unit_density"],
                dependencies=["axis"],
            ),
            # globals
            "time_num": Rule(
                self,
                lambda: self.info["time"],
                "Time",
                "/globals",
                unit=self.info["unit_time"],
            ),
            "mwa_speed": Rule(
                self,
                partial(self._vol_avg, partial(simple_getter, "vel")),
                "Mass weighted speed average",
                "/globals",
                unit=self.info["unit_velocity"],
            ),
            "mwa_sigma": Rule(
                self,
                self._mwa_sigma,
                "Mass weighted speed average",
                "/globals",
                dependencies={"mwa_speed": None},
                unit=self.info["unit_velocity"],
            ),
        }

        # Average and other
        averageables = ["coldens", "rho", "T", "Q"]
        for name in averageables:
            self.rules["rad_avg_" + name] = Rule(
                self,
                partial(self._rad_avg, name),
                "Azimuthal average of {}".format(name),
                "/radial",
                dependencies=["radial_bins", "bins_on_map", name],
            )

            self.rules["avg_map_" + name] = Rule(
                self,
                partial(self._rad_avg_map, name),
                "Interpolated map of azimuthal average of {}".format(name),
                "/maps",
                dependencies=["radial_bins", "bins_on_map", "rr", "rad_avg_" + name],
            )
            self.rules["fluct_" + name] = Rule(
                self,
                partial(self._fluct_map, name),
                "Fluctuation wrt to average of {}".format(name),
                "/maps",
                dependencies=[name, "avg_map_" + name],
            )
            self.rules["pdf_" + name] = Rule(
                self,
                partial(self._rad_fluct_pdf, name),
                "Probability density function of {} fluctuations".format(name),
                "/hist",
                dependencies=["rr", "fluct_" + name],
            )

            self.rules["fit_pdf_" + name] = Rule(
                self,
                partial(self._fit_pdf, name),
                "Fit the PDF of {} fluctuations".format(name),
                "/hist",
                dependencies=["pdf_" + name],
            )

        self._gen_rule_transform("fluct_coldens", np.max, "max", group="/globals")

        super(PostProcessor, self).def_rules()


def get_time(path, num):
    """
    Return the time of the output (code units)

    Parameters
    ----------
    num       output number
    path_out  path of the pipeline output

    Returns
    -------
    time      the time of the output (code units)
    """
    try:
        f = open(
            path
            + "/output_"
            + str(num).zfill(5)
            + "/info_"
            + str(num).zfill(5)
            + ".txt"
        )
        for line in f:
            ls = line.split()
            if len(ls) > 1 and ls[0] == "time":
                time = float(ls[2])
                break
        f.close()
        return time
    except IOError as e:
        print(e)
        return np.nan