"""
Symmetric linear games and their solutions.

This module contains the main :class:`SymmetricLinearGame` class that
knows how to solve a linear game.
"""

# These few are used only for tests.
from math import sqrt
from random import randint, uniform
from unittest import TestCase

# These are mostly actually needed.
from cvxopt import matrix, printing, solvers
from cones import CartesianProduct, IceCream, NonnegativeOrthant
from errors import GameUnsolvableException
from matrices import append_col, append_row, identity, inner_product
import options

printing.options['dformat'] = options.FLOAT_FORMAT
solvers.options['show_progress'] = options.VERBOSE


class Solution:
    """
    A representation of the solution of a linear game. It should contain
    the value of the game, and both players' strategies.

    Examples
    --------

        >>> print(Solution(10, matrix([1,2]), matrix([3,4])))
        Game value: 10.0000000
        Player 1 optimal:
          [ 1]
          [ 2]
        Player 2 optimal:
          [ 3]
          [ 4]

    """
    def __init__(self, game_value, p1_optimal, p2_optimal):
        """
        Create a new Solution object from a game value and two optimal
        strategies for the players.
        """
        self._game_value = game_value
        self._player1_optimal = p1_optimal
        self._player2_optimal = p2_optimal

    def __str__(self):
        """
        Return a string describing the solution of a linear game.

        The three data that are described are,

          * The value of the game.
          * The optimal strategy of player one.
          * The optimal strategy of player two.

        The two optimal strategy vectors are indented by two spaces.
        """
        tpl = 'Game value: {:.7f}\n' \
              'Player 1 optimal:{:s}\n' \
              'Player 2 optimal:{:s}'

        p1_str = '\n{!s}'.format(self.player1_optimal())
        p1_str = '\n  '.join(p1_str.splitlines())
        p2_str = '\n{!s}'.format(self.player2_optimal())
        p2_str = '\n  '.join(p2_str.splitlines())

        return tpl.format(self.game_value(), p1_str, p2_str)


    def game_value(self):
        """
        Return the game value for this solution.

        Examples
        --------

            >>> s = Solution(10, matrix([1,2]), matrix([3,4]))
            >>> s.game_value()
            10

        """
        return self._game_value


    def player1_optimal(self):
        """
        Return player one's optimal strategy in this solution.

        Examples
        --------

            >>> s = Solution(10, matrix([1,2]), matrix([3,4]))
            >>> print(s.player1_optimal())
            [ 1]
            [ 2]
            <BLANKLINE>

        """
        return self._player1_optimal


    def player2_optimal(self):
        """
        Return player two's optimal strategy in this solution.

        Examples
        --------

            >>> s = Solution(10, matrix([1,2]), matrix([3,4]))
            >>> print(s.player2_optimal())
            [ 3]
            [ 4]
            <BLANKLINE>

        """
        return self._player2_optimal


