################################################################################
# Copyright (C) 2013-2014 Jaakko Luttinen
#
# This file is licensed under the MIT License.
################################################################################


import numpy as np
import functools

from bayespy.utils import misc

"""
This module contains a sketch of a new implementation of the framework.
"""

def message_sum_multiply(plates_parent, dims_parent, *arrays):
    """
    Compute message to parent and sum over plates.

    Divide by the plate multiplier.
    """
    # The shape of the full message
    shapes = [np.shape(array) for array in arrays]
    shape_full = misc.broadcasted_shape(*shapes)
    # Find axes that should be summed
    shape_parent = plates_parent + dims_parent
    sum_axes = misc.axes_to_collapse(shape_full, shape_parent)
    # Compute the multiplier for cancelling the
    # plate-multiplier.  Because we are summing over the
    # dimensions already in this function (for efficiency), we
    # need to cancel the effect of the plate-multiplier
    # applied in the message_to_parent function.
    r = 1
    for j in sum_axes:
        if j >= 0 and j < len(plates_parent):
            r *= shape_full[j]
        elif j < 0 and j < -len(dims_parent):
            r *= shape_full[j]
    # Compute the sum-product
    m = misc.sum_multiply(*arrays,
                          axis=sum_axes,
                          sumaxis=True,
                          keepdims=True) / r
    # Remove extra axes
    m = misc.squeeze_to_dim(m, len(shape_parent))
    return m


class Moments():
    """
    Base class for defining the expectation of the sufficient statistics.

    The benefits:

      * Write statistic-specific features in one place only. For instance,
        covariance from Gaussian message.

      * Different nodes may have identically defined statistic so you need to
        implement related features only once. For instance, Gaussian and
        GaussianARD differ on the prior but the moments are the same.

      * General processing nodes which do not change the type of the moments may
        "inherit" the features from the parent node. For instance, slicing
        operator.

      * Conversions can be done easily in both of the above cases if the message
        conversion is defined in the moments class. For instance,
        GaussianMarkovChain to Gaussian and VaryingGaussianMarkovChain to
        Gaussian.
    """

    _converters = {}


    class NoConverterError(Exception):
        pass


    def get_instance_converter(self, **kwargs):
        """Default converter within a moments class is an identity.

        Override this method when moment class instances are not identical if
        they have different attributes.

        """
        if len(kwargs) > 0:
            raise NotImplementedError(
                "get_instance_converter not implemented for class {0}"
                .format(self.__class__.__name__)
            )
        return None


    def get_instance_conversion_kwargs(self):
        """
        Override this method when moment class instances are not identical if
        they have different attributes.
        """
        return {}


    @classmethod
    def add_converter(cls, moments_to, converter):
        cls._converters = cls._converters.copy()
        cls._converters[moments_to] = converter
        return


    def get_converter(self, moments_to):
        """
        Finds conversion to another moments type if possible.

        Note that a conversion from moments A to moments B may require
        intermediate conversions.  For instance: A->C->D->B.  This method finds
        the path which uses the least amount of conversions and returns that
        path as a single conversion.  If no conversion path is available, an
        error is raised.

        The search algorithm starts from the original moments class and applies
        all possible converters to get a new list of moments classes. This list
        is extended by adding recursively all parent classes because their
        converters are applicable. Then, all possible converters are applied to
        this list to get a new list of current moments classes. This is iterated
        until the algorithm hits the target moments class or its subclass.
        """

        # Check if there is no need for a conversion
        #
        # TODO/FIXME: This isn't sufficient. Moments can have attributes that
        # make them incompatible (e.g., ndim in GaussianMoments).
        if isinstance(self, moments_to):
            return lambda X: X

        # Initialize variables
        visited = set()
        visited.add(self.__class__)
        converted_list = [(self.__class__, [])]

        # Each iteration step consists of two parts:
        # 1) form a set of the current classes and all their parent classes
        #    recursively
        # 2) from the current set, apply possible conversions to get a new set
        #    of classes
        # Repeat these two steps until in step (1) you hit the target class.

        while len(converted_list) > 0:
            # Go through all parents recursively so we can then use all
            # converters that are available
            current_list = []
            for (moments_class, converter_path) in converted_list:
                if issubclass(moments_class, moments_to):
                    # Shortest conversion path found, return the resulting total
                    # conversion function
                    return misc.composite_function(converter_path)
                current_list.append((moments_class, converter_path))
                parents = list(moments_class.__bases__)
                for parent in parents:
                    # Recursively add parents
                    for p in parent.__bases__:
                        if isinstance(p, Moments):
                            parents.append(p)
                    # Add un-visited parents
                    if issubclass(parent, Moments) and parent not in visited:
                        visited.add(parent)
                        current_list.append((parent, converter_path))

            # Find all converters and extend the converter paths
            converted_list = []
            for (moments_class, converter_path) in current_list:
                for (conv_mom_cls, conv) in moments_class._converters.items():
                    if conv_mom_cls not in visited:
                        visited.add(conv_mom_cls)
                        converted_list.append((conv_mom_cls,
                                               converter_path + [conv]))

        raise self.NoConverterError("No conversion defined from %s to %s"
                                    % (self.__class__.__name__,
                                       moments_to.__name__))


    def compute_fixed_moments(self, x):
        # This method can't be static because the computation of the moments may
        # depend on, for instance, ndim in Gaussian arrays.
        raise NotImplementedError("compute_fixed_moments not implemented for "
                                  "%s"
                                  % (self.__class__.__name__))


    @classmethod
    def from_values(cls, x):
        raise NotImplementedError("from_values not implemented "
                                  "for %s"
                                  % (cls.__name__))

