"""
Class definitions for all of the symmetric cones (and their superclass,
SymmetricCone) supported by CVXOPT.
"""

from cvxopt import matrix
from matrices import norm

class SymmetricCone:
    """
    An instance of a symmetric (self-dual and homogeneous) cone.

    There are three types of symmetric cones supported by CVXOPT:

      * The nonnegative orthant in the real n-space.

      * The Lorentz "ice cream" cone, or the second-order cone.

      * The cone of symmetric positive-semidefinite matrices.

    This class is intended to encompass them all.
    """
    def __init__(self, dimension):
        """
        A generic constructor for symmetric cones.

        When constructing a single symmetric cone (i.e. not a cartesian
        product of them), the only information that we need is its
        dimension. We take that dimension as a parameter, and store it
        for later.
        """
        if dimension <= 0:
            raise ValueError('cones must have dimension greater than zero')

        self._dimension = dimension

    def __contains__(self, point):
        """
        Return whether or not ``point`` belongs to this cone.
        """
        raise NotImplementedError

    def contains_strict(self, point):
        """
        Return whether or not ``point`` belongs to the interior
        of this cone.
        """
        raise NotImplementedError

    def dimension(self):
        """
        Return the dimension of this symmetric cone.

        The dimension of this symmetric cone is recorded during its
        creation. This method simply returns the recorded value, and
        should not need to be overridden in subclasses. We prefer to do
        any special computation in ``__init__()`` and record the result
        in ``self._dimension``.
        """
        return self._dimension


class NonnegativeOrthant(SymmetricCone):
    """
    The nonnegative orthant in ``n`` dimensions.

    EXAMPLES:

        >>> K = NonnegativeOrthant(3)
        >>> print(K)
        Nonnegative orthant in the real 3-space

    """
    def __str__(self):
        """
        Output a human-readable description of myself.
        """
        tpl = 'Nonnegative orthant in the real {:d}-space'
        return tpl.format(self.dimension())

    def __contains__(self, point):
        """
        Return whether or not ``point`` belongs to this cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = NonnegativeOrthant(3)
            >>> matrix([1,2,3]) in K
            True

            >>> K = NonnegativeOrthant(3)
            >>> matrix([1,-0.1,3]) in K
            False

            >>> K = NonnegativeOrthant(3)
            >>> [1,2,3] in K
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = NonnegativeOrthant(3)
            >>> matrix([1,2]) in K
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        return all([x >= 0 for x in point])


    def contains_strict(self, point):
        """
        Return whether or not ``point`` belongs to the interior of this
        cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = NonnegativeOrthant(3)
            >>> K.contains_strict(matrix([1,2,3]))
            True

            >>> K = NonnegativeOrthant(3)
            >>> K.contains_strict(matrix([1,0,1]))
            False

            >>> K = NonnegativeOrthant(3)
            >>> K.contains_strict(matrix([1,-0.1,3]))
            False

            >>> K = NonnegativeOrthant(3)
            >>> K.contains_strict([1,2,3])
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = NonnegativeOrthant(3)
            >>> K.contains_strict(matrix([1,2]))
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        return all([x > 0 for x in point])



class IceCream(SymmetricCone):
    """
    The nonnegative orthant in ``n`` dimensions.

    EXAMPLES:

        >>> K = IceCream(3)
        >>> print(K)
        Lorentz "ice cream" cone in the real 3-space

    """
    def __str__(self):
        """
        Output a human-readable description of myself.
        """
        tpl = 'Lorentz "ice cream" cone in the real {:d}-space'
        return tpl.format(self.dimension())


    def __contains__(self, point):
        """
        Return whether or not ``point`` belongs to this cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = IceCream(3)
            >>> matrix([1,0.5,0.5]) in K
            True

            >>> K = IceCream(3)
            >>> matrix([1,0,1]) in K
            True

            >>> K = IceCream(3)
            >>> matrix([1,1,1]) in K
            False

            >>> K = IceCream(3)
            >>> [1,2,3] in K
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = IceCream(3)
            >>> matrix([1,2]) in K
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        height = point[0]
        if self.dimension() == 1:
            # In one dimension, the ice cream cone is the nonnegative
            # orthant.
            return height >= 0
        else:
            radius = point[1:]
            return height >= norm(radius)


    def contains_strict(self, point):
        """
        Return whether or not ``point`` belongs to the interior
        of this cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = IceCream(3)
            >>> K.contains_strict(matrix([1,0.5,0.5]))
            True

            >>> K = IceCream(3)
            >>> K.contains_strict(matrix([1,0,1]))
            False

            >>> K = IceCream(3)
            >>> K.contains_strict(matrix([1,1,1]))
            False

            >>> K = IceCream(3)
            >>> K.contains_strict([1,2,3])
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = IceCream(3)
            >>> K.contains_strict(matrix([1,2]))
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        height = point[0]
        if self.dimension() == 1:
            # In one dimension, the ice cream cone is the nonnegative
            # orthant.
            return height > 0
        else:
            radius = point[1:]
            return height > norm(radius)