class SymmetricLinearGame:
    """
    A representation of a symmetric linear game.

    The data for a linear game are,

      * A "payoff" operator ``L``.
      * A cone ``K``.
      * A point ``e`` in the interior of ``K``.
      * A point ``f`` in the interior of the dual of ``K``.

    In a symmetric game, the cone ``K`` is be self-dual. We therefore
    name the two interior points ``e1`` and ``e2`` to indicate that
    they come from the same cone but are "chosen" by players one and
    two respectively.

    The ambient space is assumed to be the span of ``K``.

    Examples
    --------

        >>> from cones import NonnegativeOrthant
        >>> K = NonnegativeOrthant(3)
        >>> L = [[1,-5,-15],[-1,2,-3],[-12,-15,1]]
        >>> e1 = [1,1,1]
        >>> e2 = [1,2,3]
        >>> SLG = SymmetricLinearGame(L, K, e1, e2)
        >>> print(SLG)
        The linear game (L, K, e1, e2) where
          L = [  1  -5 -15]
              [ -1   2  -3]
              [-12 -15   1],
          K = Nonnegative orthant in the real 3-space,
          e1 = [ 1]
               [ 1]
               [ 1],
          e2 = [ 1]
               [ 2]
               [ 3].


    Lists can (and probably should) be used for every argument::

        >>> from cones import NonnegativeOrthant
        >>> K = NonnegativeOrthant(2)
        >>> L = [[1,0],[0,1]]
        >>> e1 = [1,1]
        >>> e2 = [1,1]
        >>> G = SymmetricLinearGame(L, K, e1, e2)
        >>> print(G)
        The linear game (L, K, e1, e2) where
          L = [ 1  0]
              [ 0  1],
          K = Nonnegative orthant in the real 2-space,
          e1 = [ 1]
               [ 1],
          e2 = [ 1]
               [ 1].

    The points ``e1`` and ``e2`` can also be passed as some other
    enumerable type (of the correct length) without much harm, since
    there is no row/column ambiguity::

        >>> import cvxopt
        >>> import numpy
        >>> from cones import NonnegativeOrthant
        >>> K = NonnegativeOrthant(2)
        >>> L = [[1,0],[0,1]]
        >>> e1 = cvxopt.matrix([1,1])
        >>> e2 = numpy.matrix([1,1])
        >>> G = SymmetricLinearGame(L, K, e1, e2)
        >>> print(G)
        The linear game (L, K, e1, e2) where
          L = [ 1  0]
              [ 0  1],
          K = Nonnegative orthant in the real 2-space,
          e1 = [ 1]
               [ 1],
          e2 = [ 1]
               [ 1].

    However, ``L`` will always be intepreted as a list of rows, even
    if it is passed as a :class:`cvxopt.base.matrix` which is
    otherwise indexed by columns::

        >>> import cvxopt
        >>> from cones import NonnegativeOrthant
        >>> K = NonnegativeOrthant(2)
        >>> L = [[1,2],[3,4]]
        >>> e1 = [1,1]
        >>> e2 = e1
        >>> G = SymmetricLinearGame(L, K, e1, e2)
        >>> print(G)
        The linear game (L, K, e1, e2) where
          L = [ 1  2]
              [ 3  4],
          K = Nonnegative orthant in the real 2-space,
          e1 = [ 1]
               [ 1],
          e2 = [ 1]
               [ 1].
        >>> L = cvxopt.matrix(L)
        >>> print(L)
        [ 1  3]
        [ 2  4]
        <BLANKLINE>
        >>> G = SymmetricLinearGame(L, K, e1, e2)
        >>> print(G)
        The linear game (L, K, e1, e2) where
          L = [ 1  2]
              [ 3  4],
          K = Nonnegative orthant in the real 2-space,
          e1 = [ 1]
               [ 1],
          e2 = [ 1]
               [ 1].

    """
    def __init__(self, L, K, e1, e2):
        """
        Create a new SymmetricLinearGame object.

        INPUT:

          - ``L`` -- an square matrix represented as a list of lists
             of real numbers. ``L`` itself is interpreted as a list of
             ROWS, which agrees with (for example) SageMath and NumPy,
             but not with CVXOPT (whose matrix constructor accepts a
             list of columns).

          - ``K`` -- a SymmetricCone instance.

          - ``e1`` -- the interior point of ``K`` belonging to player one;
            it can be of any enumerable type having the correct length.

          - ``e2`` -- the interior point of ``K`` belonging to player two;
            it can be of any enumerable type having the correct length.

        """
        self._K = K
        self._e1 = matrix(e1, (K.dimension(), 1))
        self._e2 = matrix(e2, (K.dimension(), 1))

        # Our input ``L`` is indexed by rows but CVXOPT matrices are
        # indexed by columns, so we need to transpose the input before
        # feeding it to CVXOPT.
        self._L = matrix(L, (K.dimension(), K.dimension())).trans()

        if not K.contains_strict(self._e1):
            raise ValueError('the point e1 must lie in the interior of K')

        if not K.contains_strict(self._e2):
            raise ValueError('the point e2 must lie in the interior of K')

    def __str__(self):
        """
        Return a string representation of this game.
        """
        tpl = 'The linear game (L, K, e1, e2) where\n' \
              '  L = {:s},\n' \
              '  K = {!s},\n' \
              '  e1 = {:s},\n' \
              '  e2 = {:s}.'
        indented_L = '\n      '.join(str(self._L).splitlines())
        indented_e1 = '\n       '.join(str(self._e1).splitlines())
        indented_e2 = '\n       '.join(str(self._e2).splitlines())
        return tpl.format(indented_L, str(self._K), indented_e1, indented_e2)


    def solution(self):
        """
        Solve this linear game and return a Solution object.

        OUTPUT:

        If the cone program associated with this game could be
        successfully solved, then a Solution object containing the
        game's value and optimal strategies is returned. If the game
        could *not* be solved -- which should never happen -- then a
        GameUnsolvableException is raised. It can be printed to get the
        raw output from CVXOPT.

        Examples
        --------

        This example is computed in Gowda and Ravindran in the section
        "The value of a Z-transformation"::

            >>> from cones import NonnegativeOrthant
            >>> K = NonnegativeOrthant(3)
            >>> L = [[1,-5,-15],[-1,2,-3],[-12,-15,1]]
            >>> e1 = [1,1,1]
            >>> e2 = [1,1,1]
            >>> SLG = SymmetricLinearGame(L, K, e1, e2)
            >>> print(SLG.solution())
            Game value: -6.1724138
            Player 1 optimal:
              [ 0.5517241]
              [-0.0000000]
              [ 0.4482759]
            Player 2 optimal:
              [0.4482759]
              [0.0000000]
              [0.5517241]

        The value of the following game can be computed using the fact
        that the identity is invertible::

            >>> from cones import NonnegativeOrthant
            >>> K = NonnegativeOrthant(3)
            >>> L = [[1,0,0],[0,1,0],[0,0,1]]
            >>> e1 = [1,2,3]
            >>> e2 = [4,5,6]
            >>> SLG = SymmetricLinearGame(L, K, e1, e2)
            >>> print(SLG.solution())
            Game value: 0.0312500
            Player 1 optimal:
              [0.0312500]
              [0.0625000]
              [0.0937500]
            Player 2 optimal:
              [0.1250000]
              [0.1562500]
              [0.1875000]

        """
        # The cone "C" that appears in the statement of the CVXOPT
        # conelp program.
        C = CartesianProduct(self._K, self._K)

        # The column vector "b" that appears on the right-hand side of
        # Ax = b in the statement of the CVXOPT conelp program.
        b = matrix([1], tc='d')

        # A column of zeros that fits K.
        zero = matrix(0, (self._K.dimension(), 1), tc='d')

        # The column vector "h" that appears on the right-hand side of
        # Gx + s = h in the statement of the CVXOPT conelp program.
        h = matrix([zero, zero])

        # The column vector "c" that appears in the objective function
        # value <c,x> in the statement of the CVXOPT conelp program.
        c = matrix([-1, zero])

        # The matrix "G" that appears on the left-hand side of Gx + s = h
        # in the statement of the CVXOPT conelp program.
        G = append_row(append_col(zero, -identity(self._K.dimension())),
                       append_col(self._e1, -self._L))

        # The matrix "A" that appears on the right-hand side of Ax = b
        # in the statement of the CVXOPT conelp program.
        A = matrix([0, self._e2], (1, self._K.dimension() + 1), 'd')

        # Actually solve the thing and obtain a dictionary describing
        # what happened.
        soln_dict = solvers.conelp(c, G, h, C.cvxopt_dims(), A, b)

        # The "status" field contains "optimal" if everything went
        # according to plan. Other possible values are "primal
        # infeasible", "dual infeasible", "unknown", all of which
        # mean we didn't get a solution. That should never happen,
        # because by construction our game has a solution, and thus
        # the cone program should too.
        if soln_dict['status'] != 'optimal':
            raise GameUnsolvableException(soln_dict)

        p1_value = soln_dict['x'][0]
        p1_optimal = soln_dict['x'][1:]
        p2_optimal = soln_dict['z'][self._K.dimension():]

        return Solution(p1_value, p1_optimal, p2_optimal)

    def dual(self):
        """
        Return the dual game to this game.

        Examples
        --------

            >>> from cones import NonnegativeOrthant
            >>> K = NonnegativeOrthant(3)
            >>> L = [[1,-5,-15],[-1,2,-3],[-12,-15,1]]
            >>> e1 = [1,1,1]
            >>> e2 = [1,2,3]
            >>> SLG = SymmetricLinearGame(L, K, e1, e2)
            >>> print(SLG.dual())
            The linear game (L, K, e1, e2) where
              L = [  1  -1 -12]
                  [ -5   2 -15]
                  [-15  -3   1],
              K = Nonnegative orthant in the real 3-space,
              e1 = [ 1]
                   [ 2]
                   [ 3],
              e2 = [ 1]
                   [ 1]
                   [ 1].

        """
        # We pass ``self._L`` right back into the constructor, because
        # it will be transposed there. And keep in mind that ``self._K``
        # is its own dual.
        return SymmetricLinearGame(self._L,
                                   self._K,
                                   self._e2,
                                   self._e1)