def ensureparents(func):
    @functools.wraps(func)
    def wrapper(self, *parents, **kwargs):
        # Convert parents to proper nodes
        if self._parent_moments is None:
            raise ValueError(
                "Parent moments must be defined for {0}"
                .format(self.__class__.__name__)
            )
        parents = [
            Node._ensure_moments(
                parent,
                moments.__class__,
                **moments.get_instance_conversion_kwargs()
            )
            for (parent, moments) in zip(parents, self._parent_moments)
        ]
        # parents = list(parents)
        # for (ind, parent) in enumerate(parents):
        #     parents[ind] = self._ensure_moments(parent,
        #                                         self._parent_moments[ind])
        # Run the function
        return func(self, *parents, **kwargs)

    return wrapper


class Node():
    """
    Base class for all nodes.

    mask
    dims
    plates
    parents
    children
    name

    Sub-classes must implement:
    1. For computing the message to children:
       get_moments(self):
    2. For computing the message to parents:
       _get_message_and_mask_to_parent(self, index)

    Sub-classes may need to re-implement:
    1. If they manipulate plates:
       _compute_weights_to_parent(index, weights)
       _plates_to_parent(self, index)
       _plates_from_parent(self, index)
    """

    # These are objects of the _parent_moments_class. If the default way of
    # creating them is not correct, write your own creation code.
    _moments = None
    _parent_moments = None
    plates = None

    _id_counter = 0

    @ensureparents
    def __init__(self, *parents, dims=None, plates=None, name="",
                 notify_parents=True, plotter=None, plates_multiplier=None,
                 allow_dependent_parents=False):

        self.parents = parents
        self.dims = dims
        self.name = name
        self._plotter = plotter

        if not allow_dependent_parents:
            parent_id_list = []
            for parent in parents:
                parent_id_list = parent_id_list + list(parent._get_id_list())
            if len(parent_id_list) != len(set(parent_id_list)):
                raise ValueError("Parent nodes are not independent")

        # Inform parent nodes
        if notify_parents:
            for (index,parent) in enumerate(self.parents):
                parent._add_child(self, index)

        # Check plates
        parent_plates = [self._plates_from_parent(index)
                         for index in range(len(self.parents))]
        if any(p is None for p in parent_plates):
            raise ValueError("Method _plates_from_parent returned None")

        # Get and validate the plates for this node
        plates = self._total_plates(plates, *parent_plates)
        if self.plates is None:
            self.plates = plates

        # By default, ignore all plates
        self.mask = np.array(False)

        # Children
        self.children = set()

        # Get and validate the plate multiplier
        parent_plates_multiplier = [self._plates_multiplier_from_parent(index)
                                   for index in range(len(self.parents))]
        #if plates_multiplier is None:
        #    plates_multiplier = parent_plates_multiplier
        plates_multiplier = self._total_plates(plates_multiplier,
                                              *parent_plates_multiplier)
        self.plates_multiplier = plates_multiplier


    def get_pdf_nodes(self):
        return tuple(
            node
            for (child, _) in self.children
            for node in child._get_pdf_nodes_conditioned_on_parents()
        )


    def _get_pdf_nodes_conditioned_on_parents(self):
        return self.get_pdf_nodes()


    def _get_id_list(self):
        """
        Returns the stochastic ID list.

        This method is used to check that same stochastic nodes are not direct
        parents of a node several times. It is only valid if there are
        intermediate stochastic nodes.

        To put it another way: each ID corresponds to one factor q(..) in the
        posterior approximation. Different IDs mean different factors, thus they
        mean independence. The parents must have independent factors.

        Stochastic nodes should return their unique ID. Deterministic nodes
        should return the IDs of their parents. Constant nodes should return
        empty list of IDs.
        """
        raise NotImplementedError()


    @classmethod
    def _total_plates(cls, plates, *parent_plates):
        if plates is None:
            # By default, use the minimum number of plates determined
            # from the parent nodes
            try:
                return misc.broadcasted_shape(*parent_plates)
            except ValueError:
                raise ValueError(
                    "The plates of the parents do not broadcast: {0}".format(
                        parent_plates
                    )
                )
        else:
            # Check that the parent_plates are a subset of plates.
            for (ind, p) in enumerate(parent_plates):
                if not misc.is_shape_subset(p, plates):
                    raise ValueError("The plates %s of the parents "
                                     "are not broadcastable to the given "
                                     "plates %s."
                                     % (p,
                                        plates))
            return plates


    @staticmethod
    def _ensure_moments(node, moments_class, **kwargs):
        try:
            converter = node._moments.get_converter(moments_class)
        except AttributeError:
            from .constant import Constant
            return Constant(
                moments_class.from_values(node, **kwargs),
                node
            )
        else:
            node = converter(node)
            converter = node._moments.get_instance_converter(**kwargs)
            if converter is not None:
                from .converters import NodeConverter
                return NodeConverter(converter, node)
            return node


    def _compute_plates_to_parent(self, index, plates):
        # Sub-classes may want to overwrite this if they manipulate plates
        return plates


    def _compute_plates_from_parent(self, index, plates):
        # Sub-classes may want to overwrite this if they manipulate plates
        return plates


    def _compute_plates_multiplier_from_parent(self, index, plates_multiplier):
        # TODO/FIXME: How to handle this properly?
        return plates_multiplier


    def _plates_to_parent(self, index):
        return self._compute_plates_to_parent(index, self.plates)


    def _plates_from_parent(self, index):
        return self._compute_plates_from_parent(index,
                                                self.parents[index].plates)


    def _plates_multiplier_from_parent(self, index):
        return self._compute_plates_multiplier_from_parent(
            index,
            self.parents[index].plates_multiplier
        )


    @property
    def plates_multiplier(self):
        """ Plate multiplier is applied to messages to parents """
        return self.__plates_multiplier


    @plates_multiplier.setter
    def plates_multiplier(self, value):
        # TODO/FIXME: Check that multiplier is consistent with plates
        self.__plates_multiplier = value
        return


    def get_shape(self, ind):
        return self.plates + self.dims[ind]

    def _add_child(self, child, index):
        """
        Add a child node.

        Parameters
        ----------
        child : node
        index : int
           The parent index of this node for the child node.
           The child node recognizes its parents by their index
           number.
        """
        self.children.add((child, index))

    def _remove_child(self, child, index):
        """
        Remove a child node.
        """
        self.children.remove((child, index))

    def get_mask(self):
        return self.mask

    ## def _get_message_mask(self):
    ##     return self.mask

    def _set_mask(self, mask):
        # Sub-classes may overwrite this method if they have some other masks to
        # be combined (for instance, observation mask)
        self.mask = mask

    def _update_mask(self):
        # Combine masks from children
        mask = np.array(False)
        for (child, index) in self.children:
            mask = np.logical_or(mask, child._mask_to_parent(index))
        # Set the mask of this node
        self._set_mask(mask)
        if not misc.is_shape_subset(np.shape(self.mask), self.plates):

            raise ValueError("The mask of the node %s has updated "
                             "incorrectly. The plates in the mask %s are not a "
                             "subset of the plates of the node %s."
                             % (self.name,
                                np.shape(self.mask),
                                self.plates))

        # Tell parents to update their masks
        for parent in self.parents:
            parent._update_mask()


    def _compute_weights_to_parent(self, index, weights):
        """Compute the mask used for messages sent to parent[index].

        The mask tells which plates in the messages are active. This method is
        used for obtaining the mask which is used to set plates in the messages
        to parent to zero.

        Sub-classes may want to overwrite this method if they do something to
        plates so that the mask is somehow altered.

        """
        return weights


    def _mask_to_parent(self, index):
        """
        Get the mask with respect to parent[index].

        The mask tells which plate connections are active. The mask is "summed"
        (logical or) and reshaped into the plate shape of the parent. Thus, it
        can't be used for masking messages, because some plates have been summed
        already. This method is used for propagating the mask to parents.
        """
        mask = self._compute_weights_to_parent(index, self.mask) != 0

        # Check the shape of the mask
        plates_to_parent = self._plates_to_parent(index)
        if not misc.is_shape_subset(np.shape(mask), plates_to_parent):
            raise ValueError("In node %s, the mask being sent to "
                             "parent[%d] (%s) has invalid shape: The shape of "
                             "the mask %s is not a sub-shape of the plates of "
                             "the node with respect to the parent %s. It could "
                             "be that this node (%s) is manipulating plates "
                             "but has not overwritten the method "
                             "_compute_weights_to_parent."
                             % (self.name,
                                index,
                                self.parents[index].name,
                                np.shape(mask),
                                plates_to_parent,
                                self.__class__.__name__))

        # "Sum" (i.e., logical or) over the plates that have unit length in
        # the parent node.
        parent_plates = self.parents[index].plates
        s = misc.axes_to_collapse(np.shape(mask), parent_plates)
        mask = np.any(mask, axis=s, keepdims=True)
        mask = misc.squeeze_to_dim(mask, len(parent_plates))
        return mask

    def _message_to_child(self):

        u = self.get_moments()

        # Debug: Check that the message has appropriate shape
        for (ui, dim) in zip(u, self.dims):
            ndim = len(dim)
            if ndim > 0:
                if np.shape(ui)[-ndim:] != dim:
                    raise RuntimeError(
                        "A bug found by _message_to_child for %s: "
                        "The variable axes of the moments %s are not equal to "
                        "the axes %s defined by the node %s. A possible reason "
                        "is that the plates of the node are inferred "
                        "incorrectly from the parents, and the method "
                        "_plates_from_parents should be implemented."
                        % (self.__class__.__name__,
                           np.shape(ui)[-ndim:],
                           dim,
                           self.name))
                if not misc.is_shape_subset(np.shape(ui)[:-ndim],
                                            self.plates):
                    raise RuntimeError(
                        "A bug found by _message_to_child for %s: "
                        "The plate axes of the moments %s are not a subset of "
                        "the plate axes %s defined by the node %s."
                        % (self.__class__.__name__,
                           np.shape(ui)[:-ndim],
                           self.plates,
                           self.name))
            else:
                if not misc.is_shape_subset(np.shape(ui), self.plates):
                    raise RuntimeError(
                        "A bug found by _message_to_child for %s: "
                        "The plate axes of the moments %s are not a subset of "
                        "the plate axes %s defined by the node %s."
                        % (self.__class__.__name__,
                           np.shape(ui),
                           self.plates,
                           self.name))
        return u

    def _message_to_parent(self, index, u_parent=None):

        # Compute the message, check plates, apply mask and sum over some plates
        if index >= len(self.parents):
            raise ValueError("Parent index larger than the number of parents")

        # Compute the message and mask
        (m, mask) = self._get_message_and_mask_to_parent(index, u_parent=u_parent)
        mask = misc.squeeze(mask)

        # Plates in the mask
        plates_mask = np.shape(mask)

        # The parent we're sending the message to
        parent = self.parents[index]

        # Plates with respect to the parent
        plates_self = self._plates_to_parent(index)

        # Plate multiplier of the parent
        multiplier_parent = self._plates_multiplier_from_parent(index)

        # Check if m is a logpdf function (for black-box variational inference)
        if callable(m):
            return m

            def m_function(*args):
                lpdf = m(*args)
                # Log pdf only contains plate axes!
                plates_m = np.shape(lpdf)
                r = (self.broadcasting_multiplier(plates_self,
                                                  plates_m,
                                                  plates_mask,
                                                  parent.plates) *
                     self.broadcasting_multiplier(self.plates_multiplier,
                                                  multiplier_parent))
                axes_msg = misc.axes_to_collapse(plates_m, parent.plates)
                m[i] = misc.sum_multiply(mask_i, m[i], r,
                                         axis=axes_msg,
                                         keepdims=True)

                # Remove leading singular plates if the parent does not have
                # those plate axes.
                m[i] = misc.squeeze_to_dim(m[i], len(shape_parent))

            return m_function
            raise NotImplementedError()

        # Compact the message to a proper shape
        for i in range(len(m)):

            # Empty messages are given as None. We can ignore those.
            if m[i] is not None:

                try:
                    r = self.broadcasting_multiplier(self.plates_multiplier,
                                                     multiplier_parent)
                except:
                    raise ValueError("The plate multipliers are incompatible. "
                                     "This node (%s) has %s and parent[%d] "
                                     "(%s) has %s"
                                     % (self.name,
                                        self.plates_multiplier,
                                        index,
                                        parent.name,
                                        multiplier_parent))

                ndim = len(parent.dims[i])
                # Source and target shapes
                if ndim > 0:
                    dims = misc.broadcasted_shape(np.shape(m[i])[-ndim:],
                                                  parent.dims[i])
                    from_shape = plates_self + dims
                else:
                    from_shape = plates_self
                to_shape = parent.get_shape(i)
                # Add variable axes to the mask
                mask_i = misc.add_trailing_axes(mask, ndim)
                # Apply mask and sum plate axes as necessary (and apply plate
                # multiplier)
                m[i] = r * misc.sum_multiply_to_plates(np.where(mask_i, m[i], 0),
                                                       to_plates=to_shape,
                                                       from_plates=from_shape,
                                                       ndim=0)

        return m

    def _message_from_children(self, u_self=None):
        msg = [np.zeros(shape) for shape in self.dims]
        #msg = [np.array(0.0) for i in range(len(self.dims))]
        isfunction = None
        for (child,index) in self.children:
            m = child._message_to_parent(index, u_parent=u_self)
            if callable(m):
                if isfunction is False:
                    raise NotImplementedError()
                elif isfunction is None:
                    msg = m
                else:
                    def join(m1, m2):
                        return (m1[0] + m2[0], m1[1] + m2[1])
                    msg = lambda x: join(m(x), msg(x))
                    isfunction = True
            else:
                if isfunction is True:
                    raise NotImplementedError()
                else:
                    isfunction = False
                    for i in range(len(self.dims)):
                        if m[i] is not None:
                            # Check broadcasting shapes
                            sh = misc.broadcasted_shape(self.get_shape(i), np.shape(m[i]))
                            try:
                                # Try exploiting broadcasting rules
                                msg[i] += m[i]
                            except ValueError:
                                msg[i] = msg[i] + m[i]

        return msg

    def _message_from_parents(self, exclude=None):
        return [list(parent._message_to_child())
                if ind != exclude else
                None
                for (ind,parent) in enumerate(self.parents)]

    def get_moments(self):
        raise NotImplementedError()

    def delete(self):
        """
        Delete this node and the children
        """
        for (ind, parent) in enumerate(self.parents):
            parent._remove_child(self, ind)
        for (child, _) in self.children:
            child.delete()

    @staticmethod
    def broadcasting_multiplier(plates, *args):
        return misc.broadcasting_multiplier(plates, *args)
        ## """
        ## Compute the plate multiplier for given shapes.

        ## The first shape is compared to all other shapes (using NumPy
        ## broadcasting rules). All the elements which are non-unit in the first
        ## shape but 1 in all other shapes are multiplied together.

        ## This method is used, for instance, for computing a correction factor for
        ## messages to parents: If this node has non-unit plates that are unit
        ## plates in the parent, those plates are summed. However, if the message
        ## has unit axis for that plate, it should be first broadcasted to the
        ## plates of this node and then summed to the plates of the parent. In
        ## order to avoid this broadcasting and summing, it is more efficient to
        ## just multiply by the correct factor. This method computes that
        ## factor. The first argument is the full plate shape of this node (with
        ## respect to the parent). The other arguments are the shape of the message
        ## array and the plates of the parent (with respect to this node).
        ## """

        ## # Check broadcasting of the shapes
        ## for arg in args:
        ##     misc.broadcasted_shape(plates, arg)

        ## # Check that each arg-plates are a subset of plates?
        ## for arg in args:
        ##     if not misc.is_shape_subset(arg, plates):
        ##         raise ValueError("The shapes in args are not a sub-shape of "
        ##                          "plates.")

        ## r = 1
        ## for j in range(-len(plates),0):
        ##     mult = True
        ##     for arg in args:
        ##         # if -j <= len(arg) and arg[j] != 1:
        ##         if not (-j > len(arg) or arg[j] == 1):
        ##             mult = False
        ##     if mult:
        ##         r *= plates[j]
        ## return r

    def move_plates(self, from_plate, to_plate):
        return _MovePlate(self,
                          from_plate,
                          to_plate,
                          name=self.name + ".move_plates")

    def add_plate_axis(self, to_plate):
        return AddPlateAxis(to_plate)(self,
                                      name=self.name+".add_plate_axis")

    def __getitem__(self, index):
        return Slice(self, index,
                     name=(self.name+".__getitem__"))

    def has_plotter(self):
        """
        Return True if the node has a plotter
        """
        return callable(self._plotter)

    def set_plotter(self, plotter):
        self._plotter = plotter

    def plot(self, fig=None, **kwargs):
        """
        Plot the node distribution using the plotter of the node

        Because the distributions are in general very difficult to plot, the
        user must specify some functions which performs the plotting as
        wanted. See, for instance, bayespy.plot.plotting for available plotters,
        that is, functions that perform plotting for a node.
        """
        if fig is None:
            import matplotlib.pyplot as plt
            fig = plt.gcf()
        if callable(self._plotter):
            ax = self._plotter(self, fig=fig, **kwargs)
            fig.suptitle('q(%s)' % self.name)
            return ax
        else:
            raise Exception("No plotter defined, can not plot")


    @staticmethod
    def _compute_message(*arrays, plates_from=(), plates_to=(), ndim=0):
        """
        A general function for computing messages by sum-multiply

        The function computes the product of the input arrays and then sums to
        the requested plates.
        """

        # Check that the plates broadcast properly
        if not misc.is_shape_subset(plates_to, plates_from):
            raise ValueError("plates_to must be broadcastable to plates_from")

        # Compute the explicit shape of the product
        shapes = [np.shape(array) for array in arrays]
        arrays_shape = misc.broadcasted_shape(*shapes)

        # Compute plates and dims that are present
        if ndim == 0:
            arrays_plates = arrays_shape
            dims = ()
        else:
            arrays_plates = arrays_shape[:-ndim]
            dims = arrays_shape[-ndim:]

        # Compute the correction term.  If some of the plates that should be
        # summed are actually broadcasted, one must multiply by the size of the
        # corresponding plate
        r = Node.broadcasting_multiplier(plates_from, arrays_plates, plates_to)

        # For simplicity, make the arrays equal ndim
        arrays = misc.make_equal_ndim(*arrays)

        # Keys for the input plates: (N-1, N-2, ..., 0)
        nplates = len(arrays_plates)
        in_plate_keys = list(range(nplates-1, -1, -1))

        # Keys for the output plates
        out_plate_keys = [key
                          for key in in_plate_keys
                          if key < len(plates_to) and plates_to[-key-1] != 1]

        # Keys for the dims
        dim_keys = list(range(nplates, nplates+ndim))

        # Total input and output keys
        in_keys = len(arrays) * [in_plate_keys + dim_keys]
        out_keys = out_plate_keys + dim_keys

        # Compute the sum-product with correction
        einsum_args = misc.zipper_merge(arrays, in_keys) + [out_keys]
        y = r * np.einsum(*einsum_args)

        # Reshape the result and apply correction
        nplates_result = min(len(plates_to), len(arrays_plates))
        if nplates_result == 0:
            plates_result = []
        else:
            plates_result = [min(plates_to[ind], arrays_plates[ind])
                             for ind in range(-nplates_result, 0)]

        y = np.reshape(y, plates_result + list(dims))

        return y


