"""
Unit test for Linear Programming
"""
from __future__ import division, print_function, absolute_import

import numpy as np
from numpy.testing import (assert_, assert_allclose, assert_equal,
                           assert_array_less)
from pytest import raises as assert_raises
from scipy.optimize import linprog, OptimizeWarning
from scipy._lib._numpy_compat import _assert_warns, suppress_warnings
from scipy.sparse.linalg import MatrixRankWarning
from scipy.linalg import LinAlgWarning
import pytest

has_umfpack = True
try:
    from scikits.umfpack import UmfpackWarning
except ImportError:
    has_umfpack = False

has_cholmod = True
try:
    import sksparse
except ImportError:
    has_cholmod = False


def _assert_iteration_limit_reached(res, maxiter):
    assert_(not res.success, "Incorrectly reported success")
    assert_(res.success < maxiter, "Incorrectly reported number of iterations")
    assert_equal(res.status, 1, "Failed to report iteration limit reached")


def _assert_infeasible(res):
    # res: linprog result object
    assert_(not res.success, "incorrectly reported success")
    assert_equal(res.status, 2, "failed to report infeasible status")


def _assert_unbounded(res):
    # res: linprog result object
    assert_(not res.success, "incorrectly reported success")
    assert_equal(res.status, 3, "failed to report unbounded status")


def _assert_unable_to_find_basic_feasible_sol(res):
    # res: linprog result object
    assert_(not res.success, "incorrectly reported success")
    assert_equal(res.status, 2, "failed to report optimization failure")


def _assert_success(res, desired_fun=None, desired_x=None,
                    rtol=1e-8, atol=1e-8):
    # res: linprog result object
    # desired_fun: desired objective function value or None
    # desired_x: desired solution or None
    if not res.success:
        msg = "linprog status {0}, message: {1}".format(res.status,
                                                        res.message)
        raise AssertionError(msg)

    assert_equal(res.status, 0)
    if desired_fun is not None:
        assert_allclose(res.fun, desired_fun,
                        err_msg="converged to an unexpected objective value",
                        rtol=rtol, atol=atol)
    if desired_x is not None:
        assert_allclose(res.x, desired_x,
                        err_msg="converged to an unexpected solution",
                        rtol=rtol, atol=atol)


def magic_square(n):
    """
    Generates a linear program for which integer solutions represent an
    n x n magic square; binary decision variables represent the presence
    (or absence) of an integer 1 to n^2 in each position of the square.
    """

    np.random.seed(0)
    M = n * (n**2 + 1) / 2

    numbers = np.arange(n**4) // n**2 + 1

    numbers = numbers.reshape(n**2, n, n)

    zeros = np.zeros((n**2, n, n))

    A_list = []
    b_list = []

    # Rule 1: use every number exactly once
    for i in range(n**2):
        A_row = zeros.copy()
        A_row[i, :, :] = 1
        A_list.append(A_row.flatten())
        b_list.append(1)

    # Rule 2: Only one number per square
    for i in range(n):
        for j in range(n):
            A_row = zeros.copy()
            A_row[:, i, j] = 1
            A_list.append(A_row.flatten())
            b_list.append(1)

    # Rule 3: sum of rows is M
    for i in range(n):
        A_row = zeros.copy()
        A_row[:, i, :] = numbers[:, i, :]
        A_list.append(A_row.flatten())
        b_list.append(M)

    # Rule 4: sum of columns is M
    for i in range(n):
        A_row = zeros.copy()
        A_row[:, :, i] = numbers[:, :, i]
        A_list.append(A_row.flatten())
        b_list.append(M)

    # Rule 5: sum of diagonals is M
    A_row = zeros.copy()
    A_row[:, range(n), range(n)] = numbers[:, range(n), range(n)]
    A_list.append(A_row.flatten())
    b_list.append(M)
    A_row = zeros.copy()
    A_row[:, range(n), range(-1, -n - 1, -1)] = \
        numbers[:, range(n), range(-1, -n - 1, -1)]
    A_list.append(A_row.flatten())
    b_list.append(M)

    A = np.array(np.vstack(A_list), dtype=float)
    b = np.array(b_list, dtype=float)
    c = np.random.rand(A.shape[1])

    return A, b, c, numbers


