""" Test functions for linalg.decomp module

"""
from __future__ import division, print_function, absolute_import

__usage__ = """
Build linalg:
  python setup_linalg.py build
Run tests if scipy is installed:
  python -c 'import scipy;scipy.linalg.test()'
"""

import itertools
import numpy as np
from numpy.testing import (assert_equal, assert_almost_equal,
                           assert_array_almost_equal, assert_array_equal,
                           assert_, assert_allclose)

import pytest
from pytest import raises as assert_raises

from scipy._lib.six import xrange

from scipy.linalg import (eig, eigvals, lu, svd, svdvals, cholesky, qr,
     schur, rsf2csf, lu_solve, lu_factor, solve, diagsvd, hessenberg, rq,
     eig_banded, eigvals_banded, eigh, eigvalsh, qr_multiply, qz, orth, ordqz,
     subspace_angles, hadamard, eigvalsh_tridiagonal, eigh_tridiagonal,
     null_space, cdf2rdf)
from scipy.linalg.lapack import dgbtrf, dgbtrs, zgbtrf, zgbtrs, \
     dsbev, dsbevd, dsbevx, zhbevd, zhbevx
from scipy.linalg.misc import norm
from scipy.linalg._decomp_qz import _select_function

from numpy import array, transpose, diag, ones, linalg, \
     argsort, zeros, arange, float32, complex64, dot, conj, identity, \
     ravel, sqrt, iscomplex, shape, sort, conjugate, sign, \
     asarray, isfinite, ndarray, outer, eye, dtype, empty,\
     triu, tril

from numpy.random import normal, seed, random

from scipy.linalg._testutils import assert_no_overwrite
from scipy.sparse.sputils import bmat, matrix

# digit precision to use in asserts for different types
DIGITS = {'d':11, 'D':11, 'f':4, 'F':4}

def clear_fuss(ar, fuss_binary_bits=7):
    """Clears trailing `fuss_binary_bits` of mantissa of a floating number"""
    x = np.asanyarray(ar)
    if np.iscomplexobj(x):
        return clear_fuss(x.real) + 1j * clear_fuss(x.imag)

    significant_binary_bits = np.finfo(x.dtype).nmant
    x_mant, x_exp = np.frexp(x)
    f = 2.0**(significant_binary_bits - fuss_binary_bits)
    x_mant *= f
    np.rint(x_mant, out=x_mant)
    x_mant /= f

    return np.ldexp(x_mant, x_exp)


# XXX: This function should be available through numpy.testing


def assert_dtype_equal(act, des):
    if isinstance(act, ndarray):
        act = act.dtype
    else:
        act = dtype(act)

    if isinstance(des, ndarray):
        des = des.dtype
    else:
        des = dtype(des)

    assert_(act == des, 'dtype mismatch: "%s" (should be "%s") ' % (act, des))

# XXX: This function should not be defined here, but somewhere in
#      scipy.linalg namespace


def symrand(dim_or_eigv):
    """Return a random symmetric (Hermitian) matrix.

    If 'dim_or_eigv' is an integer N, return a NxN matrix, with eigenvalues
        uniformly distributed on (-1,1).

    If 'dim_or_eigv' is  1-D real array 'a', return a matrix whose
                      eigenvalues are 'a'.
    """
    if isinstance(dim_or_eigv, int):
        dim = dim_or_eigv
        d = random(dim)*2 - 1
    elif (isinstance(dim_or_eigv, ndarray) and
          len(dim_or_eigv.shape) == 1):
        dim = dim_or_eigv.shape[0]
        d = dim_or_eigv
    else:
        raise TypeError("input type not supported.")

    v = random_rot(dim)
    h = dot(dot(v.T.conj(), diag(d)), v)
    # to avoid roundoff errors, symmetrize the matrix (again)
    h = 0.5*(h.T+h)
    return h

# XXX: This function should not be defined here, but somewhere in
#      scipy.linalg namespace


def random_rot(dim):
    """Return a random rotation matrix, drawn from the Haar distribution
    (the only uniform distribution on SO(n)).
    The algorithm is described in the paper
    Stewart, G.W., 'The efficient generation of random orthogonal
    matrices with an application to condition estimators', SIAM Journal
    on Numerical Analysis, 17(3), pp. 403-409, 1980.
    For more information see
    https://en.wikipedia.org/wiki/Orthogonal_matrix#Randomization"""
    H = eye(dim)
    D = ones((dim,))
    for n in range(1, dim):
        x = normal(size=(dim-n+1,))
        D[n-1] = sign(x[0])
        x[0] -= D[n-1]*sqrt((x*x).sum())
        # Householder transformation

        Hx = eye(dim-n+1) - 2.*outer(x, x)/(x*x).sum()
        mat = eye(dim)
        mat[n-1:,n-1:] = Hx
        H = dot(H, mat)
    # Fix the last sign such that the determinant is 1
    D[-1] = -D.prod()
    H = (D*H.T).T
    return H