class SymmetricLinearGameTest(TestCase):
    """
    Tests for the SymmetricLinearGame and Solution classes.
    """

    def assert_within_tol(self, first, second):
        """
        Test that ``first`` and ``second`` are equal within our default
        tolerance.
        """
        self.assertTrue(abs(first - second) < options.ABS_TOL)


    def assert_solution_exists(self, L, K, e1, e2):
        """
        Given the parameters needed to construct a SymmetricLinearGame,
        ensure that that game has a solution.
        """
        G = SymmetricLinearGame(L, K, e1, e2)
        soln = G.solution()
        L_matrix = matrix(L).trans()
        expected = inner_product(L_matrix*soln.player1_optimal(),
                                 soln.player2_optimal())
        self.assert_within_tol(soln.game_value(), expected)

    def test_solution_exists_nonnegative_orthant(self):
        """
        Every linear game has a solution, so we should be able to solve
        every symmetric linear game over the NonnegativeOrthant. Pick
        some parameters randomly and give it a shot. The resulting
        optimal solutions should give us the optimal game value when we
        apply the payoff operator to them.
        """
        ambient_dim = randint(1, 10)
        K = NonnegativeOrthant(ambient_dim)
        e1 = [uniform(0.1, 10) for idx in range(K.dimension())]
        e2 = [uniform(0.1, 10) for idx in range(K.dimension())]
        L = [[uniform(-10, 10) for i in range(K.dimension())]
             for j in range(K.dimension())]
        self.assert_solution_exists(L, K, e1, e2)

    def test_solution_exists_ice_cream(self):
        """
        Like :meth:`test_solution_exists_nonnegative_orthant`, except
        over the ice cream cone.
        """
        # Use a minimum dimension of two to avoid divide-by-zero in
        # the fudge factor we make up later.
        ambient_dim = randint(2, 10)
        K = IceCream(ambient_dim)
        e1 = [1] # Set the "height" of e1 to one
        e2 = [1] # And the same for e2

        # If we choose the rest of the components of e1,e2 randomly
        # between 0 and 1, then the largest the squared norm of the
        # non-height part of e1,e2 could be is the 1*(dim(K) - 1). We
        # need to make it less than one (the height of the cone) so
        # that the whole thing is in the cone. The norm of the
        # non-height part is sqrt(dim(K) - 1), and we can divide by
        # twice that.
        fudge_factor = 1.0 / (2.0*sqrt(K.dimension() - 1.0))
        e1 += [fudge_factor*uniform(0, 1) for idx in range(K.dimension() - 1)]
        e2 += [fudge_factor*uniform(0, 1) for idx in range(K.dimension() - 1)]
        L = [[uniform(-10, 10) for i in range(K.dimension())]
             for j in range(K.dimension())]
        self.assert_solution_exists(L, K, e1, e2)