from .deterministic import Deterministic


def slicelen(s, length=None):
    if length is not None:
        s = slice(*(s.indices(length)))
    return max(0, misc.ceildiv(s.stop - s.start, s.step))

class Slice(Deterministic):

    """
    Basic slicing for plates.

    Slicing occurs when index is a slice object (constructed by start:stop:step
    notation inside of brackets), an integer, or a tuple of slice objects and
    integers.

    Currently, accept slices, newaxis, ellipsis and integers. For instance, does
    not accept lists/tuples to pick multiple indices of the same axis.

    Ellipsis expand to the number of : objects needed to make a selection tuple
    of the same length as x.ndim. Only the first ellipsis is expanded, any
    others are interpreted as :.

    Similar to:
    http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#basic-slicing
    """


    def __init__(self, X, slices, **kwargs):

        self._moments = X._moments
        self._parent_moments = (X._moments,)

        # Force a list
        if not isinstance(slices, tuple):
            slices = [slices]
        else:
            slices = list(slices)

        #
        # Expand Ellipsis
        #

        # Compute the number of required axes and how Ellipsis is expanded
        num_axis = 0
        ellipsis_index = None
        for (k, s) in enumerate(slices):

            if misc.is_scalar_integer(s) or isinstance(s, slice):
                num_axis += 1

            elif s is None:
                pass

            elif s is Ellipsis:
                # Index is an ellipsis, e.g., [...]

                if ellipsis_index is None:
                    # Expand ...
                    ellipsis_index = k
                else:
                    # Interpret ... as :
                    num_axis += 1
                    slices[k] = slice(None)

            else:
                raise TypeError("Invalid argument type: {0}".format(s.__class__))

        if num_axis > len(X.plates):
            raise IndexError("Too many indices")

        # The number of plates that were not given explicit slicing (either
        # Ellipsis was used or the number of slices was smaller than the number
        # of plate axes)
        expand_len = len(X.plates) - num_axis

        if ellipsis_index is not None:
            # Replace Ellipsis with correct number of :
            k = ellipsis_index
            del slices[k]
            slices = slices[:k] + [slice(None)] * expand_len + slices[k:]
        else:
            # Add trailing : so that each plate has explicit slicing
            slices = slices + [slice(None)] * expand_len

        #
        # Preprocess indexing:
        # - integer indices to non-negative values
        # - slice start/stop values to non-negative
        # - slice start/stop values based on the size of the plate
        #

        # Index for parent plates
        j = 0

        for (k, s) in enumerate(slices):

            if misc.is_scalar_integer(s):
                # Index is an integer, e.g., [3]

                if s < 0:
                    # Handle negative index
                    s += X.plates[j]
                if s < 0 or s >= X.plates[j]:
                    raise IndexError("Index out of range")
                # Store the preprocessed integer index
                slices[k] = s
                j += 1

            elif isinstance(s, slice):
                # Index is a slice, e.g., [2:6]

                # Normalize the slice
                s = slice(*(s.indices(X.plates[j])))
                if slicelen(s) <= 0:
                    raise IndexError("Slicing leads to empty plates")
                slices[k] = s
                j += 1

        self.slices = slices

        super().__init__(X,
                         dims=X.dims,
                         **kwargs)

    def _plates_to_parent(self, index):
        return self.parents[index].plates

    def _plates_from_parent(self, index):

        plates = list(self.parents[index].plates)

        # Compute the plates. Note that Ellipsis has already been preprocessed
        # to a proper number of :
        k = 0
        for s in self.slices:
            # Then, each case separately: slice, newaxis, integer

            if isinstance(s, slice):
                # Slice, e.g., [2:5]
                N = slicelen(s)
                if N <= 0:
                    raise IndexError("Slicing leads to empty plates")
                plates[k] = N
                k += 1

            elif s is None:
                # [np.newaxis]
                plates = plates[:k] + [1] + plates[k:]
                k += 1

            elif misc.is_scalar_integer(s):
                # Integer, e.g., [3]
                del plates[k]
            else:
                raise RuntimeError("BUG: Unknown index type. Should capture earlier.")

        return tuple(plates)

    @staticmethod
    def __reverse_indexing(slices, m_child, plates, dims):
        """
        A helpful function for performing reverse indexing/slicing
        """

        j = -1 # plate index for parent
        i = -1 # plate index for child
        child_slices = ()
        parent_slices = ()
        msg_plates = ()

        # Compute plate axes in the message from children
        ndim = len(dims)
        if ndim > 0:
            m_plates = np.shape(m_child)[:-ndim]
        else:
            m_plates = np.shape(m_child)

        for s in reversed(slices):

            if misc.is_scalar_integer(s):
                # Case: integer
                parent_slices = (s,) + parent_slices
                msg_plates = (plates[j],) + msg_plates
                j -= 1
            elif s is None:
                # Case: newaxis
                if -i <= len(m_plates):
                    child_slices = (0,) + child_slices
                i -= 1
            elif isinstance(s, slice):
                # Case: slice
                if -i <= len(m_plates):
                    child_slices = (slice(None),) + child_slices
                parent_slices = (s,) + parent_slices
                if ((-i > len(m_plates) or m_plates[i] == 1)
                    and slicelen(s) == plates[j]):
                    # Broadcasting can be applied. The message does not need
                    # to be explicitly shaped to the full size
                    msg_plates = (1,) + msg_plates
                else:
                    # No broadcasting. Must explicitly form the full size
                    # axis
                    msg_plates = (plates[j],) + msg_plates
                j -= 1
                i -= 1
            else:
                raise RuntimeError("BUG: Unknown index type. Should capture earlier.")

        # Set the elements of the message
        m_parent = np.zeros(msg_plates + dims)
        if np.ndim(m_parent) == 0 and np.ndim(m_child) == 0:
            m_parent = m_child
        elif np.ndim(m_parent) == 0:
            m_parent = m_child[child_slices]
        elif np.ndim(m_child) == 0:
            m_parent[parent_slices] = m_child
        else:
            m_parent[parent_slices] = m_child[child_slices]

        return m_parent


    def _compute_weights_to_parent(self, index, weights):
        """
        Compute the mask to the parent node.
        """
        if index != 0:
            raise ValueError("Invalid index")
        parent = self.parents[0]

        return self.__reverse_indexing(self.slices,
                                       weights,
                                       parent.plates,
                                       ())


    def _compute_message_to_parent(self, index, m, u):
        """
        Compute the message to a parent node.
        """

        if index != 0:
            raise ValueError("Invalid index")
        parent = self.parents[0]

        # Apply reverse indexing for the message arrays
        msg = [self.__reverse_indexing(self.slices,
                                       m_child,
                                       parent.plates,
                                       dims)
               for (m_child, dims) in zip(m, parent.dims)]

        return msg

    def _compute_moments(self, u):
        """
        Get the moments with an added plate axis.
        """

        # Process each moment
        for n in range(len(u)):

            # Compute the effective plates in the message/moment
            ndim = len(self.dims[n])
            if ndim > 0:
                shape = np.shape(u[n])[:-ndim]
            else:
                shape = np.shape(u[n])

            # Construct a list of slice objects
            u_slices = []

            # Index for the shape
            j = -len(self.parents[0].plates)

            for (k, s) in enumerate(self.slices):

                if s is None:
                    # [np.newaxis]
                    if -j < len(shape):
                        # Only add newaxis if there are some axes before
                        # this. It does not make any difference if you added
                        # leading unit axes
                        u_slices.append(s)

                else:
                    # slice or integer index

                    if -j <= len(shape):
                        # The moment has this axis, so it is not broadcasting it
                        if shape[j] != 1:
                            # Use the slice as it is
                            u_slices.append(s)
                        elif isinstance(s, slice):
                            # Slice.
                            # The moment is using broadcasting, just pick the
                            # first element but use slice in order to keep the
                            # axis
                            u_slices.append(slice(0,1,1))
                        else:
                            # Integer.
                            # The moment is using broadcasting, just pick the
                            # first element
                            u_slices.append(0)

                    j += 1

            # Slice the message/moment
            u[n] = u[n][tuple(u_slices)]

        return u