def lpgen_2d(m, n):
    """ -> A b c LP test: m*n vars, m+n constraints
        row sums == n/m, col sums == 1
        https://gist.github.com/denis-bz/8647461
    """
    np.random.seed(0)
    c = - np.random.exponential(size=(m, n))
    Arow = np.zeros((m, m * n))
    brow = np.zeros(m)
    for j in range(m):
        j1 = j + 1
        Arow[j, j * n:j1 * n] = 1
        brow[j] = n / m

    Acol = np.zeros((n, m * n))
    bcol = np.zeros(n)
    for j in range(n):
        j1 = j + 1
        Acol[j, j::n] = 1
        bcol[j] = 1

    A = np.vstack((Arow, Acol))
    b = np.hstack((brow, bcol))

    return A, b, c.ravel()


def nontrivial_problem():
    c = [-1, 8, 4, -6]
    A_ub = [[-7, -7, 6, 9],
            [1, -1, -3, 0],
            [10, -10, -7, 7],
            [6, -1, 3, 4]]
    b_ub = [-3, 6, -6, 6]
    A_eq = [[-10, 1, 1, -8]]
    b_eq = [-4]
    x_star = [101 / 1391, 1462 / 1391, 0, 752 / 1391]
    f_star = 7083 / 1391
    return c, A_ub, b_ub, A_eq, b_eq, x_star, f_star


def generic_callback_test(self):
    # Check that callback is as advertised
    last_cb = {}

    def cb(res):
        message = res.pop('message')
        complete = res.pop('complete')

        assert_(res.pop('phase') in (1, 2))
        assert_(res.pop('status') in range(4))
        assert_(isinstance(res.pop('nit'), int))
        assert_(isinstance(complete, bool))
        assert_(isinstance(message, str))

        last_cb['x'] = res['x']
        last_cb['fun'] = res['fun']
        last_cb['slack'] = res['slack']
        last_cb['con'] = res['con']

    c = np.array([-3, -2])
    A_ub = [[2, 1], [1, 1], [1, 0]]
    b_ub = [10, 8, 4]
    res = linprog(c, A_ub=A_ub, b_ub=b_ub, callback=cb, method=self.method)

    _assert_success(res, desired_fun=-18.0, desired_x=[2, 6])
    assert_allclose(last_cb['fun'], res['fun'])
    assert_allclose(last_cb['x'], res['x'])
    assert_allclose(last_cb['con'], res['con'])
    assert_allclose(last_cb['slack'], res['slack'])


def test_unknown_solver():
    c = np.array([-3, -2])
    A_ub = [[2, 1], [1, 1], [1, 0]]
    b_ub = [10, 8, 4]

    assert_raises(ValueError, linprog,
                  c, A_ub=A_ub, b_ub=b_ub, method='ekki-ekki-ekki')


A_ub = None
b_ub = None
A_eq = None
b_eq = None
bounds = None

################
# Common Tests #
################


