from sympy import Matrix, zeros, ones, Integer

from sympy.physics.quantum.matrixutils import (
    to_sympy, to_numpy, to_scipy_sparse, matrix_tensor_product,
    matrix_to_zero
)

from sympy.external import import_module
from sympy.utilities.pytest import skip

m = Matrix([[1,2],[3,4]])


def test_sympy_to_sympy():
    assert to_sympy(m) == m


def test_matrix_to_zero():
    assert matrix_to_zero(m) == m
    assert matrix_to_zero(Matrix([[0,0],[0,0]])) == Integer(0)

np = import_module('numpy', min_python_version=(2, 6))

def test_to_numpy():
    if not np:
        skip("numpy not installed or Python too old.")

    result = np.matrix([[1,2],[3,4]], dtype='complex')
    assert (to_numpy(m) == result).all()

def test_matrix_tensor_product():
    if not np:
        skip("numpy not installed or Python too old.")

    l1 = zeros(4)
    for i in range(16):
        l1[i] = 2**i
    l2 = zeros(4)
    for i in range(16):
        l2[i] = i
    l3 = zeros(2)
    for i in range(4):
        l3[i] = i
    vec = Matrix([1,2,3])

    #test for Matrix known 4x4 matricies
    numpyl1 = np.matrix(l1.tolist())
    numpyl2 = np.matrix(l2.tolist())
    numpy_product = np.kron(numpyl1,numpyl2)
    args = [l1, l2]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()
    numpy_product = np.kron(numpyl2,numpyl1)
    args = [l2, l1]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()

    #test for other known matrix of different dimensions
    numpyl2 = np.matrix(l3.tolist())
    numpy_product = np.kron(numpyl1,numpyl2)
    args = [l1, l3]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()
    numpy_product = np.kron(numpyl2,numpyl1)
    args = [l3, l1]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()

    #test for non square matrix
    numpyl2 = np.matrix(vec.tolist())
    numpy_product = np.kron(numpyl1,numpyl2)
    args = [l1, vec]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()
    numpy_product = np.kron(numpyl2,numpyl1)
    args = [vec, l1]
    sympy_product = matrix_tensor_product(*args)
    assert numpy_product.tolist() == sympy_product.tolist()

    #test for random matrix with random values that are floats
    random_matrix1 = np.random.rand(np.random.rand()*5+1,np.random.rand()*5+1)
    random_matrix2 = np.random.rand(np.random.rand()*5+1,np.random.rand()*5+1)
    numpy_product = np.kron(random_matrix1,random_matrix2)
    args = [Matrix(random_matrix1.tolist()),Matrix(random_matrix2.tolist())]
    sympy_product = matrix_tensor_product(*args)
    assert not (sympy_product - Matrix(numpy_product.tolist())).tolist() > \
    (ones((sympy_product.rows,sympy_product.cols))*epsilon).tolist()

    #test for three matrix kronecker
    sympy_product = matrix_tensor_product(l1,vec,l2)

    numpy_product = np.kron(l1,np.kron(vec,l2))
    assert numpy_product.tolist() == sympy_product.tolist()


scipy = import_module('scipy', __import__kwargs={'fromlist':['sparse']})

def test_to_scipy_sparse():
    if not np:
        skip("numpy not installed or Python too old.")
    if not scipy:
        skip("scipy not installed.")
    else:
        sparse = scipy.sparse

    result = sparse.csr_matrix([[1,2],[3,4]], dtype='complex')
    assert np.linalg.norm((to_scipy_sparse(m) - result).todense()) == 0.0

epsilon = .000001
