#
# Copyright 2020 Antoine Sanner
#           2020 Lars Pastewka
#
# ### MIT license
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

"""
Defines the interface for Adhesion systems
"""

import numpy as np

import Adhesion
import ContactMechanics
import SurfaceTopography
from ContactMechanics.Tools import compare_containers
from ContactMechanics.Systems import IncompatibleResolutionError, SystemBase

class SmoothContactSystem(SystemBase):
    """
    For smooth contact mechanics (i.e. the ones for which optimization is only
    kinda-hell
    """

    def __init__(self, substrate, interaction, surface):
        """ Represents a contact problem
        Parameters
        ----------
        substrate: An instance of HalfSpace.
            Defines the solid mechanics in the substrate
        interaction: Adhesion.Interactions.SoftWall
            Defines the contact formulation.
            If this computes interaction energies, forces etc,
            these are supposed to be expressed per unit area in whatever units
             you use. The conversion is performed by the system
        surface: SurfaceTopography.Topography
            Defines the profile.
        """
        if surface.has_undefined_data:
            raise ValueError("The topography you provided contains undefined "
                             "data")
        super().__init__(substrate=substrate, surface=surface)
        self.interaction = interaction
        if not compare_containers(surface.nb_grid_pts, substrate.nb_grid_pts):
            raise IncompatibleResolutionError(
                ("the substrate ({}) and the surface ({}) have incompatible "
                 "nb_grid_ptss.").format(
                    substrate.nb_grid_pts, surface.nb_grid_pts))  # nopep8
        self.dim = len(self.substrate.nb_grid_pts)
        self.energy = None
        self.force = None
        self.force_h = None
        self.interaction_energy = None
        self.interaction_force = None
        self.heights_k = None
        self.engine = self.substrate.fftengine

        if hasattr(substrate.fftengine, "register_halfcomplex_field") and self.engine.communicator.size == 1:
            # avoids the initialization to fail if we use an fftengine without these hcffts implemnented
            # preconditionning is not parallelized yet, this we have no parallelized wrapper for hcfft in muFFT yet
            self.real_buffer = self.engine.register_halfcomplex_field("hc-real-space", 1)
            self.fourier_buffer = self.engine.register_halfcomplex_field("hc-fourier-space", 1)

            self.stiffness_k = self._compute_stiffness_k()

    @property
    def nb_grid_pts(self):
        return self.surface.nb_grid_pts

    @staticmethod
    def handles(substrate_type, interaction_type, surface_type, comm):
        is_ok = True
        # any periodic type of substrate formulation should do
        is_ok &= issubclass(substrate_type,
                            ContactMechanics.Substrate)

        # only soft interactions allowed
        is_ok &= issubclass(interaction_type,
                            Adhesion.SoftWall)

        # any surface should do
        is_ok &= issubclass(surface_type,
                            SurfaceTopography.UniformTopographyInterface)
        return is_ok

    def compute_repulsive_force(self):
        "computes and returns the sum of all repulsive forces"
        return self.reduction.sum(np.where(
            self.interaction_force > 0, self.interaction_force, 0
        ))

    def compute_attractive_force(self):
        "computes and returns the sum of all attractive forces"
        return self.reduction.sum(np.where(
            self.interaction_force < 0, self.interaction_force, 0
        ))

    def _compute_stiffness_k(self):
        """
        computes and returns the wavevectors q that exist for the surfaces
        physical_sizes and nb_grid_pts as one vector of components per
        dimension
        """

        vectors = []
        q = []
        nb_dims = len(self.substrate.nb_grid_pts)
        # if nb_dims == 1:
        #     nb_grid_pts = [self.substrate.nb_grid_pts]
        #     physical_sizes = [self.substrate.physical_sizes]
        for dim in range(nb_dims):
            vectors.append(2 * np.pi * np.fft.fftfreq(
                self.substrate.nb_grid_pts[dim],
                self.substrate.physical_sizes[dim] /
                self.substrate.nb_grid_pts[dim]))
        if nb_dims == 1:
            q = vectors[0]
            q[0] = q[1]
        elif nb_dims == 2:
            qx = vectors[0]
            qy = vectors[1]
            q = np.sqrt(
                (qx * qx).reshape((-1, 1)) + (qy * qy).reshape((1, -1)))
            q[0, 0] = (q[0, 1] + q[1, 0]) / 2

        return 0.5 * self.substrate.contact_modulus * abs(q)

    def compute_normal_force(self):
        "computes and returns the sum of all forces"
        return self.reduction.sum(self.interaction_force)

    def compute_repulsive_contact_area(self):
        "computes and returns the area where contact pressure is repulsive"
        return self.compute_nb_repulsive_pts() * self.area_per_pt

    def compute_attractive_contact_area(self):
        "computes and returns the are where contact pressure is attractive"
        return self.compute_nb_attractive_pts() * self.area_per_pt

    def compute_nb_contact_pts(self):
        """
        compute and return the number of contact points. Note that this is of
        no physical interest, as it is a purely numerical artefact
        """
        return self.reduction.sum(
            np.where(self.interaction_force != 0., 1., 0.))

    def compute_nb_repulsive_pts(self):
        """
        compute and return the number of contact points under repulsive
        pressure. Note that this is of no physical interest, as it is a
        purely numerical artefact
        """
        return self.reduction.sum(
            np.where(self.interaction_force > 0., 1., 0.))

    def compute_nb_attractive_pts(self):
        """
        compute and return the number of contact points under attractive
        pressure. Note that this is of no physical interest, as it is a
        purely numerical artefact
        """
        return self.reduction.sum(
            np.where(self.interaction_force < 0., 1., 0.))

    def compute_repulsive_coordinates(self):
        """
        returns an array of all coordinates, where contact pressure is
        repulsive. Useful for evaluating the number of contact islands etc.
        """
        return np.argwhere(self.interaction_force > 0.)

    def compute_attractive_coordinates(self):
        """
        returns an array of all coordinates, where contact pressure is
        attractive. Useful for evaluating the number of contact islands etc.
        """
        return np.argwhere(self.interaction_force < 0.)

    def compute_mean_gap(self):
        """
        mean of the gap in the the physical domain (means excluding padding
        region for the FreeFFTElasticHalfspace)
        """
        return self.reduction.sum(self.gap) / np.prod(self.nb_grid_pts)

    def logger_input(self):
        """

        Returns
        -------
        headers: list of strings
        values: list
        """
        tot_nb_grid_pts = np.prod(self.nb_grid_pts)
        rel_rep_area = self.compute_nb_repulsive_pts() / tot_nb_grid_pts
        rel_att_area = self.compute_nb_attractive_pts() / tot_nb_grid_pts
        # TODO: eventually put a flag to turn
        #  reductions off since this is an additional communication.
        return (['energy',
                 'max. abs. grad.',
                 'mean gap',
                 'frac. rep. area',
                 'frac. att. area',
                 'frac. int. area', 'substrate force', 'interaction force'],
                [self.energy,
                 self.reduction.max(np.abs(self.force)),
                 self.compute_mean_gap(),
                 rel_rep_area,
                 rel_att_area,
                 rel_rep_area + rel_att_area,
                 -self.reduction.sum(self.substrate.force),
                 self.reduction.sum(self.interaction_force)])

    def evaluate(self, disp, offset, pot=True, forces=False, logger=None):
        """
        Compute the energies and forces in the system for a given displacement
        field

        Parameters:
        -----------
        disp: ndarray
            displacement field, in the shape of
            system.substrate.nb_subdomain_grid_pts
        offset: float
            determines indentation depth,
            constant value added to the heights (system.topography)
        pot: bool, optional
            Wether to evaluate the energy, default True
        forces: bool, optional
            Wether to evaluate the forces, default False
        logger: ContactMechanics.Tools.Logger
            informations of current state of the system will be passed to
            logger at every evaluation
        """
        # attention: the substrate may have a higher nb_grid_pts than the gap
        # and the interaction (e.g. FreeElasticHalfSpace)
        self.gap = self.compute_gap(disp, offset)
        interaction_energies, self.interaction_force, _ = \
            self.interaction.evaluate(self.gap,
                                      potential=pot,
                                      gradient=forces,
                                      curvature=False)

        self.interaction_energy = \
            self.reduction.sum(interaction_energies) * self.area_per_pt

        self.substrate.compute(disp, pot, forces)
        self.energy = (self.interaction_energy +
                       self.substrate.energy
                       if pot else None)
        if forces:
            self.interaction_force *= -self.area_per_pt
            #                       ^ gradient to force per pixel
            self.force = self.substrate.force.copy()
            self.force[self.comp_slice] += \
                self.interaction_force
        else:
            self.force = None

        if logger is not None:
            logger.st(*self.logger_input())
        return (self.energy, self.force)

    def objective(self, offset, disp0=None, gradient=False, disp_scale=1.,
                  logger=None):
        r"""
        This helper method exposes a scipy.optimize-friendly interface to the
        evaluate() method. Use this for optimization purposes, it makes sure
        that the shape of disp is maintained and lets you set the offset and
        'forces' flag without using scipy's cumbersome argument passing
        interface. Returns a function of only disp

        Parameters:
        -----------
        disp0: ndarray
            unused variable, present only for interface compatibility
            with inheriting classes
        offset: float
            determines indentation depth,
            constant value added to the heights (system.topography)
        gradient: bool, optional
            Wether to evaluate the gradient, default False
        disp_scale : float, optional
            (default 1.) allows to specify a scaling of the
            dislacement before evaluation. This can be necessary when
            using dumb minimizers with hardcoded convergence criteria
            such as scipy's L-BFGS-B.
        logger: ContactMechanics.Tools.Logger
            informations of current state of the system will be passed to
            logger at every evaluation

        Returns:
            function(disp)
                Parameters:
                disp: an ndarray of shape
                      `system.substrate.nb_subdomain_grid_pts`
                      displacements
                Returns:
                    energy or energy, gradient
        """
        res = self.substrate.nb_subdomain_grid_pts
        if gradient:
            def fun(disp):
                # pylint: disable=missing-docstring
                try:
                    self.evaluate(
                        disp_scale * disp.reshape(res), offset, forces=True,
                        logger=logger)
                except ValueError as err:
                    raise ValueError(
                        "{}: disp.shape: {}, res: {}".format(
                            err, disp.shape, res))
                return (self.energy, -self.force.reshape(-1) * disp_scale)
        else:
            def fun(disp):
                # pylint: disable=missing-docstring
                return self.evaluate(
                    disp_scale * disp.reshape(res), offset, forces=False,
                    logger=logger)[0]

        return fun

    def primal_evaluate(self, disp, gap, pot=True, forces=False, logger=None):
        """
        Compute the energies and forces in the system for a given
        displacement and gap..

        Parameters:
        -----------
        disp: ndarray
            displacement field, in the shape of
            system.substrate.nb_subdomain_grid_pts
        gap: ndarray
            gap , in the shape of
            system.substrate.nb_subdomain_grid_pts
        pot: bool, optional
            Whether to evaluate the energy, default True
        forces: bool, optional
            Whether to evaluate the forces, default False
        logger: ContactMechanics.Tools.Logger
            information of current state of the system will be
            passed to
            logger at every evaluation
        """
        # attention: the substrate may have a higher nb_grid_pts than the gap
        # and the interaction (e.g. FreeElasticHalfSpace)

        self.gap = gap
        interaction_energies, self.interaction_force, _ = \
            self.interaction.evaluate(self.gap,
                                      potential=pot,
                                      gradient=forces,
                                      curvature=False)

        self.interaction_energy = \
            self.reduction.sum(interaction_energies) * self.area_per_pt

        self.substrate.compute(disp, pot, forces)
        self.energy = (self.interaction_energy +
                       self.substrate.energy
                       if pot else None)
        if forces:
            self.interaction_force *= -self.area_per_pt
            #                       ^ gradient to force per pixel
            self.force = self.substrate.force.copy()

            self.force[self.comp_slice] += self.interaction_force
        else:
            self.force = None

        if logger is not None:
            logger.st(*self.logger_input())

        return (self.energy, self.force)

    def primal_objective(self, offset, disp0=None, gradient=False,
                         logger=None):
        r"""To solve the primal objective using gap as the variable.
        Can be fed directly to standard solvers ex: scipy solvers etc
        and returns the elastic energy and it's gradient (negative of
        the forces) as a function of the gap.

        Parameters
        __________

        gap : float
              gap between the contact surfaces.
        offset : float
                constant value to add to the surface heights
        pot : (default False)

        gradient : (default True)

        Returns
        _______
        energy : float
                value of energy(scalar value).
        force : float,array
                value of force(array).

        Notes
        _____

        Objective:

        .. math ::

            \min_u f = \frac{1}{2} u_i K_{ij} u_j + \phi (u_{ij})\\
            \\
            \nabla f = K_{ij} u_j + \phi^{\prime} \text{    which is the force} \\

        """

        res = self.substrate.nb_subdomain_grid_pts
        if gradient:
            def fun(gap):
                disp = gap.reshape(res) + self.surface.heights() + offset
                try:
                    self.primal_evaluate(
                        disp, gap.reshape(res), forces=True, logger=logger)
                except ValueError as err:
                    raise ValueError(
                        "{}: gap.shape: {}, res: {}".format(
                            err, gap.shape, res))
                return (self.energy, -self.force.reshape(-1))
        else:
            def fun(gap):
                disp = gap.reshape(res) + self.surface.heights() + offset
                return self.primal_evaluate(
                    disp, gap.reshape(res), forces=False, logger=logger)[0]

        return fun

    def primal_hessian_product(self, gap, des_dir):
        """Returns the hessian product of the primal_objective function.
        """
        _, _, adh_curv = self.interaction.evaluate(gap, curvature=True)

        hessp_val = - self.substrate.evaluate_force(
            des_dir.reshape(self.substrate.nb_subdomain_grid_pts)
        ).reshape(np.shape(des_dir)) \
            + adh_curv * des_dir * self.substrate.area_per_pt

        return hessp_val.reshape(des_dir.shape)

    def hessian_product_function(self, offset):
        def hessp(disp, des_dir):
            gap = disp.reshape(self.substrate.nb_subdomain_grid_pts
                               )[self.comp_slice] \
                  - (self.surface.heights() + offset)
            _, _, adh_curv = self.interaction.evaluate(gap, curvature=True)
            hessp_val = - self.substrate.evaluate_force(
                des_dir.reshape(self.substrate.nb_subdomain_grid_pts)
            )

            hessp_val[self.comp_slice] += adh_curv \
                * des_dir.reshape(
                self.substrate.nb_subdomain_grid_pts)[self.comp_slice] * \
                self.substrate.area_per_pt
            return hessp_val.reshape(des_dir.shape)

        return hessp

    def _fourier_coefficients(self):
        """
        Returns the coefficients for elasticity matrix when working in fourier
        space for both 1D and 2D system.
        """

        nx = self.substrate.nb_grid_pts[0]
        nb_dims = len(self.substrate.nb_grid_pts)

        if nb_dims == 2:
            ny = self.substrate.nb_grid_pts[1]
            coeffs = np.zeros(self.substrate.nb_grid_pts)
            if np.logical_and((nx % 2 == 0), (ny % 2 == 0)):
                coeffs[0, 0] = 1 / (nx * ny)
                coeffs[0, 1:ny // 2] = 2 / (nx * ny)
                coeffs[0, ny // 2 + 1:] = 2 / (nx * ny)
                coeffs[1:nx // 2, 0] = 2 / (nx * ny)
                coeffs[nx // 2 + 1:, 0] = 2 / (nx * ny)
                coeffs[:nx // 2, ny // 2] = 2 / (nx * ny)
                coeffs[nx // 2 + 1:, ny // 2] = 2 / (nx * ny)
                coeffs[nx // 2, :ny // 2] = 2 / (nx * ny)
                coeffs[nx // 2, ny // 2 + 1:] = 2 / (nx * ny)
                coeffs[1:nx // 2, 1:ny // 2] = 4 / (nx * ny)
                coeffs[nx // 2 + 1:, 1:ny // 2] = 4 / (nx * ny)
                coeffs[1:nx // 2, ny // 2 + 1:] = 4 / (nx * ny)
                coeffs[nx // 2 + 1:, ny // 2 + 1:] = 4 / (nx * ny)
                coeffs[nx // 2, ny // 2] = 1 / (nx * ny)
                coeffs[nx // 2, 0] = 1 / (nx * ny)
                coeffs[0, ny // 2] = 1 / (nx * ny)
            else:
                coeffs[0, 0] = 1 / (nx * ny)
                coeffs[0, 1:] = 2 / (nx * ny)
                coeffs[1:, 0] = 2 / (nx * ny)
                coeffs[1:, 1:] = 4 / (nx * ny)
        elif nb_dims == 1:
            coeffs = np.zeros(self.substrate.nb_grid_pts)
            if (nx % 2 == 0):
                coeffs[0] = 1 / nx
                coeffs[1:nx // 2] = 2 / nx
                coeffs[nx // 2 + 1:] = 2 / nx
                coeffs[nx // 2] = 1 / nx
            else:
                coeffs[0] = 1 / nx
                coeffs[1:] = 2 / nx

        return coeffs

    def evaluate_k(self, disp_k, gap, offset, mw=False, pot=True, forces=False,
                   logger=None):

        """
        Compute the energies and forces in the system for a given displacement
        field in fourier space.

        Parameters
        -----------

        disp_k: ndarray
            displacement field in fourier space.
        gap:  ndarray
            displacement field in real space, in the shape of
            system.substrate.nb_subdomain_grid_pts
        offset: float
            determines indentation depth,
            constant value added to the heights (system.topography)
        mw: bool, optional
            when mass weighting is required then set this to TRUE.
        pot: bool, optional
            Wether to evaluate the energy, default True
        forces: bool, optional
            Wether to evaluate the forces, default False
        logger: ContactMechanics.Tools.Logger
            informations of current state of the system will be passed to
            logger at every evaluation.
        """

        # self.gap = self.compute_gap(disp, offset)
        self.gap = gap
        interaction_energies, self.interaction_force, _ = \
            self.interaction.evaluate(self.gap,
                                      potential=pot,
                                      gradient=forces,
                                      curvature=False)
        self.interaction_energy = \
            self.reduction.sum(interaction_energies) * self.area_per_pt

        self.grad_k = np.zeros(self.substrate.nb_grid_pts)

        coeff = self._fourier_coefficients()

        if mw:
            self.grad_k = disp_k * coeff
        else:
            self.grad_k = disp_k * coeff * self.stiffness_k

        self.grad_k *= self.area_per_pt

        # ENERGY FROM SUBSTRATE
        self.energy = 0.5 * (np.sum(self.grad_k * disp_k))

        self.substrate.energy = self.energy

        self.force_h = -self.grad_k

        # TOTAL ENERGY
        self.energy += self.interaction_energy

        if forces:
            self.interaction_force *= -self.area_per_pt
            #                     ^ gradient to force per pixel

            self.real_buffer.array()[...] = self.interaction_force
            self.engine.hcfft(self.real_buffer, self.fourier_buffer)
            interaction_force_float_k = self.fourier_buffer.array()[...].copy()

            adh_coeffs = self._fourier_coefficients()
            interaction_force_float_k *= adh_coeffs

            if mw:
                k = np.sqrt(self.stiffness_k.copy() * self.area_per_pt)
                interaction_force_float_k = interaction_force_float_k * (1 / k)

            self.force_h += interaction_force_float_k
        else:
            self.force_h = None

        if logger is not None:
            disp_real = self.gap + self.surface.heights().copy() + offset
            force_real = self.substrate.evaluate_force(disp_real)
            force_real = force_real + self.interaction_force
            logger.st(*(['energy',
                         'max. abs. grad.',
                         'max. abs. grad. real'],
                        [self.energy,
                         self.reduction.max(np.abs(self.force_h)),
                         self.reduction.max(np.abs(force_real))
                         ]))

        return (self.energy, self.force_h)

    # def hessian_product_k(self, dispk, des_dir_k):
    #     """Returns the hessian product of the fourier space
    #     objective_k function.
    #     """
    #     self.substrate.fourier_buffer.array()[...] = dispk.copy()
    #     self.substrate.fftengine.ifft(self.substrate.fourier_buffer,
    #                                   self.substrate.real_buffer)
    #     disp = self.substrate.real_buffer.array()[...].copy() \
    #         * self.substrate.fftengine.normalisation
    #
    #     gap = self.compute_gap(disp)
    #     _, _, adh_curv = self.interaction.evaluate(gap, curvature=True)
    #
    #     self.substrate.real_buffer.array()[...] = adh_curv.reshape(
    #         self.substrate.nb_grid_pts).copy()
    #     self.substrate.fftengine.fft(self.substrate.real_buffer,
    #                                  self.substrate.fourier_buffer)
    #     adh_curv_k = self.substrate.fourier_buffer.array()[...].copy()
    #
    #     hessp_val_k = -self.substrate.evaluate_k_force_k(des_dir_k) + \
    #         adh_curv_k * des_dir_k * self.substrate.area_per_pt
    #
    #     return hessp_val_k

    # def hessian_product_preconditioned(self, offset):

    #     def hessp_precond(disp_h, des_dir):
    #
    #         self.real_buffer.array()[...] = offset
    #         self.engine.hcfft(self.real_buffer, self.fourier_buffer)
    #         offset_k = self.fourier_buffer.array()[...].copy()
    #
    #         self.real_buffer.array()[...] = self.surface.heights().copy()
    #         self.engine.hcfft(self.real_buffer, self.fourier_buffer)
    #         self.heights_k_float = self.fourier_buffer.array()[...].copy()
    #
    #         disp_float_k = disp_h.copy()
    #         disp_float_k = disp_float_k.reshape(self.substrate.nb_grid_pts)
    #         gap_float_k = (disp_float_k / np.sqrt(self.stiffness_k *
    #                                               self.area_per_pt)) - \
    #                       self.heights_k_float - offset_k
    #
    #         self.fourier_buffer.array()[...] = gap_float_k.copy()
    #         self.engine.ihcfft(self.fourier_buffer, self.real_buffer)
    #         gap = self.real_buffer.array()[...].copy() * \
    #               self.engine.normalisation
    #
    #         _, _, adh_curv = \
    #             self.interaction.evaluate(gap,
    #                                       curvature=True)
    #
    #         coeff = self._fourier_coefficients()
    #
    #         el_hess_k = coeff @ coeff
    #         el_hess_k *= self.area_per_pt
    #
    #         orig_shape = np.shape(des_dir)
    #
    #         el_hess_k *= des_dir #.reshape(el_hess_k)
    #
    #         adh_curv *= self.area_per_pt
    #
    #         self.real_buffer.array()[...] = adh_curv
    #         self.engine.hcfft(self.real_buffer, self.fourier_buffer)
    #         adh_curv_float_k = self.fourier_buffer.array()[...].copy()
    #
    #         adh_curv_float_k *= coeff @ coeff
    #
    #         k = np.sqrt(self.stiffness_k.copy() * self.area_per_pt)
    #
    #         adh_hess_k = adh_curv_float_k * des_dir * (1 / k) * (1 / k)
    #
    #         hess_k = el_hess_k.reshape(orig_shape) + \
    #                  adh_hess_k.reshape(orig_shape)
    #
    #         return hess_k
    #
    #     return hessp_precond

    def preconditioned_objective(self, offset, gradient=False, logger=None):
        r"""
        This helper method interface to the evaluate_k() method with
        preconditioning active. That is, it tries to solve a simpler problem
        formulated using,

        original problem:

        .. math ::

            \frac{1}{2(n_x n_y)} \tilde{u}\tilde{K} \bar{\tilde{u}} +
            \phi(F^{-1}(\tilde{u} - \tilde{h}))

        preconditioned problem:

        .. math ::
             \tilde{v} = \tilde{k}^{\frac{1}{2}} \tilde{u} \\

             \frac{1}{2(n_x n_y)} \tilde{v}\bar{\tilde{v}} +
            \phi(F^{-1}(\frac{\tilde{v}}{\tilde{k}^{\frac{1}{2}}} - \tilde{h}))

        we solve for variable :math:`\tilde{v}`.

        Parameters:
        -----------

        offset: float
            determines indentation depth,
            constant value added to the heights (system.topography)
        gradient: bool, optional
            Whether to evaluate the gradient, default False
        logger: ContactMechanics.Tools.Logger
            information of current state of the system will be passed to
            logger at every evaluation

        Returns
        _______

            function(disp_k)

                Parameters
                __________

                disp_k: an ndarray in fourier halfcomplex space

                Returns
                _______

                energy: scalar
                        energy of the system

                force_h: an halfcomplex array of shape(disp_k)
                        force of the system
        """

        # TODO: fourier transforming the offset is useless
        self.real_buffer.array()[...] = offset
        self.engine.hcfft(self.real_buffer, self.fourier_buffer)
        offset_k = self.fourier_buffer.array()[...].copy()

        self.real_buffer.array()[...] = self.surface.heights().copy()
        self.engine.hcfft(self.real_buffer, self.fourier_buffer)
        self.heights_k_float = self.fourier_buffer.array()[...].copy()

        if gradient:
            def fun(disp_):
                disp_float_k = disp_.copy()
                orig_shape = np.shape(disp_float_k)
                disp_float_k = disp_float_k.reshape(self.substrate.nb_grid_pts)
                gap_float_k = (disp_float_k / np.sqrt(self.stiffness_k *
                                                      self.area_per_pt)) - \
                    self.heights_k_float - offset_k

                self.fourier_buffer.array()[...] = gap_float_k.copy()
                self.engine.ihcfft(self.fourier_buffer, self.real_buffer)
                gap = self.real_buffer.array()[...].copy() * self.engine.normalisation

                self.energy, self.force_h = self.evaluate_k(disp_float_k,
                                                            gap, offset,
                                                            mw=True,
                                                            forces=True,
                                                            logger=logger
                                                            )

                return (self.energy, -self.force_h.reshape(orig_shape))
        else:
            raise NotImplementedError

        return fun

    def objective_k_float(self, offset, gradient=False, logger=None):
        r"""

        Returns callable objective as needed by scipy minimizers.

        The optimisation varialbe is the halfcomplex fourier transform of the gap.

        This helper method interface to the evaluate_k() method without
        preconditioning active.

        .. math ::

            \frac{1}{2(n_x n_y)} \tilde{u}\tilde{K} \bar{\tilde{u}} +
            \phi(F^{-1}(\tilde{u} - \tilde{h}))  \\

        preconditioned problem:

        .. math ::

             \tilde{v} = \tilde{k}^{\frac{1}{2}} \tilde{u} \\

             \frac{1}{2(n_x n_y)} \tilde{v} \bar{\tilde{v}} +
            \phi(F^{-1}(\frac{\tilde{v}}{\tilde{k}^{\frac{1}{2}}} - \tilde{h}))
        \\
        we solve for variable :math:`\tilde{v}`.

        Parameters:
        -----------

        offset: float
            determines indentation depth,
            constant value added to the heights (system.topography)
        gradient: bool, optional
            Whether to evaluate the gradient, default False
        logger: ContactMechanics.Tools.Logger
            information of current state of the system will be passed to
            logger at every evaluation

        Returns
        _______

            function(disp_k)

                Parameters
                __________

                disp_k: an ndarray in fourier space

                Returns
                _______

                    energy, gradient_k_float
        """

        self.real_buffer.array()[...] = offset
        self.engine.hcfft(self.real_buffer, self.fourier_buffer)
        offset_k = self.fourier_buffer.array()[...].copy()

        self.real_buffer.array()[...] = self.surface.heights().copy()
        self.engine.hcfft(self.real_buffer, self.fourier_buffer)
        self.heights_k_float = self.fourier_buffer.array()[...].copy()

        if gradient:
            def fun(disp_k):
                disp_float_k = disp_k.copy()
                orig_shape = np.shape(disp_float_k)
                disp_float_k = disp_float_k.reshape(self.substrate.nb_grid_pts)
                gap_float_k = disp_float_k - self.heights_k_float - offset_k
                self.fourier_buffer.array()[...] = gap_float_k.copy()
                self.engine.ihcfft(self.fourier_buffer, self.real_buffer)
                gap = self.real_buffer.array()[...].copy() \
                    * self.engine.normalisation

                self.energy, self.force_h = self.evaluate_k(disp_float_k,
                                                            gap, offset,
                                                            forces=True,
                                                            logger=logger
                                                            )
                return (self.energy, -self.force_h.reshape(orig_shape))
        else:
            def fun(disp_k):
                # pylint: disable=missing-docstring
                disp_float_k = disp_k.copy()
                disp_float_k = disp_float_k.reshape(self.substrate.nb_grid_pts)
                gap_float_k = disp_float_k - self.heights_k_float - offset_k
                self.fourier_buffer.array()[...] = gap_float_k.copy()
                self.engine.ihcfft(self.fourier_buffer, self.real_buffer)
                gap = self.real_buffer.array()[...].copy() \
                    * self.engine.normalisation

                return self.evaluate_k(disp_float_k, gap, offset, forces=True,
                                       logger=logger)[0]

        return fun

    def callback(self, force=False):
        """
        Simple callback function that can be handed over to scipy's minimize to
        get updates during minimisation
        Parameters:
        ----------
        force: bool, optional
            whether to include the norm of the force
            vector in the update message
            (default False)
        """
        counter = 0
        if force:
            def fun(dummy):
                "includes the force norm in its output"
                nonlocal counter
                counter += 1
                print("at it {}, e = {}, |f| = {}".format(
                    counter, self.energy,
                    np.linalg.norm(np.ravel(self.force))))
        else:
            def fun(dummy):
                "prints messages without force information"
                nonlocal counter
                counter += 1
                print("at it {}, e = {}".format(
                    counter, self.energy))
        return fun

class BoundedSmoothContactSystem(SmoothContactSystem):
    @staticmethod
    def handles(*args, **kwargs):  # FIXME work around, see issue #208
        return False

    def compute_nb_contact_pts(self):
        """
        compute and return the number of contact points.
        """
        return self.reduction.sum(np.where(self.gap == 0., 1., 0.))

    def logger_input(self):
        """

        Returns
        -------
        headers: list of strings
        values: list
        """
        tot_nb_grid_pts = np.prod(self.nb_grid_pts)
        rel_rep_area = self.compute_nb_repulsive_pts() / tot_nb_grid_pts
        rel_att_area = self.compute_nb_attractive_pts() / tot_nb_grid_pts
        # TODO: eventually put a flag to turn
        #  reductions off since this is an additional communication.

        contacting_points = self.gap == 0.
        mask = np.ones(self.substrate.nb_subdomain_grid_pts)
        mask[self.substrate.local_topography_subdomain_slices][
            contacting_points] = 0
        max_proj_grad = self.reduction.max(abs(mask * self.force))

        return (['energy',
                 'max. proj. grad.',
                 'mean gap',
                 'frac. cont. area',
                 'frac. rep. area',
                 'frac. att. area',
                 'frac. int. area',
                 'substrate force',
                 'interaction force'],
                [self.energy,
                 max_proj_grad,
                 self.compute_mean_gap(),
                 self.compute_nb_contact_pts() / np.prod(self.nb_grid_pts),
                 rel_rep_area,
                 rel_att_area,
                 rel_rep_area + rel_att_area,
                 -self.reduction.sum(self.substrate.force),
                 self.reduction.sum(self.interaction_force)])

    def compute_normal_force(self):
        "computes and returns the sum of all forces"

        # sum of the jacobian in the contact area (Lagrange multiplier)
        # and the ineraction forces.
        # can also be computed easily from the substrate forces,
        # what we do here
        return self.reduction.sum(
            - self.substrate.force[self.substrate.local_topography_subdomain_slices])

    def compute_repulsive_force(self):
        """computes and returns the sum of all repulsive forces

        Assumptions: there
        """
        return self.reduction.sum(
            np.where(
                - self.substrate.force[
                    self.substrate.local_topography_subdomain_slices] > 0,
                - self.substrate.force[
                    self.substrate.local_topography_subdomain_slices], 0.))

    def compute_attractive_force(self):
        "computes and returns the sum of all attractive forces"
        return self.reduction.sum(
            np.where(
                - self.substrate.force[
                    self.substrate.local_topography_subdomain_slices] < 0,
                - self.substrate.force[
                    self.substrate.local_topography_subdomain_slices],
                0.))

    def compute_nb_repulsive_pts(self):
        """
        compute and return the number of contact points under repulsive
        pressure.

        """
        return self.reduction.sum(
            np.where(
                np.logical_and(
                    self.gap == 0.,
                    - self.substrate.force[
                        self.substrate.local_topography_subdomain_slices] > 0),
                1., 0.))

    def compute_nb_attractive_pts(self):
        """
        compute and return the number of contact points under attractive
        pressure.
        """

        # Compute points where substrate force is negative
        # or there is no contact
        pts = np.logical_or(- self.substrate.force[
            self.substrate.local_topography_subdomain_slices] < 0,
                            self.gap > 0.)

        # exclude points where there is no contact
        # and the interaction force is 0.
        pts[np.logical_and(self.gap > 0.,
                           self.interaction_force == 0.)] = 0.

        return self.reduction.sum(pts)

    def compute_repulsive_coordinates(self):
        """
        returns an array of all coordinates, where contact pressure is
        repulsive. Useful for evaluating the number of contact islands etc.
        """
        return np.argwhere(
            np.logical_and(
                self.gap == 0.,
                - self.substrate.force[
                    self.substrate.local_topography_subdomain_slices]
                > 0))

    def compute_attractive_coordinates(self):
        """
        returns an array of all coordinates, where contact pressure is
        attractive. Useful for evaluating the number of contact islands etc.
        """

        # Compute points where substrate force is negative
        # or there is no contact
        pts = np.logical_or(
            - self.substrate.force[self.substrate.local_topography_subdomain_slices]
            < 0,
            self.gap > 0.)

        # exclude points where there is no contact
        # and the interaction force is 0.
        pts[np.logical_and(self.gap > 0.,
                           self.interaction_force == 0.)] = 0.

        return np.argwhere(pts)


"""
Code to export from yt to Sunrise

"""
from __future__ import print_function

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

try:
    import pyfits
except ImportError:
    pass

import time
import numpy as np
from yt.funcs import *
import yt.utilities.lib.api as amr_utils
from yt.utilities.physical_constants import \
    kpc_per_cm, \
    sec_per_year
from yt.mods import *

def export_to_sunrise(ds, fn, star_particle_type, fc, fwidth, ncells_wide=None,
        debug=False,dd=None,**kwargs):
    r"""Convert the contents of a dataset to a FITS file format that Sunrise
    understands.

    This function will accept a dataset, and from that dataset
    construct a depth-first octree containing all of the data in the parameter
    file.  This octree will be written to a FITS file.  It will probably be
    quite big, so use this function with caution!  Sunrise is a tool for
    generating synthetic spectra, available at
    http://sunrise.googlecode.com/ .

    Parameters
    ----------
    ds : `Dataset`
       The dataset to convert.
    fn : string
       The filename of the output FITS file.
    fc : array
       The center of the extraction region
    fwidth  : array
       Ensure this radius around the center is enclosed
       Array format is (nx,ny,nz) where each element is floating point
       in unitary position units where 0 is leftmost edge and 1
       the rightmost.

    Notes
    -----

    Note that the process of generating simulated images from Sunrise will
    require substantial user input; see the Sunrise wiki at
    http://sunrise.googlecode.com/ for more information.

    """
    fc = np.array(fc)
    fwidth = np.array(fwidth)

    #we must round the dle,dre to the nearest root grid cells
    ile,ire,super_level,ncells_wide= \
            round_ncells_wide(ds.domain_dimensions,fc-fwidth,fc+fwidth,nwide=ncells_wide)

    assert np.all((ile-ire)==(ile-ire)[0])
    mylog.info("rounding specified region:")
    mylog.info("from [%1.5f %1.5f %1.5f]-[%1.5f %1.5f %1.5f]"%(tuple(fc-fwidth)+tuple(fc+fwidth)))
    mylog.info("to   [%07i %07i %07i]-[%07i %07i %07i]"%(tuple(ile)+tuple(ire)))
    fle,fre = ile*1.0/ds.domain_dimensions, ire*1.0/ds.domain_dimensions
    mylog.info("to   [%1.5f %1.5f %1.5f]-[%1.5f %1.5f %1.5f]"%(tuple(fle)+tuple(fre)))

    #Create a list of the star particle properties in PARTICLE_DATA
    #Include ID, parent-ID, position, velocity, creation_mass,
    #formation_time, mass, age_m, age_l, metallicity, L_bol
    particle_data,nstars = prepare_star_particles(ds,star_particle_type,fle=fle,fre=fre,
                                           dd=dd,**kwargs)

    #Create the refinement hilbert octree in GRIDSTRUCTURE
    #For every leaf (not-refined) cell we have a column n GRIDDATA
    #Include mass_gas, mass_metals, gas_temp_m, gas_teff_m, cell_volume, SFR
    #since the octree always starts with one cell, an our 0-level mesh
    #may have many cells, we must create the octree region sitting
    #ontop of the first mesh by providing a negative level
    output, refinement,dd,nleaf = prepare_octree(ds,ile,start_level=super_level,
            debug=debug,dd=dd,center=fc)

    create_fits_file(ds,fn, refinement,output,particle_data,fle,fre)

    return fle,fre,ile,ire,dd,nleaf,nstars

def export_to_sunrise_from_halolist(ds,fni,star_particle_type,
                                        halo_list,domains_list=None,**kwargs):
    """
    Using the center of mass and the virial radius
    for a halo, calculate the regions to extract for sunrise.
    The regions are defined on the root grid, and so individual
    octs may span a large range encompassing many halos
    and subhalos. Instead of repeating the oct extraction for each
    halo, arrange halos such that we only calculate what we need to.

    Parameters
    ----------
    ds : `Dataset`
        The dataset to convert. We use the root grid to specify the domain.
    fni : string
        The filename of the output FITS file, but depends on the domain. The
        dle and dre are appended to the name.
    particle_type : int
        The particle index for stars
    halo_list : list of halo objects
        The halo list objects must have halo.CoM and halo.Rvir,
        both of which are assumed to be in unitary length units.
    frvir (optional) : float
        Ensure that CoM +/- frvir*Rvir is contained within each domain
    domains_list (optiona): dict of halos
        Organize halos into a dict of domains. Keys are DLE/DRE tuple
        values are a list of halos
    """
    dn = ds.domain_dimensions
    if domains_list is None:
        domains_list = domains_from_halos(ds,halo_list,**kwargs)
    if fni.endswith('.fits'):
        fni = fni.replace('.fits','')

    for (num_halos, domain, halos) in domains_list:
        dle,dre = domain
        print('exporting: ')
        print("[%03i %03i %03i] -"%tuple(dle), end=' ')
        print("[%03i %03i %03i] "%tuple(dre), end=' ')
        print(" with %i halos"%num_halos)
        dle,dre = domain
        dle, dre = np.array(dle),np.array(dre)
        fn = fni
        fn += "%03i_%03i_%03i-"%tuple(dle)
        fn += "%03i_%03i_%03i"%tuple(dre)
        fnf = fn + '.fits'
        fnt = fn + '.halos'
        if os.path.exists(fnt):
            os.remove(fnt)
        fh = open(fnt,'w')
        for halo in halos:
            fh.write("%i "%halo.ID)
            fh.write("%6.6e "%(halo.CoM[0]*ds['kpc']))
            fh.write("%6.6e "%(halo.CoM[1]*ds['kpc']))
            fh.write("%6.6e "%(halo.CoM[2]*ds['kpc']))
            fh.write("%6.6e "%(halo.Mvir))
            fh.write("%6.6e \n"%(halo.Rvir*ds['kpc']))
        fh.close()
        export_to_sunrise(ds, fnf, star_particle_type, dle*1.0/dn, dre*1.0/dn)

def domains_from_halos(ds,halo_list,frvir=0.15):
    domains = {}
    dn = ds.domain_dimensions
    for halo in halo_list:
        fle, fre = halo.CoM-frvir*halo.Rvir,halo.CoM+frvir*halo.Rvir
        dle,dre = np.floor(fle*dn), np.ceil(fre*dn)
        dle,dre = tuple(dle.astype('int')),tuple(dre.astype('int'))
        if (dle,dre) in domains.keys():
            domains[(dle,dre)] += halo,
        else:
            domains[(dle,dre)] = [halo,]
    #for niceness, let's process the domains in order of
    #the one with the most halos
    domains_list = [(len(v),k,v) for k,v in domains.iteritems()]
    domains_list.sort()
    domains_list.reverse() #we want the most populated domains first
    return domains_list

def prepare_octree(ds,ile,start_level=0,debug=True,dd=None,center=None):
    if dd is None:
        #we keep passing dd around to not regenerate the data all the time
        dd = ds.all_data()
    try:
        dd['MetalMass']
    except KeyError:
        add_fields() #add the metal mass field that sunrise wants
    def _temp_times_mass(field, data):
        return data["Temperature"]*data["CellMassMsun"]
    add_field("TemperatureTimesCellMassMsun", function=_temp_times_mass)
    fields = ["CellMassMsun","TemperatureTimesCellMassMsun",
              "MetalMass","CellVolumeCode"]

    #gather the field data from octs
    pbar = get_pbar("Retrieving field data",len(fields))
    field_data = []
    for fi,f in enumerate(fields):
        field_data += dd[f],
        pbar.update(fi)
    pbar.finish()
    del field_data

    #first we cast every cell as an oct
    #ngrids = np.max([g.id for g in ds._grids])
    grids = {}
    levels_all = {}
    levels_finest = {}
    for l in range(100):
        levels_finest[l]=0
        levels_all[l]=0
    pbar = get_pbar("Initializing octs ",len(ds.index.grids))
    for gi,g in enumerate(ds.index.grids):
        ff = np.array([g[f] for f in fields])
        og = amr_utils.OctreeGrid(
                g.child_index_mask.astype('int32'),
                ff.astype("float64"),
                g.LeftEdge.astype("float64"),
                g.ActiveDimensions.astype("int32"),
                np.ones(1,dtype="float64")*g.dds[0],
                g.Level,
                g.id)
        grids[g.id] = og
        #how many refinement cells will we have?
        #measure the 'volume' of each mesh, but many
        #cells do not exist. an overstimate
        levels_all[g.Level] += g.ActiveDimensions.prod()
        #how many leaves do we have?
        #this overestimates. a child of -1 means no child,
        #but that cell may still be expanded on a submesh because
        #(at least in ART) the meshes are inefficient.
        g.clear_data()
        pbar.update(gi)
    pbar.finish()

    #create the octree grid list
    #oct_list =  amr_utils.OctreeGridList(grids)

    #initialize arrays to be passed to the recursion algo
    o_length = np.sum(levels_all.values())
    r_length = np.sum(levels_all.values())
    output   = np.zeros((o_length,len(fields)), dtype='float64')
    refined  = np.zeros(r_length, dtype='int32')
    levels   = np.zeros(r_length, dtype='int32')
    ids      = np.zeros(r_length, dtype='int32')
    pos = position()
    hs       = hilbert_state()
    start_time = time.time()
    if debug:
        printing = lambda x: print_oct(x)
    else:
        printing = None
    pbar = get_pbar("Building Hilbert DFO octree",len(refined))
    RecurseOctreeDepthFirstHilbert(
            ile,
            pos,
            grids[0], #we always start on the root grid
            hs,
            output,refined,levels,
            grids,
            start_level,
            ids,
            debug=printing,
            tracker=pbar)
    pbar.finish()
    #by time we get it here the 'current' position is actually
    #for the next spot, so we're off by 1
    print('took %1.2e seconds'%(time.time()-start_time))
    print('refinement tree # of cells %i, # of leaves %i'%(pos.refined_pos,pos.output_pos))
    print('first few entries :',refined[:12])
    output  = output[:pos.output_pos]
    refined = refined[:pos.refined_pos]
    levels = levels[:pos.refined_pos]
    return output,refined,dd,pos.refined_pos

def print_oct(data,nd=None,nc=None):
    ci = data['cell_index']
    l  = data['level']
    g  = data['grid']
    o  = g.offset
    fle = g.left_edges+g.dx*ci
    fre = g.left_edges+g.dx*(ci+1)
    if nd is not None:
        fle *= nd
        fre *= nd
        if nc is not None:
            fle -= nc
            fre -= nc
    txt  = '%+1i '
    txt += '%+1i '
    txt += '%+1.3f '*3+'- '
    txt += '%+1.3f '*3
    if l<2:
        print(txt%((l,)+(o,)+tuple(fle)+tuple(fre)))

def RecurseOctreeDepthFirstHilbert(cell_index, #integer (rep as a float) on the [grid_index]
                            pos, #the output hydro data position and refinement position
                            grid,  #grid that this oct lives on (not its children)
                            hilbert,  #the hilbert state
                            output, #holds the hydro data
                            refined, #holds the refinement status  of Octs, 0s and 1s
                            levels, #For a given Oct, what is the level
                            grids, #list of all patch grids available to us
                            level, #starting level of the oct (not the children)
                            ids, #record the oct ID
                            debug=None,tracker=True):
    if tracker is not None:
        if pos.refined_pos%1000 == 500 : tracker.update(pos.refined_pos)
    if debug is not None:
        debug(vars())
    child_grid_index = grid.child_indices[cell_index[0],cell_index[1],cell_index[2]]
    #record the refinement state
    levels[pos.refined_pos]  = level
    is_leaf = (child_grid_index==-1) and (level>0)
    refined[pos.refined_pos] = not is_leaf #True is oct, False is leaf
    ids[pos.refined_pos] = child_grid_index #True is oct, False is leaf
    pos.refined_pos+= 1
    if is_leaf: #never subdivide if we are on a superlevel
        #then we have hit a leaf cell; write it out
        for field_index in range(grid.fields.shape[0]):
            output[pos.output_pos,field_index] = \
                    grid.fields[field_index,cell_index[0],cell_index[1],cell_index[2]]
        pos.output_pos+= 1
    else:
        assert child_grid_index>-1
        #find the grid we descend into
        #then find the eight cells we break up into
        subgrid = grids[child_grid_index]
        #calculate the floating point LE of the children
        #then translate onto the subgrid integer index
        parent_fle  = grid.left_edges + cell_index*grid.dx
        subgrid_ile = np.floor((parent_fle - subgrid.left_edges)/subgrid.dx)
        for (vertex, hilbert_child) in hilbert:
            #vertex is a combination of three 0s and 1s to
            #denote each of the 8 octs
            if level < 0:
                subgrid = grid #we don't actually descend if we're a superlevel
                #child_ile = cell_index + np.array(vertex)*2**(-level)
                child_ile = cell_index + np.array(vertex)*2**(-(level+1))
                child_ile = child_ile.astype('int')
            else:
                child_ile = subgrid_ile+np.array(vertex)
                child_ile = child_ile.astype('int')

            RecurseOctreeDepthFirstHilbert(child_ile,pos,
                subgrid,hilbert_child,output,refined,levels,grids,
                level+1,ids = ids,
                debug=debug,tracker=tracker)

def create_fits_file(ds,fn, refined,output,particle_data,fle,fre):
    #first create the grid structure
    structure = pyfits.Column("structure", format="B", array=refined.astype("bool"))
    cols = pyfits.ColDefs([structure])
    st_table = pyfits.new_table(cols)
    st_table.name = "GRIDSTRUCTURE"
    st_table.header.update("hierarch lengthunit", "kpc", comment="Length unit for grid")
    fdx = fre-fle
    for i,a in enumerate('xyz'):
        st_table.header.update("min%s" % a, fle[i] * ds['kpc'])
        st_table.header.update("max%s" % a, fre[i] * ds['kpc'])
        st_table.header.update("n%s" % a, fdx[i])
        st_table.header.update("subdiv%s" % a, 2)
    st_table.header.update("subdivtp", "OCTREE", "Type of grid subdivision")

    #not the hydro grid data
    fields = ["CellMassMsun","TemperatureTimesCellMassMsun",
              "MetalMass","CellVolumeCode"]
    fd = {}
    for i,f in enumerate(fields):
        fd[f]=output[:,i]
    del output
    col_list = []
    size = fd["CellMassMsun"].size
    tm = fd["CellMassMsun"].sum()
    col_list.append(pyfits.Column("mass_gas", format='D',
                    array=fd['CellMassMsun'], unit="Msun"))
    col_list.append(pyfits.Column("mass_metals", format='D',
                    array=fd['MetalMass'], unit="Msun"))
    # col_list.append(pyfits.Column("mass_stars", format='D',
    #                 array=np.zeros(size,dtype='D'),unit="Msun"))
    # col_list.append(pyfits.Column("mass_stellar_metals", format='D',
    #                 array=np.zeros(size,dtype='D'),unit="Msun"))
    # col_list.append(pyfits.Column("age_m", format='D',
    #                 array=np.zeros(size,dtype='D'),unit="yr*Msun"))
    # col_list.append(pyfits.Column("age_l", format='D',
    #                 array=np.zeros(size,dtype='D'),unit="yr*Msun"))
    # col_list.append(pyfits.Column("L_bol", format='D',
    #                 array=np.zeros(size,dtype='D')))
    # col_list.append(pyfits.Column("L_lambda", format='D',
    #                 array=np.zeros(size,dtype='D')))
    # The units for gas_temp are really K*Msun. For older Sunrise versions
    # you must set the unit to just K
    col_list.append(pyfits.Column("gas_temp_m", format='D',
                    array=fd['TemperatureTimesCellMassMsun'], unit="K*Msun"))
    col_list.append(pyfits.Column("gas_teff_m", format='D',
                    array=fd['TemperatureTimesCellMassMsun'], unit="K*Msun"))
    col_list.append(pyfits.Column("cell_volume", format='D',
                    array=fd['CellVolumeCode'].astype('float64')*ds['kpc']**3.0,
                    unit="kpc^3"))
    col_list.append(pyfits.Column("SFR", format='D',
                    array=np.zeros(size, dtype='D')))
    cols = pyfits.ColDefs(col_list)
    mg_table = pyfits.new_table(cols)
    mg_table.header.update("M_g_tot", tm)
    mg_table.header.update("timeunit", "yr")
    mg_table.header.update("tempunit", "K")
    mg_table.name = "GRIDDATA"

    # Add a dummy Primary; might be a better way to do this!
    col_list = [pyfits.Column("dummy", format="F", array=np.zeros(1, dtype='float32'))]
    cols = pyfits.ColDefs(col_list)
    md_table = pyfits.new_table(cols)
    md_table.header.update("snaptime", ds.current_time*ds['years'])
    md_table.name = "YT"

    phdu = pyfits.PrimaryHDU()
    phdu.header.update('nbodycod','yt')
    hls = [phdu, st_table, mg_table,md_table]
    hls.append(particle_data)
    hdus = pyfits.HDUList(hls)
    hdus.writeto(fn, clobber=True)

def nearest_power(x):
    #round to the nearest power of 2
    x-=1
    x |= x >> 1
    x |= x >> 2
    x |= x >> 4
    x |= x >> 8
    x |= x >> 16
    x+=1
    return x

def round_ncells_wide(dds,fle,fre,nwide=None):
    fc = (fle+fre)/2.0
    assert np.all(fle < fc)
    assert np.all(fre > fc)
    ic = np.rint(fc*dds) #nearest vertex to the center
    ile,ire = ic.astype('int'),ic.astype('int')
    cfle,cfre = fc.copy(),fc.copy()
    idx = np.array([0,0,0]) #just a random non-equal array
    width = 0.0
    if nwide is None:
        #expand until borders are included and
        #we have an equaly-sized, non-zero box
        idxq,out=False,True
        while not out or not idxq:
            cfle,cfre = fc-width, fc+width
            ile = np.rint(cfle*dds).astype('int')
            ire = np.rint(cfre*dds).astype('int')
            idx = ire-ile
            width += 0.1/dds
            #quit if idxq is true:
            idxq = idx[0]>0 and np.all(idx==idx[0])
            out  = np.all(fle>cfle) and np.all(fre<cfre)
            out &= abs(np.log2(idx[0])-np.rint(np.log2(idx[0])))<1e-5 #nwide should be a power of 2
            assert width[0] < 1.1 #can't go larger than the simulation volume
        nwide = idx[0]
    else:
        #expand until we are nwide cells span
        while not np.all(idx==nwide):
            assert np.any(idx<=nwide)
            cfle,cfre = fc-width, fc+width
            ile = np.rint(cfle*dds).astype('int')
            ire = np.rint(cfre*dds).astype('int')
            idx = ire-ile
            width += 1e-2*1.0/dds
    assert np.all(idx==nwide)
    assert idx[0]>0
    maxlevel = -np.rint(np.log2(nwide)).astype('int')
    assert abs(np.log2(nwide)-np.rint(np.log2(nwide)))<1e-5 #nwide should be a power of 2
    return ile,ire,maxlevel,nwide

def round_nearest_edge(ds,fle,fre):
    dds = ds.domain_dimensions
    ile = np.floor(fle*dds).astype('int')
    ire = np.ceil(fre*dds).astype('int')

    #this is the number of cells the super octree needs to expand to
    #must round to the nearest power of 2
    width = np.max(ire-ile)
    width = nearest_power(width)

    maxlevel = -np.rint(np.log2(width)).astype('int')
    return ile,ire,maxlevel

def prepare_star_particles(ds,star_type,pos=None,vel=None, age=None,
                          creation_time=None,initial_mass=None,
                          current_mass=None,metallicity=None,
                          radius = None,
                          fle=[0.,0.,0.],fre=[1.,1.,1.],
                          dd=None):
    if dd is None:
        dd = ds.all_data()
    idxst = dd["particle_type"] == star_type

    #make sure we select more than a single particle
    assert na.sum(idxst)>0
    if pos is None:
        pos = np.array([dd["particle_position_%s" % ax]
                        for ax in 'xyz']).transpose()
    idx = idxst & np.all(pos>fle,axis=1) & np.all(pos<fre,axis=1)
    assert np.sum(idx)>0
    pos = pos[idx]*ds['kpc'] #unitary units -> kpc
    if age is None:
        age = dd["particle_age"][idx]*ds['years'] # seconds->years
    if vel is None:
        vel = np.array([dd["particle_velocity_%s" % ax][idx]
                        for ax in 'xyz']).transpose()
        # Velocity is cm/s, we want it to be kpc/yr
        #vel *= (ds["kpc"]/ds["cm"]) / (365*24*3600.)
        vel *= kpc_per_cm * sec_per_year
    if initial_mass is None:
        #in solar masses
        initial_mass = dd["particle_mass_initial"][idx]*ds['Msun']
    if current_mass is None:
        #in solar masses
        current_mass = dd["particle_mass"][idx]*ds['Msun']
    if metallicity is None:
        #this should be in dimensionless units, metals mass / particle mass
        metallicity = dd["particle_metallicity"][idx]
        assert np.all(metallicity>0.0)
    if radius is None:
        radius = initial_mass*0.0+10.0/1000.0 #10pc radius
    formation_time = ds.current_time*ds['years']-age
    #create every column
    col_list = []
    col_list.append(pyfits.Column("ID", format="J", array=np.arange(current_mass.size).astype('int32')))
    col_list.append(pyfits.Column("parent_ID", format="J", array=np.arange(current_mass.size).astype('int32')))
    col_list.append(pyfits.Column("position", format="3D", array=pos, unit="kpc"))
    col_list.append(pyfits.Column("velocity", format="3D", array=vel, unit="kpc/yr"))
    col_list.append(pyfits.Column("creation_mass", format="D", array=initial_mass, unit="Msun"))
    col_list.append(pyfits.Column("formation_time", format="D", array=formation_time, unit="yr"))
    col_list.append(pyfits.Column("radius", format="D", array=radius, unit="kpc"))
    col_list.append(pyfits.Column("mass", format="D", array=current_mass, unit="Msun"))
    col_list.append(pyfits.Column("age", format="D", array=age,unit='yr'))
    #For particles, Sunrise takes
    #the dimensionless metallicity, not the mass of the metals
    col_list.append(pyfits.Column("metallicity", format="D",
        array=metallicity,unit="Msun"))

    #make the table
    cols = pyfits.ColDefs(col_list)
    pd_table = pyfits.new_table(cols)
    pd_table.name = "PARTICLEDATA"

    #make sure we have nonzero particle number
    assert pd_table.data.shape[0]>0
    return pd_table,na.sum(idx)

def add_fields():
    """Add three Eulerian fields Sunrise uses"""
    def _MetalMass(field, data):
        return data["Metallicity"] * data["CellMassMsun"]

    def _convMetalMass(data):
        return 1.0
    add_field("MetalMass", function=_MetalMass,
              convert_function=_convMetalMass)
    def _initial_mass_cen_ostriker(field, data):
        # SFR in a cell. This assumes stars were created by the Cen & Ostriker algorithm
        # Check Grid_AddToDiskProfile.C and star_maker7.src
        star_mass_ejection_fraction = data.ds.get_parameter("StarMassEjectionFraction",float)
        star_maker_minimum_dynamical_time = 3e6 # years, which will get divided out
        dtForSFR = star_maker_minimum_dynamical_time / data.ds["years"]
        xv1 = ((data.ds["InitialTime"] - data["creation_time"])
                / data["dynamical_time"])
        xv2 = ((data.ds["InitialTime"] + dtForSFR - data["creation_time"])
                / data["dynamical_time"])
        denom = (1.0 - star_mass_ejection_fraction * (1.0 - (1.0 + xv1)*np.exp(-xv1)))
        minitial = data["ParticleMassMsun"] / denom
        return minitial

    add_field("InitialMassCenOstriker", function=_initial_mass_cen_ostriker)

class position:
    def __init__(self):
        self.output_pos = 0
        self.refined_pos = 0

class hilbert_state():
    def __init__(self,dim=None,sgn=None,octant=None):
        if dim is None: dim = [0,1,2]
        if sgn is None: sgn = [1,1,1]
        if octant is None: octant = 5
        self.dim = dim
        self.sgn = sgn
        self.octant = octant
    def flip(self,i):
        self.sgn[i]*=-1
    def swap(self,i,j):
        temp = self.dim[i]
        self.dim[i]=self.dim[j]
        self.dim[j]=temp
        axis = self.sgn[i]
        self.sgn[i] = self.sgn[j]
        self.sgn[j] = axis
    def reorder(self,i,j,k):
        ndim = [self.dim[i],self.dim[j],self.dim[k]]
        nsgn = [self.sgn[i],self.sgn[j],self.sgn[k]]
        self.dim = ndim
        self.sgn = nsgn
    def copy(self):
        return hilbert_state([self.dim[0],self.dim[1],self.dim[2]],
                             [self.sgn[0],self.sgn[1],self.sgn[2]],
                             self.octant)
    def descend(self,o):
        child = self.copy()
        child.octant = o
        if o==0:
            child.swap(0,2)
        elif o==1:
            child.swap(1,2)
        elif o==2:
            pass
        elif o==3:
            child.flip(0)
            child.flip(2)
            child.reorder(2,0,1)
        elif o==4:
            child.flip(0)
            child.flip(1)
            child.reorder(2,0,1)
        elif o==5:
            pass
        elif o==6:
            child.flip(1)
            child.flip(2)
            child.swap(1,2)
        elif o==7:
            child.flip(0)
            child.flip(2)
            child.swap(0,2)
        return child

    def __iter__(self):
        vertex = [0,0,0]
        j=0
        for i in range(3):
            vertex[self.dim[i]] = 0 if self.sgn[i]>0 else 1
        yield vertex, self.descend(j)
        vertex[self.dim[0]] += self.sgn[0]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[1]] += self.sgn[1]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[0]] -= self.sgn[0]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[2]] += self.sgn[2]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[0]] += self.sgn[0]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[1]] -= self.sgn[1]
        j+=1
        yield vertex, self.descend(j)
        vertex[self.dim[0]] -= self.sgn[0]
        j+=1
        yield vertex, self.descend(j)