class LinprogCommonTests(object):
    """
    Base class for `linprog` tests. Generally, each test will be performed
    once for every derived class of LinprogCommonTests, each of which will
    typically change self.options and/or self.method. Effectively, these tests
    are run for many combination of method (simplex, revised simplex, and
    interior point) and options (such as pivoting rule or sparse treatment).
    """

    ##################
    # Targeted Tests #
    ##################

    def test_callback(self):
        generic_callback_test(self)

    def test_disp(self):
        # test that display option does not break anything.
        A, b, c = lpgen_2d(20, 20)
        res = linprog(c, A_ub=A, b_ub=b, method=self.method,
                      options={"disp": True})
        _assert_success(res, desired_fun=-64.049494229)

    def test_docstring_example(self):
        # Example from linprog docstring.
        c = [-1, 4]
        A = [[-3, 1], [1, 2]]
        b = [6, 4]
        x0_bounds = (None, None)
        x1_bounds = (-3, None)
        res = linprog(c, A_ub=A, b_ub=b, bounds=(x0_bounds, x1_bounds),
                      options=self.options, method=self.method)
        _assert_success(res, desired_fun=-22)

    def test_type_error(self):
        # (presumably) checks that linprog recognizes type errors
        # This is tested more carefully in test__linprog_clean_inputs.py
        c = [1]
        A_eq = [[1]]
        b_eq = "hello"
        assert_raises(TypeError, linprog,
                      c, A_eq=A_eq, b_eq=b_eq,
                      method=self.method, options=self.options)

    def test_aliasing_b_ub(self):
        # (presumably) checks that linprog does not modify b_ub
        # This is tested more carefully in test__linprog_clean_inputs.py
        c = np.array([1.0])
        A_ub = np.array([[1.0]])
        b_ub_orig = np.array([3.0])
        b_ub = b_ub_orig.copy()
        bounds = (-4.0, np.inf)
        res = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds,
                      method=self.method, options=self.options)
        _assert_success(res, desired_fun=-4, desired_x=[-4])
        assert_allclose(b_ub_orig, b_ub)

    def test_aliasing_b_eq(self):
        # (presumably) checks that linprog does not modify b_eq
        # This is tested more carefully in test__linprog_clean_inputs.py
        c = np.array([1.0])
        A_eq = np.array([[1.0]])
        b_eq_orig = np.array([3.0])
        b_eq = b_eq_orig.copy()
        bounds = (-4.0, np.inf)
        res = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds,
                      method=self.method, options=self.options)
        _assert_success(res, desired_fun=3, desired_x=[3])
        assert_allclose(b_eq_orig, b_eq)

    def test_non_ndarray_args(self):
        # (presumably) checks that linprog accepts list in place of arrays
        # This is tested more carefully in test__linprog_clean_inputs.py
        c = [1.0]
        A_ub = [[1.0]]
        b_ub = [3.0]
        A_eq = [[1.0]]
        b_eq = [2.0]
        bounds = (-1.0, 10.0)
        res = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds,
                      method=self.method, options=self.options)
        _assert_success(res, desired_fun=2, desired_x=[2])

    def test_unknown_options(self):
        c = np.array([-3, -2])
        A_ub = [[2, 1], [1, 1], [1, 0]]
        b_ub = [10, 8, 4]

        def f(c, A_ub=None, b_ub=None, A_eq=None,
              b_eq=None, bounds=None, options={}):
            linprog(c, A_ub, b_ub, A_eq, b_eq, bounds,
                    method=self.method, options=options)

        o = {key: self.options[key] for key in self.options}
        o['spam'] = 42

        _assert_warns(OptimizeWarning, f,
                      c, A_ub=A_ub, b_ub=b_ub, options=o)

    def test_invalid_inputs(self):

        def f(c, A_ub=None, b_ub=None, A_eq=None, b_eq=None, bounds=None):
            linprog(c, A_ub, b_ub, A_eq, b_eq, bounds,
                    method=self.method, options=self.options)

        for bad_bound in [[(5, 0), (1, 2), (3, 4)],
                          [(1, 2), (3, 4)],
                          [(1, 2), (3, 4), (3, 4, 5)],
                          [(1, 2), (np.inf, np.inf), (3, 4)],
                          [(1, 2), (-np.inf, -np.inf), (3, 4)],
                          ]:
            assert_raises(ValueError, f, [1, 2, 3], bounds=bad_bound)

        assert_raises(ValueError, f, [1, 2], A_ub=[[1, 2]], b_ub=[1, 2])