""" Test functions for linalg module

"""
from __future__ import division, absolute_import, print_function

import os
import sys
import itertools
import traceback

import numpy as np
from numpy import array, single, double, csingle, cdouble, dot, identity
from numpy import multiply, atleast_2d, inf, asarray, matrix
from numpy import linalg
from numpy.linalg import matrix_power, norm, matrix_rank
from numpy.testing import (
    assert_, assert_equal, assert_raises, assert_array_equal,
    assert_almost_equal, assert_allclose, run_module_suite,
    dec
)


def ifthen(a, b):
    return not a or b


def imply(a, b):
    return not a or b


old_assert_almost_equal = assert_almost_equal


def assert_almost_equal(a, b, **kw):
    if asarray(a).dtype.type in (single, csingle):
        decimal = 6
    else:
        decimal = 12
    old_assert_almost_equal(a, b, decimal=decimal, **kw)


def get_real_dtype(dtype):
    return {single: single, double: double,
            csingle: single, cdouble: double}[dtype]


def get_complex_dtype(dtype):
    return {single: csingle, double: cdouble,
            csingle: csingle, cdouble: cdouble}[dtype]

def get_rtol(dtype):
    # Choose a safe rtol
    if dtype in (single, csingle):
        return 1e-5
    else:
        return 1e-11

class LinalgCase(object):
    def __init__(self, name, a, b, exception_cls=None):
        assert isinstance(name, str)
        self.name = name
        self.a = a
        self.b = b
        self.exception_cls = exception_cls

    def check(self, do):
        if self.exception_cls is None:
            do(self.a, self.b)
        else:
            assert_raises(self.exception_cls, do, self.a, self.b)

    def __repr__(self):
        return "<LinalgCase: %s>" % (self.name,)


#
# Base test cases
#

np.random.seed(1234)

SQUARE_CASES = [
    LinalgCase("single",
               array([[1., 2.], [3., 4.]], dtype=single),
               array([2., 1.], dtype=single)),
    LinalgCase("double",
               array([[1., 2.], [3., 4.]], dtype=double),
               array([2., 1.], dtype=double)),
    LinalgCase("double_2",
               array([[1., 2.], [3., 4.]], dtype=double),
               array([[2., 1., 4.], [3., 4., 6.]], dtype=double)),
    LinalgCase("csingle",
               array([[1.+2j, 2+3j], [3+4j, 4+5j]], dtype=csingle),
               array([2.+1j, 1.+2j], dtype=csingle)),
    LinalgCase("cdouble",
               array([[1.+2j, 2+3j], [3+4j, 4+5j]], dtype=cdouble),
               array([2.+1j, 1.+2j], dtype=cdouble)),
    LinalgCase("cdouble_2",
               array([[1.+2j, 2+3j], [3+4j, 4+5j]], dtype=cdouble),
               array([[2.+1j, 1.+2j, 1+3j], [1-2j, 1-3j, 1-6j]], dtype=cdouble)),
    LinalgCase("empty",
               atleast_2d(array([], dtype = double)),
               atleast_2d(array([], dtype = double)),
               linalg.LinAlgError),
    LinalgCase("8x8",
               np.random.rand(8, 8),
               np.random.rand(8)),
    LinalgCase("1x1",
               np.random.rand(1, 1),
               np.random.rand(1)),
    LinalgCase("nonarray",
               [[1, 2], [3, 4]],
               [2, 1]),
    LinalgCase("matrix_b_only",
               array([[1., 2.], [3., 4.]]),
               matrix([2., 1.]).T),
    LinalgCase("matrix_a_and_b",
               matrix([[1., 2.], [3., 4.]]),
               matrix([2., 1.]).T),
]