def AddPlateAxis(to_plate):

    if to_plate >= 0:
        raise Exception("Give negative value for axis index to_plate.")

    class _AddPlateAxis(Deterministic):

        def __init__(self, X, **kwargs):

            nonlocal to_plate

            N = len(X.plates) + 1

            # Check the parameters
            if to_plate >= 0 or to_plate < -N:
                raise ValueError("Invalid plate position to add.")

            # Use positive indexing only
            ## if to_plate < 0:
            ##     to_plate += N
            # Use negative indexing only
            if to_plate >= 0:
                to_plate -= N
                #self.to_plate = to_plate

            super().__init__(X,
                             dims=X.dims,
                             **kwargs)

        def _plates_to_parent(self, index):
            plates = list(self.plates)
            plates.pop(to_plate)
            return tuple(plates)
        #return self.plates[:to_plate] + self.plates[(to_plate+1):]

        def _plates_from_parent(self, index):
            plates = list(self.parents[index].plates)
            plates.insert(len(plates)-to_plate+1, 1)
            return tuple(plates)


        def _compute_weights_to_parent(self, index, weights):
            # Remove the added mask plate
            if abs(to_plate) <= np.ndim(weights):
                sh_weighs = list(np.shape(weights))
                sh_weights.pop(to_plate)
                weights = np.reshape(weights, sh_weights)
            return weights


        def _compute_message_to_parent(self, index, m, *u_parents):
            """
            Compute the message to a parent node.
            """

            # Remove the added message plate
            for i in range(len(m)):
                # Remove the axis
                if np.ndim(m[i]) >= abs(to_plate) + len(self.dims[i]):
                    axis = to_plate - len(self.dims[i])
                    sh_m = list(np.shape(m[i]))
                    sh_m.pop(axis)
                    m[i] = np.reshape(m[i], sh_m)

            return m

        def _compute_moments(self, u):
            """
            Get the moments with an added plate axis.
            """

            # Get parents' moments
            #u = self.parents[0].message_to_child()

            # Move a plate axis
            u = list(u)
            for i in range(len(u)):
                # Make sure the moments have all the axes
                #diff = len(self.plates) + len(self.dims[i]) - np.ndim(u[i]) - 1
                #u[i] = misc.add_leading_axes(u[i], diff)

                # The location of the new axis/plate:
                axis = np.ndim(u[i]) - abs(to_plate) - len(self.dims[i]) + 1
                if axis > 0:
                    # Add one axes to the correct position
                    sh_u = list(np.shape(u[i]))
                    sh_u.insert(axis, 1)
                    u[i] = np.reshape(u[i], sh_u)

            return u

    return _AddPlateAxis


class NodeConstantScalar(Node):
    @staticmethod
    def compute_fixed_u_and_f(x):
        """ Compute u(x) and f(x) for given x. """
        return ([x], 0)

    def __init__(self, a, **kwargs):
        self.u = [a]
        super().__init__(self,
                         plates=np.shape(a),
                         dims=[()],
                         **kwargs)

    def start_optimization(self):
        # FIXME: Set the plate sizes appropriately!!
        x0 = self.u[0]
        #self.gradient = np.zeros(np.shape(x0))
        def transform(x):
            # E.g., for positive scalars you could have exp here.
            self.gradient = np.zeros(np.shape(x0))
            self.u[0] = x
        def gradient():
            # This would need to apply the gradient of the
            # transformation to the computed gradient
            return self.gradient

        return (x0, transform, gradient)

    def add_to_gradient(self, d):
        self.gradient += d

    def message_to_child(self, gradient=False):
        if gradient:
            return (self.u, [ [np.ones(np.shape(self.u[0])),
                               #self.gradient] ])
                               self.add_to_gradient] ])
        else:
            return self.u


    def stop_optimization(self):
        #raise Exception("Not implemented for " + str(self.__class__))
        pass