"""
Copyright 2013 Steven Diamond

This file is part of CVXPY.

CVXPY is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

CVXPY is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with CVXPY.  If not, see <http://www.gnu.org/licenses/>.
"""


import cvxpy.utilities as u
from cvxpy.atoms.atom import Atom
from cvxpy.atoms.affine.index import index
import numpy as np
import numbers


from ..utilities.power_tools import fracify, decompose, approx_error, lower_bound, over_bound, prettydict, gm, gm_constrs
import cvxpy.lin_ops.lin_utils as lu


class geo_mean(Atom):
    """ The (weighted) geometric mean of vector ``x``, with optional powers given by ``p``:

    .. math::

        \\left(x_1^{p_1} \cdots x_n^{p_n} \\right)^{\\frac{1}{\mathbf{1}^Tp}}

    The powers ``p`` can be a ``list``, ``tuple``, or ``numpy.array`` of nonnegative
    ``int``, ``float``, or ``Fraction`` objects with nonzero sum.

    If not specified, ``p`` defaults to a vector of all ones, giving the
    **unweighted** geometric mean

    .. math::

        x_1^{1/n} \cdots x_n^{1/n}.

    The geometric mean includes an implicit constraint that :math:`x_i \geq 0`
    whenever :math:`p_i > 0`. If :math:`p_i = 0`, :math:`x_i` will be unconstrained.

    The only exception to this rule occurs when
    ``p`` has exactly one nonzero element, say, ``p_i``, in which case
    ``geo_mean(x, p)`` is equivalent to ``x_i`` (without the nonnegativity constraint).
    A specific case of this is when :math:`x \in \mathbf{R}^1`.


    .. note::

        Generally, ``p`` cannot be represented exactly, so a rational,
        i.e., fractional, **approximation** must be made.

        Internally, ``geo_mean`` immediately computes an approximate normalized
        weight vector :math:`w \\approx p/\mathbf{1}^Tp`
        and the ``geo_mean`` atom is represented as

        .. math::

            x_1^{w_1} \cdots x_n^{w_n},

        where the elements of ``w`` are ``Fraction`` objects that sum to exactly 1.

        The maximum denominator used in forming the rational approximation
        is given by ``max_denom``, which defaults to 1024, but can be adjusted
        to modify the accuracy of the approximations.

        The approximating ``w`` and the approximation error can be
        found through the attributes ``geo_mean.w`` and ``geo_mean.approx_error``.


    Examples
    --------

    The weights ``w`` can be seen from the string representation of the ``geo_mean`` object, or through
    the ``w`` attribute.

    >>> from cvxpy import Variable, geo_mean, Problem, Maximize
    >>> x = Variable(3, name='x')
    >>> print(geo_mean(x))
    geo_mean(x, (1/3, 1/3, 1/3))
    >>> g = geo_mean(x, [1, 2, 1])
    >>> g.w
    (Fraction(1, 4), Fraction(1, 2), Fraction(1, 4))

    Floating point numbers with few decimal places can sometimes be represented exactly. The approximation
    error between ``w`` and ``p/sum(p)`` is given by the ``approx_error`` attribute.

    >>> import numpy as np
    >>> x = Variable(4, name='x')
    >>> p = np.array([.12, .34, .56, .78])
    >>> g = geo_mean(x, p)
    >>> g.w
    (Fraction(1, 15), Fraction(17, 90), Fraction(14, 45), Fraction(13, 30))
    >>> g.approx_error
    0.0

    In general, the approximation is not exact.

    >>> p = [.123, .456, .789, .001]
    >>> g = geo_mean(x, p)
    >>> g.w
    (Fraction(23, 256), Fraction(341, 1024), Fraction(295, 512), Fraction(1, 1024))
    >>> 1e-4 <= g.approx_error <= 1e-3
    True

    The weight vector ``p`` can contain combinations of ``int``, ``float``, and ``Fraction`` objects.

    >>> from fractions import Fraction
    >>> x = Variable(4, name='x')
    >>> g = geo_mean(x, [.1, Fraction(1,3), 0, 2])
    >>> print(g)
    geo_mean(x, (3/73, 10/73, 0, 60/73))
    >>> g.approx_error <= 1e-10
    True

    Sequences of ``Fraction`` and ``int`` powers can often be represented **exactly**.

    >>> p = [Fraction(1,17), Fraction(4,9), Fraction(1,3), Fraction(25,153)]
    >>> x = Variable(4, name='x')
    >>> print(geo_mean(x, p))
    geo_mean(x, (1/17, 4/9, 1/3, 25/153))

    Terms with a zero power will not have an implicit nonnegativity constraint.

    >>> p = [1, 0, 1]
    >>> x = Variable(3, name='x')
    >>> obj = Maximize(geo_mean(x,p))
    >>> constr = [sum(x) <= 1, -1 <= x, x <= 1]
    >>> val = Problem(obj, constr).solve()
    >>> x = np.array(x.value).flatten()
    >>> print(x)
    [ 1. -1.  1.]


    Parameters
    ----------
    x : cvxpy.Variable
        A column or row vector whose elements we will take the geometric mean of.

    p : Sequence (list, tuple, numpy.array, ...) of ``int``, ``float``, or ``Fraction`` objects
        A vector of weights for the weighted geometric mean

        When ``p`` is a sequence of ``int`` and/or ``Fraction`` objects, ``w`` can often be an **exact** representation
        of the weights. An exact representation is sometimes possible when ``p`` has ``float`` elements with only a few
        decimal places.

    max_denom : int
        The maximum denominator to use in approximating ``p/sum(p)`` with ``geo_mean.w``. If ``w`` is not an exact
        representation, increasing ``max_denom`` **may** offer a more accurate representation, at the cost of requiring
        more convex inequalities to represent the geometric mean.


    Attributes
    ----------
    w : tuple of ``Fractions``
        A rational approximation of ``p/sum(p)``.
    approx_error : float
        The error in approximating ``p/sum(p)`` with ``w``, given by :math:`\|p/\mathbf{1}^T p - w \|_\infty`
    """

    def __init__(self, x, p=None, max_denom=1024):
        """ Implementation details of geo_mean.

        Attributes
        ----------

        w_dyad : tuple of ``Fractions`` whose denominators are all a power of two
            The dyadic completion of ``w``, which is used internally to form the inequalities representing the
            geometric mean.

        tree : ``dict``
            keyed by dyadic tuples, whose values are Sequences of children.
            The children are also dyadic tuples.
            This represents the graph that needs to be formed to represent the weighted geometric mean.

        cone_lb : int
            A known lower bound (which is not always tight) on the number of cones needed to represent this
            geometric mean.

        cone_num_over : int
            The number of cones beyond the lower bound that this geometric mean used.
            If 0, we know that it used the minimum possible number of cones.
            Since cone_lb is not always tight, it may be using the minimum number of cones even if
            cone_num_over is not 0.

        cone_num : int
            The number of second order cones used to form this geometric mean

        """
        super(geo_mean, self).__init__(x)

        x = self.args[0]
        if x.size[0] == 1:
            n = x.size[1]
        elif x.size[1] == 1:
            n = x.size[0]
        else:
            raise ValueError('x must be a row or column vector.')

        if p is None:
            p = [1]*n

        if len(p) != n:
            raise ValueError('x and p must have the same number of elements.')

        if any(v < 0 for v in p) or sum(p) <= 0:
            raise ValueError('powers must be nonnegative and not all zero.')

        self.w, self.w_dyad = fracify(p, max_denom)
        self.approx_error = approx_error(p, self.w)

        self.tree = decompose(self.w_dyad)

        # known lower bound on number of cones needed to represent w_dyad
        self.cone_lb = lower_bound(self.w_dyad)

        # number of cones used past known lower bound
        self.cone_num_over = over_bound(self.w_dyad, self.tree)

        # number of cones used
        self.cone_num = self.cone_lb + self.cone_num_over

    # Returns the (weighted) geometric mean of the elements of x.
    @Atom.numpy_numeric
    def numeric(self, values):
        values = np.array(values[0]).flatten()
        val = 1.0
        for x, p in zip(values, self.w):
            val *= x**p
        return val

    def name(self):
        return "%s(%s, (%s))" % (self.__class__.__name__,
                                 self.args[0].name(),
                                 ', '.join(str(v) for v in self.w))

    def pretty_tree(self):
        print(prettydict(self.tree))

    def shape_from_args(self):
        return u.Shape(1, 1)

    def sign_from_args(self):
        return u.Sign.POSITIVE

    def func_curvature(self):
        return u.Curvature.CONCAVE

    def monotonicity(self):
        return [u.monotonicity.INCREASING]

    def validate_arguments(self):
        # since correctly validating arguments with this function is tricky,
        # we do it in __init__ instead.
        pass

    def get_data(self):
        return self.w, self.w_dyad, self.tree

    @staticmethod
    def graph_implementation(arg_objs, size, data=None):
        """Reduces the atom to an affine expression and list of constraints.

        Parameters
        ----------
        arg_objs : list
            LinExpr for each argument.
        size : tuple
            The size of the resulting expression.
        data :
            Additional data required by the atom.

        Returns
        -------
        tuple
            (LinOp for objective, list of constraints)
        """
        w, w_dyad, tree = data
        t = lu.create_var((1, 1))
        x_list = [index.get_index(arg_objs[0], [], i, 0) for i in range(len(w))]

        #todo: catch cases where we have (0, 0, 1)?
        #todo: what about curvature case (should be affine) in trivial case of (0, 0 , 1), should this behavior match with what we do in power?

        return t, gm_constrs(t, x_list, w)