NONSQUARE_CASES = [
    LinalgCase("single_nsq_1",
               array([[1., 2., 3.], [3., 4., 6.]], dtype=single),
               array([2., 1.], dtype=single)),
    LinalgCase("single_nsq_2",
               array([[1., 2.], [3., 4.], [5., 6.]], dtype=single),
               array([2., 1., 3.], dtype=single)),
    LinalgCase("double_nsq_1",
               array([[1., 2., 3.], [3., 4., 6.]], dtype=double),
               array([2., 1.], dtype=double)),
    LinalgCase("double_nsq_2",
               array([[1., 2.], [3., 4.], [5., 6.]], dtype=double),
               array([2., 1., 3.], dtype=double)),
    LinalgCase("csingle_nsq_1",
               array([[1.+1j, 2.+2j, 3.-3j], [3.-5j, 4.+9j, 6.+2j]], dtype=csingle),
               array([2.+1j, 1.+2j], dtype=csingle)),
    LinalgCase("csingle_nsq_2",
               array([[1.+1j, 2.+2j], [3.-3j, 4.-9j], [5.-4j, 6.+8j]], dtype=csingle),
               array([2.+1j, 1.+2j, 3.-3j], dtype=csingle)),
    LinalgCase("cdouble_nsq_1",
               array([[1.+1j, 2.+2j, 3.-3j], [3.-5j, 4.+9j, 6.+2j]], dtype=cdouble),
               array([2.+1j, 1.+2j], dtype=cdouble)),
    LinalgCase("cdouble_nsq_2",
               array([[1.+1j, 2.+2j], [3.-3j, 4.-9j], [5.-4j, 6.+8j]], dtype=cdouble),
               array([2.+1j, 1.+2j, 3.-3j], dtype=cdouble)),
    LinalgCase("cdouble_nsq_1_2",
               array([[1.+1j, 2.+2j, 3.-3j], [3.-5j, 4.+9j, 6.+2j]], dtype=cdouble),
               array([[2.+1j, 1.+2j], [1-1j, 2-2j]], dtype=cdouble)),
    LinalgCase("cdouble_nsq_2_2",
               array([[1.+1j, 2.+2j], [3.-3j, 4.-9j], [5.-4j, 6.+8j]], dtype=cdouble),
               array([[2.+1j, 1.+2j], [1-1j, 2-2j], [1-1j, 2-2j]], dtype=cdouble)),
    LinalgCase("8x11",
               np.random.rand(8, 11),
               np.random.rand(11)),
    LinalgCase("1x5",
               np.random.rand(1, 5),
               np.random.rand(5)),
    LinalgCase("5x1",
               np.random.rand(5, 1),
               np.random.rand(1)),
]

HERMITIAN_CASES = [
    LinalgCase("hsingle",
               array([[1., 2.], [2., 1.]], dtype=single),
               None),
    LinalgCase("hdouble",
               array([[1., 2.], [2., 1.]], dtype=double),
               None),
    LinalgCase("hcsingle",
               array([[1., 2+3j], [2-3j, 1]], dtype=csingle),
               None),
    LinalgCase("hcdouble",
               array([[1., 2+3j], [2-3j, 1]], dtype=cdouble),
               None),
    LinalgCase("hempty",
               atleast_2d(array([], dtype = double)),
               None,
               linalg.LinAlgError),
    LinalgCase("hnonarray",
               [[1, 2], [2, 1]],
               None),
    LinalgCase("matrix_b_only",
               array([[1., 2.], [2., 1.]]),
               None),
    LinalgCase("hmatrix_a_and_b",
               matrix([[1., 2.], [2., 1.]]),
               None),
    LinalgCase("hmatrix_1x1",
               np.random.rand(1, 1),
               None),
]


#
# Gufunc test cases
#

GENERALIZED_SQUARE_CASES = []
GENERALIZED_NONSQUARE_CASES = []
GENERALIZED_HERMITIAN_CASES = []

for tgt, src in ((GENERALIZED_SQUARE_CASES, SQUARE_CASES),
                 (GENERALIZED_NONSQUARE_CASES, NONSQUARE_CASES),
                 (GENERALIZED_HERMITIAN_CASES, HERMITIAN_CASES)):
    for case in src:
        if not isinstance(case.a, np.ndarray):