class TestEigVals(object):

    def test_simple(self):
        a = [[1,2,3],[1,2,3],[2,5,6]]
        w = eigvals(a)
        exact_w = [(9+sqrt(93))/2,0,(9-sqrt(93))/2]
        assert_array_almost_equal(w,exact_w)

    def test_simple_tr(self):
        a = array([[1,2,3],[1,2,3],[2,5,6]],'d')
        a = transpose(a).copy()
        a = transpose(a)
        w = eigvals(a)
        exact_w = [(9+sqrt(93))/2,0,(9-sqrt(93))/2]
        assert_array_almost_equal(w,exact_w)

    def test_simple_complex(self):
        a = [[1,2,3],[1,2,3],[2,5,6+1j]]
        w = eigvals(a)
        exact_w = [(9+1j+sqrt(92+6j))/2,
                   0,
                   (9+1j-sqrt(92+6j))/2]
        assert_array_almost_equal(w,exact_w)

    def test_finite(self):
        a = [[1,2,3],[1,2,3],[2,5,6]]
        w = eigvals(a, check_finite=False)
        exact_w = [(9+sqrt(93))/2,0,(9-sqrt(93))/2]
        assert_array_almost_equal(w,exact_w)


class TestEig(object):

    def test_simple(self):
        a = [[1,2,3],[1,2,3],[2,5,6]]
        w,v = eig(a)
        exact_w = [(9+sqrt(93))/2,0,(9-sqrt(93))/2]
        v0 = array([1,1,(1+sqrt(93)/3)/2])
        v1 = array([3.,0,-1])
        v2 = array([1,1,(1-sqrt(93)/3)/2])
        v0 = v0 / sqrt(dot(v0,transpose(v0)))
        v1 = v1 / sqrt(dot(v1,transpose(v1)))
        v2 = v2 / sqrt(dot(v2,transpose(v2)))
        assert_array_almost_equal(w,exact_w)
        assert_array_almost_equal(v0,v[:,0]*sign(v[0,0]))
        assert_array_almost_equal(v1,v[:,1]*sign(v[0,1]))
        assert_array_almost_equal(v2,v[:,2]*sign(v[0,2]))
        for i in range(3):
            assert_array_almost_equal(dot(a,v[:,i]),w[i]*v[:,i])
        w,v = eig(a,left=1,right=0)
        for i in range(3):
            assert_array_almost_equal(dot(transpose(a),v[:,i]),w[i]*v[:,i])

    def test_simple_complex_eig(self):
        a = [[1,2],[-2,1]]
        w,vl,vr = eig(a,left=1,right=1)
        assert_array_almost_equal(w, array([1+2j, 1-2j]))
        for i in range(2):
            assert_array_almost_equal(dot(a,vr[:,i]),w[i]*vr[:,i])
        for i in range(2):
            assert_array_almost_equal(dot(conjugate(transpose(a)),vl[:,i]),
                                      conjugate(w[i])*vl[:,i])

    def test_simple_complex(self):
        a = [[1,2,3],[1,2,3],[2,5,6+1j]]
        w,vl,vr = eig(a,left=1,right=1)
        for i in range(3):
            assert_array_almost_equal(dot(a,vr[:,i]),w[i]*vr[:,i])
        for i in range(3):
            assert_array_almost_equal(dot(conjugate(transpose(a)),vl[:,i]),
                                      conjugate(w[i])*vl[:,i])

    def test_gh_3054(self):
        a = [[1]]
        b = [[0]]
        w, vr = eig(a, b, homogeneous_eigvals=True)
        assert_allclose(w[1,0], 0)
        assert_(w[0,0] != 0)
        assert_allclose(vr, 1)

        w, vr = eig(a, b)
        assert_equal(w, np.inf)
        assert_allclose(vr, 1)

    def _check_gen_eig(self, A, B):
        if B is not None:
            A, B = asarray(A), asarray(B)
            B0 = B
        else:
            A = asarray(A)
            B0 = B
            B = np.eye(*A.shape)
        msg = "\n%r\n%r" % (A, B)

        # Eigenvalues in homogeneous coordinates
        w, vr = eig(A, B0, homogeneous_eigvals=True)
        wt = eigvals(A, B0, homogeneous_eigvals=True)
        val1 = dot(A, vr) * w[1,:]
        val2 = dot(B, vr) * w[0,:]
        for i in range(val1.shape[1]):
            assert_allclose(val1[:,i], val2[:,i], rtol=1e-13, atol=1e-13, err_msg=msg)

        if B0 is None:
            assert_allclose(w[1,:], 1)
            assert_allclose(wt[1,:], 1)

        perm = np.lexsort(w)
        permt = np.lexsort(wt)
        assert_allclose(w[:,perm], wt[:,permt], atol=1e-7, rtol=1e-7,
                        err_msg=msg)

        length = np.empty(len(vr))
        for i in xrange(len(vr)):
            length[i] = norm(vr[:,i])
        assert_allclose(length, np.ones(length.size), err_msg=msg,
                        atol=1e-7, rtol=1e-7)

        # Convert homogeneous coordinates
        beta_nonzero = (w[1,:] != 0)
        wh = w[0,beta_nonzero] / w[1,beta_nonzero]

        # Eigenvalues in standard coordinates
        w, vr = eig(A, B0)
        wt = eigvals(A, B0)
        val1 = dot(A, vr)
        val2 = dot(B, vr) * w
        res = val1 - val2
        for i in range(res.shape[1]):
            if np.all(isfinite(res[:,i])):
                assert_allclose(res[:,i], 0, rtol=1e-13, atol=1e-13, err_msg=msg)

        w_fin = w[isfinite(w)]
        wt_fin = wt[isfinite(wt)]
        perm = argsort(clear_fuss(w_fin))
        permt = argsort(clear_fuss(wt_fin))
        assert_allclose(w[perm], wt[permt],
                        atol=1e-7, rtol=1e-7, err_msg=msg)

        length = np.empty(len(vr))
        for i in xrange(len(vr)):
            length[i] = norm(vr[:,i])
        assert_allclose(length, np.ones(length.size), err_msg=msg)

        # Compare homogeneous and nonhomogeneous versions
        assert_allclose(sort(wh), sort(w[np.isfinite(w)]))

    @pytest.mark.xfail(reason="See gh-2254.")
    def test_singular(self):
        # Example taken from
        # https://web.archive.org/web/20040903121217/http://www.cs.umu.se/research/nla/singular_pairs/guptri/matlab.html
        A = array(([22,34,31,31,17], [45,45,42,19,29], [39,47,49,26,34],
            [27,31,26,21,15], [38,44,44,24,30]))
        B = array(([13,26,25,17,24], [31,46,40,26,37], [26,40,19,25,25],
            [16,25,27,14,23], [24,35,18,21,22]))

        olderr = np.seterr(all='ignore')
        try:
            self._check_gen_eig(A, B)
        finally:
            np.seterr(**olderr)

    def test_falker(self):
        # Test matrices giving some Nan generalized eigenvalues.
        M = diag(array(([1,0,3])))
        K = array(([2,-1,-1],[-1,2,-1],[-1,-1,2]))
        D = array(([1,-1,0],[-1,1,0],[0,0,0]))
        Z = zeros((3,3))
        I3 = identity(3)
        A = bmat([[I3, Z], [Z, -K]])
        B = bmat([[Z, I3], [M, D]])

        olderr = np.seterr(all='ignore')
        try:
            self._check_gen_eig(A, B)
        finally:
            np.seterr(**olderr)

    def test_bad_geneig(self):
        # Ticket #709 (strange return values from DGGEV)

        def matrices(omega):
            c1 = -9 + omega**2
            c2 = 2*omega
            A = [[1, 0, 0, 0],
                 [0, 1, 0, 0],
                 [0, 0, c1, 0],
                 [0, 0, 0, c1]]
            B = [[0, 0, 1, 0],
                 [0, 0, 0, 1],
                 [1, 0, 0, -c2],
                 [0, 1, c2, 0]]
            return A, B

        # With a buggy LAPACK, this can fail for different omega on different
        # machines -- so we need to test several values
        olderr = np.seterr(all='ignore')
        try:
            for k in xrange(100):
                A, B = matrices(omega=k*5./100)
                self._check_gen_eig(A, B)
        finally:
            np.seterr(**olderr)

    def test_make_eigvals(self):
        # Step through all paths in _make_eigvals
        seed(1234)