class SymmetricPSD(SymmetricCone):
    """
    The nonnegative orthant in ``n`` dimensions.

    EXAMPLES:

        >>> K = SymmetricPSD(3)
        >>> print(K)
        Cone of symmetric positive-semidefinite matrices on the real 3-space

    """
    def __str__(self):
        """
        Output a human-readable description of myself.
        """
        tpl = 'Cone of symmetric positive-semidefinite matrices ' \
              'on the real {:d}-space'
        return tpl.format(self.dimension())


class CartesianProduct(SymmetricCone):
    """
    A cartesian product of symmetric cones, which is itself a symmetric
    cone.

    EXAMPLES:

        >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(2))
        >>> print(K)
        Cartesian product of dimension 5 with 2 factors:
          * Nonnegative orthant in the real 3-space
          * Lorentz "ice cream" cone in the real 2-space

    """
    def __init__(self, *factors):
        my_dimension = sum([f.dimension() for f in factors])
        super().__init__(my_dimension)
        self._factors = factors

    def __str__(self):
        """
        Output a human-readable description of myself.
        """
        tpl = 'Cartesian product of dimension {:d} with {:d} factors:'
        tpl += '\n  * {!s}' * len(self.factors())
        format_args = [self.dimension(), len(self.factors())]
        format_args += list(self.factors())
        return tpl.format(*format_args)

    def __contains__(self, point):
        """
        Return whether or not ``point`` belongs to this cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> matrix([1,2,3,1,0.5,0.5]) in K
            True

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> matrix([0,0,0,1,0,1]) in K
            True

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> matrix([1,1,1,1,1,1]) in K
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> matrix([1,-1,1,1,0,1]) in K
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> [1,2,3,4,5,6] in K
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> matrix([1,2]) in K
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        for factor in self.factors():
            # Split off the components of ``point`` corresponding to
            # ``factor``.
            factor_part = point[0:factor.dimension()]
            if not factor_part in factor:
                return False
            point = point[factor.dimension():]

        return True


    def contains_strict(self, point):
        """
        Return whether or not ``point`` belongs to the interior
        of this cone.

        INPUT:

        An instance of the ``cvxopt.base.matrix`` class having
        dimensions ``(n,1)`` where ``n`` is the dimension of this cone.

        EXAMPLES:

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([1,2,3,1,0.5,0.5]))
            True

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([1,2,3,1,0,1]))
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([0,1,1,1,0.5,0.5]))
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([1,1,1,1,1,1]))
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([1,-1,1,1,0,1]))
            False

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict([1,2,3,4,5,6])
            Traceback (most recent call last):
            ...
            TypeError: the given point is not a cvxopt.base.matrix

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(3))
            >>> K.contains_strict(matrix([1,2]))
            Traceback (most recent call last):
            ...
            TypeError: the given point has the wrong dimensions

        """
        if not isinstance(point, matrix):
            raise TypeError('the given point is not a cvxopt.base.matrix')
        if not point.size == (self.dimension(), 1):
            raise TypeError('the given point has the wrong dimensions')

        for factor in self.factors():
            # Split off the components of ``point`` corresponding to
            # ``factor``.
            factor_part = point[0:factor.dimension()]
            if not factor.contains_strict(factor_part):
                return False
            point = point[factor.dimension():]

        return True


    def factors(self):
        """
        Return a tuple containing the factors (in order) of this
        cartesian product.

        EXAMPLES:

            >>> K = CartesianProduct(NonnegativeOrthant(3), IceCream(2))
            >>> len(K.factors())
            2

        """
        return self._factors

    def cvxopt_dims(self):
        """
        Return a dictionary of dimensions corresponding to the factors
        of this cartesian product. The format of this dictionary is
        described in the CVXOPT user's guide:

          http://cvxopt.org/userguide/coneprog.html#linear-cone-programs

        EXAMPLES:

            >>> K = CartesianProduct(NonnegativeOrthant(3),
            ...                      IceCream(2),
            ...                      IceCream(3))
            >>> d = K.cvxopt_dims()
            >>> (d['l'], d['q'], d['s'])
            (3, [2, 3], [])

        """
        dims = {'l':0, 'q':[], 's':[]}
        dims['l'] += sum([K.dimension()
                          for K in self.factors()
                          if isinstance(K, NonnegativeOrthant)])
        dims['q'] = [K.dimension()
                     for K in self.factors()
                     if isinstance(K, IceCream)]
        dims['s'] = [K.dimension()
                     for K in self.factors()
                     if isinstance(K, SymmetricPSD)]
        return dims