import tensorflow as tf
import numpy as np
import scipy as sp
import jax
import pytest
from tensornetwork.backends.jax import jax_backend
import jax.config as config
# pylint: disable=no-member
config.update("jax_enable_x64", True)
np_randn_dtypes = [np.float32, np.float16, np.float64]
np_dtypes = np_randn_dtypes + [np.complex64, np.complex128]
np_not_half = [np.float32, np.float64, np.complex64, np.complex128]



def test_tensordot():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(2 * np.ones((2, 3, 4)))
  b = backend.convert_to_tensor(np.ones((2, 3, 4)))
  actual = backend.tensordot(a, b, ((1, 2), (1, 2)))
  expected = np.array([[24.0, 24.0], [24.0, 24.0]])
  np.testing.assert_allclose(expected, actual)


def test_reshape():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(np.ones((2, 3, 4)))
  actual = backend.shape_tuple(backend.reshape(a, np.array((6, 4, 1))))
  assert actual == (6, 4, 1)


def test_transpose():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(
      np.array([[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]]]))
  actual = backend.transpose(a, [2, 0, 1])
  expected = np.array([[[1.0, 3.0], [5.0, 7.0]], [[2.0, 4.0], [6.0, 8.0]]])
  np.testing.assert_allclose(expected, actual)


def test_transpose_noperm():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(
      np.array([[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]]]))
  actual = backend.transpose(a) # [2, 1, 0]
  actual = backend.transpose(actual, perm=[0, 2, 1])
  expected = np.array([[[1.0, 3.0], [5.0, 7.0]], [[2.0, 4.0], [6.0, 8.0]]])
  np.testing.assert_allclose(expected, actual)


def test_shape_concat():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(2 * np.ones((1, 3, 1)))
  b = backend.convert_to_tensor(np.ones((1, 2, 1)))
  expected = backend.shape_concat((a, b), axis=1)
  actual = np.array([[[2.0], [2.0], [2.0], [1.0], [1.0]]])
  np.testing.assert_allclose(expected, actual)


def test_slice():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(
      np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]))
  actual = backend.slice(a, (1, 1), (2, 2))
  expected = np.array([[5., 6.], [8., 9.]])
  np.testing.assert_allclose(expected, actual)


def test_slice_raises_error():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(
      np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]))
  with pytest.raises(ValueError):
    backend.slice(a, (1, 1), (2, 2, 2))


def test_shape_tensor():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(np.ones([2, 3, 4]))
  assert isinstance(backend.shape_tensor(a), tuple)
  actual = backend.shape_tensor(a)
  expected = np.array([2, 3, 4])
  np.testing.assert_allclose(expected, actual)


def test_shape_tuple():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(np.ones([2, 3, 4]))
  actual = backend.shape_tuple(a)
  assert actual == (2, 3, 4)


def test_shape_prod():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(2 * np.ones([1, 2, 3, 4]))
  actual = np.array(backend.shape_prod(a))
  assert actual == 2**24


def test_sqrt():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(np.array([4., 9.]))
  actual = backend.sqrt(a)
  expected = np.array([2, 3])
  np.testing.assert_allclose(expected, actual)


def test_convert_to_tensor():
  backend = jax_backend.JaxBackend()
  array = np.ones((2, 3, 4))
  actual = backend.convert_to_tensor(array)
  expected = jax.jit(lambda x: x)(array)
  assert isinstance(actual, type(expected))
  np.testing.assert_allclose(expected, actual)

def test_outer_product():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(2 * np.ones((2, 1)))
  b = backend.convert_to_tensor(np.ones((1, 2, 2)))
  actual = backend.outer_product(a, b)
  expected = np.array([[[[[2.0, 2.0], [2.0, 2.0]]]], [[[[2.0, 2.0], [2.0,
                                                                     2.0]]]]])
  np.testing.assert_allclose(expected, actual)


def test_einsum():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(2 * np.ones((2, 1)))
  b = backend.convert_to_tensor(np.ones((1, 2, 2)))
  actual = backend.einsum('ij,jil->l', a, b)
  expected = np.array([4.0, 4.0])
  np.testing.assert_allclose(expected, actual)


def test_convert_bad_test():
  backend = jax_backend.JaxBackend()
  with pytest.raises(TypeError, match="Expected"):
    backend.convert_to_tensor(tf.ones((2, 2)))


def test_norm():
  backend = jax_backend.JaxBackend()
  a = backend.convert_to_tensor(np.ones((2, 2)))
  assert backend.norm(a) == 2


@pytest.mark.parametrize("dtype", np_dtypes)
def test_eye(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.eye(N=4, M=5, dtype=dtype)
  np.testing.assert_allclose(np.eye(N=4, M=5, dtype=dtype), a)


@pytest.mark.parametrize("dtype", np_dtypes)
def test_ones(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.ones((4, 4), dtype=dtype)
  np.testing.assert_allclose(np.ones((4, 4), dtype=dtype), a)


@pytest.mark.parametrize("dtype", np_dtypes)
def test_zeros(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.zeros((4, 4), dtype=dtype)
  np.testing.assert_allclose(np.zeros((4, 4), dtype=dtype), a)


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_randn(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.randn((4, 4), dtype=dtype)
  assert a.shape == (4, 4)


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_random_uniform(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), dtype=dtype)
  assert a.shape == (4, 4)


@pytest.mark.parametrize("dtype", [np.complex64, np.complex128])
def test_randn_non_zero_imag(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.randn((4, 4), dtype=dtype)
  assert np.linalg.norm(np.imag(a)) != 0.0


@pytest.mark.parametrize("dtype", [np.complex64, np.complex128])
def test_random_uniform_non_zero_imag(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), dtype=dtype)
  assert np.linalg.norm(np.imag(a)) != 0.0


@pytest.mark.parametrize("dtype", np_dtypes)
def test_eye_dtype(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.eye(N=4, M=4, dtype=dtype)
  assert a.dtype == dtype


@pytest.mark.parametrize("dtype", np_dtypes)
def test_ones_dtype(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.ones((4, 4), dtype=dtype)
  assert a.dtype == dtype


@pytest.mark.parametrize("dtype", np_dtypes)
def test_zeros_dtype(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.zeros((4, 4), dtype=dtype)
  assert a.dtype == dtype


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_randn_dtype(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.randn((4, 4), dtype=dtype)
  assert a.dtype == dtype


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_random_uniform_dtype(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), dtype=dtype)
  assert a.dtype == dtype


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_randn_seed(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.randn((4, 4), seed=10, dtype=dtype)
  b = backend.randn((4, 4), seed=10, dtype=dtype)
  np.testing.assert_allclose(a, b)


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_random_uniform_seed(dtype):
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), seed=10, dtype=dtype)
  b = backend.random_uniform((4, 4), seed=10, dtype=dtype)
  np.testing.assert_allclose(a, b)


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def test_random_uniform_boundaries(dtype):
  lb = 1.2
  ub = 4.8
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), seed=10, dtype=dtype)
  b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype)
  assert ((a >= 0).all() and (a <= 1).all() and (b >= lb).all() and
          (b <= ub).all())


def test_random_uniform_behavior():
  seed = 10
  key = jax.random.PRNGKey(seed)
  backend = jax_backend.JaxBackend()
  a = backend.random_uniform((4, 4), seed=seed)
  b = jax.random.uniform(key, (4, 4))
  np.testing.assert_allclose(a, b)


def test_conj():
  backend = jax_backend.JaxBackend()
  real = np.random.rand(2, 2, 2)
  imag = np.random.rand(2, 2, 2)
  a = backend.convert_to_tensor(real + 1j * imag)
  actual = backend.conj(a)
  expected = real - 1j * imag
  np.testing.assert_allclose(expected, actual)


@pytest.mark.parametrize("dtype", np_randn_dtypes)
def index_update(dtype):
  backend = jax_backend.JaxBackend()
  tensor = backend.randn((4, 2, 3), dtype=dtype, seed=10)
  out = backend.index_update(tensor, tensor > 0.1, 0.0)
  tensor = np.array(tensor)
  tensor[tensor > 0.1] = 0.0
  np.testing.assert_allclose(tensor, out)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
def test_eigsh_valid_init_operator_with_shape(dtype):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  init = backend.randn((D,), dtype=dtype, seed=10)
  tmp = backend.randn((D, D), dtype=dtype, seed=10)
  H = tmp + backend.transpose(backend.conj(tmp), (1, 0))

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
  eta2, U2 = np.linalg.eigh(H)
  v2 = U2[:, 0]
  v2 = v2 / sum(v2)
  v1 = np.reshape(U1[0], (D))
  v1 = v1 / sum(v1)
  np.testing.assert_allclose(eta1[0], min(eta2))
  np.testing.assert_allclose(v1, v2)


def test_eigsh_small_number_krylov_vectors():
  backend = jax_backend.JaxBackend()
  init = np.array([1, 1], dtype=np.float64)
  H = np.array([[1, 2], [2, 4]], dtype=np.float64)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta1, _ = backend.eigsh_lanczos(mv, [H], init, numeig=1, num_krylov_vecs=2)
  np.testing.assert_almost_equal(eta1, [0])


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
def test_eigsh_lanczos_1(dtype):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  init = backend.randn((D,), dtype=dtype, seed=10)
  tmp = backend.randn((D, D), dtype=dtype, seed=10)
  H = tmp + backend.transpose(backend.conj(tmp), (1, 0))

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta1, U1 = backend.eigsh_lanczos(mv, [H], init)
  eta2, U2 = np.linalg.eigh(H)
  v2 = U2[:, 0]
  v2 = v2 / sum(v2)
  v1 = np.reshape(U1[0], (D))
  v1 = v1 / sum(v1)
  np.testing.assert_allclose(eta1[0], min(eta2))
  np.testing.assert_allclose(v1, v2)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
def test_eigsh_lanczos_2(dtype):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  tmp = backend.randn((D, D), dtype=dtype, seed=10)
  H = tmp + backend.transpose(backend.conj(tmp), (1, 0))

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta1, U1 = backend.eigsh_lanczos(mv, [H], shape=(D,), dtype=dtype)
  eta2, U2 = np.linalg.eigh(H)
  v2 = U2[:, 0]
  v2 = v2 / sum(v2)
  v1 = np.reshape(U1[0], (D))
  v1 = v1 / sum(v1)
  np.testing.assert_allclose(eta1[0], min(eta2))
  np.testing.assert_allclose(v1, v2)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("numeig", [1, 2, 3, 4])
def test_eigsh_lanczos_reorthogonalize(dtype, numeig):
  backend = jax_backend.JaxBackend()
  D = 24
  np.random.seed(10)
  tmp = backend.randn((D, D), dtype=dtype, seed=10)
  H = tmp + backend.transpose(backend.conj(tmp), (1, 0))

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta1, U1 = backend.eigsh_lanczos(
      mv, [H],
      shape=(D,),
      dtype=dtype,
      numeig=numeig,
      num_krylov_vecs=D,
      reorthogonalize=True,
      ndiag=1,
      tol=1E-12,
      delta=1E-12)
  eta2, U2 = np.linalg.eigh(H)

  np.testing.assert_allclose(eta1[0:numeig], eta2[0:numeig])
  for n in range(numeig):
    v2 = U2[:, n]
    v2 /= np.sum(v2)  #fix phases
    v1 = np.reshape(U1[n], (D))
    v1 /= np.sum(v1)

    np.testing.assert_allclose(v1, v2, rtol=1E-5, atol=1E-5)


def test_eigsh_lanczos_raises():
  backend = jax_backend.JaxBackend()
  with pytest.raises(
      ValueError, match='`num_krylov_vecs` >= `numeig` required!'):
    backend.eigsh_lanczos(lambda x: x, numeig=10, num_krylov_vecs=9)
  with pytest.raises(
      ValueError,
      match="Got numeig = 2 > 1 and `reorthogonalize = False`. "
      "Use `reorthogonalize=True` for `numeig > 1`"):
    backend.eigsh_lanczos(lambda x: x, numeig=2, reorthogonalize=False)
  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigsh_lanczos(lambda x: x, shape=(10,), dtype=None)
  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigsh_lanczos(lambda x: x, shape=None, dtype=np.float64)
  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigsh_lanczos(lambda x: x)
  with pytest.raises(
      TypeError, match="Expected a `jax.array`. Got <class 'list'>"):
    backend.eigsh_lanczos(lambda x: x, initial_state=[1, 2, 3])


@pytest.mark.parametrize("dtype", np_dtypes)
def test_index_update(dtype):
  backend = jax_backend.JaxBackend()
  tensor = backend.randn((4, 2, 3), dtype=dtype, seed=10)
  out = backend.index_update(tensor, tensor > 0.1, 0.0)
  np_tensor = np.array(tensor)
  np_tensor[np_tensor > 0.1] = 0.0
  np.testing.assert_allclose(out, np_tensor)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
def test_broadcast_right_multiplication(dtype):
  backend = jax_backend.JaxBackend()
  tensor1 = backend.randn((2, 3), dtype=dtype, seed=10)
  tensor2 = backend.randn((3,), dtype=dtype, seed=10)
  out = backend.broadcast_right_multiplication(tensor1, tensor2)
  np.testing.assert_allclose(out, np.array(tensor1) * np.array(tensor2))


def test_broadcast_right_multiplication_raises():
  backend = jax_backend.JaxBackend()
  tensor1 = backend.randn((2, 3))
  tensor2 = backend.randn((3, 3))
  with pytest.raises(ValueError):
    backend.broadcast_right_multiplication(tensor1, tensor2)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
def test_broadcast_left_multiplication(dtype):
  backend = jax_backend.JaxBackend()
  tensor1 = backend.randn((3,), dtype=dtype, seed=10)
  tensor2 = backend.randn((3, 4, 2), dtype=dtype, seed=10)
  out = backend.broadcast_left_multiplication(tensor1, tensor2)
  np.testing.assert_allclose(out, np.reshape(tensor1, (3, 1, 1)) * tensor2)


def test_broadcast_left_multiplication_raises():
  dtype = np.float64
  backend = jax_backend.JaxBackend()
  tensor1 = backend.randn((3, 3), dtype=dtype, seed=10)
  tensor2 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
  with pytest.raises(ValueError):
    backend.broadcast_left_multiplication(tensor1, tensor2)


def test_sparse_shape():
  dtype = np.float64
  backend = jax_backend.JaxBackend()
  tensor = backend.randn((2, 3, 4), dtype=dtype, seed=10)
  np.testing.assert_allclose(backend.sparse_shape(tensor), tensor.shape)


@pytest.mark.parametrize("dtype,method", [(np.float64, "sin"),
                                          (np.complex128, "sin"),
                                          (np.float64, "cos"),
                                          (np.complex128, "cos"),
                                          (np.float64, "exp"),
                                          (np.complex128, "exp"),
                                          (np.float64, "log"),
                                          (np.complex128, "log")])
def test_elementwise_ops(dtype, method):
  backend = jax_backend.JaxBackend()
  tensor = backend.randn((4, 3, 2), dtype=dtype, seed=10)
  if method == "log":
    tensor = np.abs(tensor)
  tensor1 = getattr(backend, method)(tensor)
  tensor2 = getattr(np, method)(tensor)
  np.testing.assert_almost_equal(tensor1, tensor2)


@pytest.mark.parametrize("dtype,method", [(np.float64, "expm"),
                                          (np.complex128, "expm")])
def test_matrix_ops(dtype, method):
  backend = jax_backend.JaxBackend()
  matrix = backend.randn((4, 4), dtype=dtype, seed=10)
  matrix1 = getattr(backend, method)(matrix)
  matrix2 = getattr(sp.linalg, method)(matrix)
  np.testing.assert_almost_equal(matrix1, matrix2)


@pytest.mark.parametrize("dtype,method", [(np.float64, "expm"),
                                          (np.complex128, "expm")])
def test_matrix_ops_raises(dtype, method):
  backend = jax_backend.JaxBackend()
  matrix = backend.randn((4, 4, 4), dtype=dtype, seed=10)
  with pytest.raises(ValueError, match=r".*Only matrices.*"):
    getattr(backend, method)(matrix)
  matrix = backend.randn((4, 3), dtype=dtype, seed=10)
  with pytest.raises(ValueError, match=r".*N\*N matrix.*"):
    getattr(backend, method)(matrix)


def test_jit():
  backend = jax_backend.JaxBackend()

  def fun(x, A, y):
    return jax.numpy.dot(x, jax.numpy.dot(A, y))

  fun_jit = backend.jit(fun)
  x = jax.numpy.array(np.random.rand(4))
  y = jax.numpy.array(np.random.rand(4))
  A = jax.numpy.array(np.random.rand(4, 4))
  res1 = fun(x, A, y)
  res2 = fun_jit(x, A, y)
  np.testing.assert_allclose(res1, res2)


def test_jit_args():
  backend = jax_backend.JaxBackend()

  def fun(x, A, y):
    return jax.numpy.dot(x, jax.numpy.dot(A, y))

  fun_jit = backend.jit(fun)
  x = jax.numpy.array(np.random.rand(4))
  y = jax.numpy.array(np.random.rand(4))
  A = jax.numpy.array(np.random.rand(4, 4))

  res1 = fun(x, A, y)
  res2 = fun_jit(x, A, y)
  res3 = fun_jit(x, y=y, A=A)
  np.testing.assert_allclose(res1, res2)
  np.testing.assert_allclose(res1, res3)


def compare_eigvals_and_eigvecs(U, eta, U_exact, eta_exact, thresh=1E-8):
  _, iy = np.nonzero(np.abs(eta[:, None] - eta_exact[None, :]) < thresh)
  U_exact_perm = U_exact[:, iy]
  U_exact_perm = U_exact_perm / np.expand_dims(np.sum(U_exact_perm, axis=0), 0)
  U = U / np.expand_dims(np.sum(U, axis=0), 0)
  np.testing.assert_allclose(U_exact_perm, U)
  np.testing.assert_allclose(eta, eta_exact[iy])


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("which", ["LR", "LM"])
def test_eigs_all_eigvals_with_init(dtype, which):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  init = backend.randn((D,), dtype=dtype, seed=10)
  H = backend.randn((D, D), dtype=dtype, seed=10)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta, U = backend.eigs(mv, [H], init, numeig=D, num_krylov_vecs=D, which=which)
  eta_exact, U_exact = np.linalg.eig(H)
  compare_eigvals_and_eigvecs(
      np.stack(U, axis=1), eta, U_exact, eta_exact, thresh=1E-8)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("which", ["LR", "LM"])
def test_eigs_all_eigvals_no_init(dtype, which):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  H = backend.randn((D, D), dtype=dtype, seed=10)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta, U = backend.eigs(
      mv, [H],
      shape=(D,),
      dtype=dtype,
      numeig=D,
      num_krylov_vecs=D,
      which=which)
  eta_exact, U_exact = np.linalg.eig(H)
  compare_eigvals_and_eigvecs(
      np.stack(U, axis=1), eta, U_exact, eta_exact, thresh=1E-8)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("which", ["LR", "LM"])
def test_eigs_few_eigvals_with_init(dtype, which):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  init = backend.randn((D,), dtype=dtype, seed=10)
  H = backend.randn((D, D), dtype=dtype, seed=10)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta, U = backend.eigs(
      mv, [H], init, numeig=4, num_krylov_vecs=16, maxiter=50, which=which)
  eta_exact, U_exact = np.linalg.eig(H)
  compare_eigvals_and_eigvecs(
      np.stack(U, axis=1), eta, U_exact, eta_exact, thresh=1E-8)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("which", ["LR", "LM"])
def test_eigs_few_eigvals_no_init(dtype, which):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  H = backend.randn((D, D), dtype=dtype, seed=10)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta, U = backend.eigs(
      mv, [H],
      shape=(D,),
      dtype=dtype,
      numeig=4,
      num_krylov_vecs=16,
      which=which)
  eta_exact, U_exact = np.linalg.eig(H)
  compare_eigvals_and_eigvecs(
      np.stack(U, axis=1), eta, U_exact, eta_exact, thresh=1E-8)


@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.parametrize("which", ["LR", "LM"])
def test_eigs_large_ncv_with_init(dtype, which):
  backend = jax_backend.JaxBackend()
  D = 16
  np.random.seed(10)
  init = backend.randn((D,), dtype=dtype, seed=10)
  H = backend.randn((D, D), dtype=dtype, seed=10)

  def mv(x, H):
    return jax.numpy.dot(H, x)

  eta, U = backend.eigs(
      mv, [H], init, numeig=4, num_krylov_vecs=50, maxiter=50, which=which)
  eta_exact, U_exact = np.linalg.eig(H)
  compare_eigvals_and_eigvecs(
      np.stack(U, axis=1), eta, U_exact, eta_exact, thresh=1E-8)

  
def test_eigs_raises():
  backend = jax_backend.JaxBackend()
  with pytest.raises(
      ValueError, match='`num_krylov_vecs` >= `numeig` required!'):
    backend.eigs(lambda x: x, numeig=10, num_krylov_vecs=9)

  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigs(lambda x: x, shape=(10,), dtype=None)
  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigs(lambda x: x, shape=None, dtype=np.float64)
  with pytest.raises(
      ValueError,
      match="if no `initial_state` is passed, then `shape` and"
      "`dtype` have to be provided"):
    backend.eigs(lambda x: x)
  with pytest.raises(
      TypeError, match="Expected a `jax.array`. Got <class 'list'>"):
    backend.eigs(lambda x: x, initial_state=[1, 2, 3])
  for which in ('SI', 'LI', 'SM', 'SR'):
    with pytest.raises(
        ValueError, match=f"which = {which}"
        f" is currently not supported."):
      backend.eigs(lambda x: x, which=which)
  with pytest.raises(KeyError, match="dtype"):
    backend.eigs(lambda x: x, shape=(10,), dtype=np.int32)

##################################################################
#############  This test should just not crash    ################
##################################################################
@pytest.mark.parametrize("dtype",
                         [np.float64, np.complex128, np.float32, np.complex64])
def test_eigs_bugfix(dtype):
  backend = jax_backend.JaxBackend()
  D = 200
  mat = jax.numpy.array(np.random.rand(D, D).astype(dtype))
  x = jax.numpy.array(np.random.rand(D).astype(dtype))

  def matvec_jax(vector, matrix):
    return matrix @ vector

  backend.eigs(
      matvec_jax, [mat],
      numeig=1,
      initial_state=x,
      which='LR',
      maxiter=10,
      num_krylov_vecs=100,
      tol=0.0001)
  #this test will cause some annoying output to std buffer
  with pytest.raises(np.linalg.LinAlgError):
    backend.eigs(
        matvec_jax, [mat],
        numeig=1,
        initial_state=x,
        which='LR',
        maxiter=10,
        num_krylov_vecs=100,
        tol=0.0001,
        res_thresh=0.0)


def test_sum():
  np.random.seed(10)
  backend = jax_backend.JaxBackend()
  tensor = np.random.rand(2, 3, 4)
  a = backend.convert_to_tensor(tensor)
  actual = backend.sum(a, axis=(1, 2))
  expected = np.sum(tensor, axis=(1, 2))
  np.testing.assert_allclose(expected, actual)

  actual = backend.sum(a, axis=(1, 2), keepdims=True)
  expected = np.sum(a, axis=(1, 2), keepdims=True)
  np.testing.assert_allclose(expected, actual)


def test_matmul():
  np.random.seed(10)
  backend = jax_backend.JaxBackend()
  t1 = np.random.rand(10, 2, 3)
  t2 = np.random.rand(10, 3, 4)
  a = backend.convert_to_tensor(t1)
  b = backend.convert_to_tensor(t2)
  actual = backend.matmul(a, b)
  expected = np.matmul(t1, t2)
  np.testing.assert_allclose(expected, actual)
  t3 = np.random.rand(10)
  t4 = np.random.rand(11)
  c = backend.convert_to_tensor(t3)
  d = backend.convert_to_tensor(t4)
  with pytest.raises(ValueError, match="inputs to"):
    backend.matmul(c, d)

def test_gmres_raises():
  backend = jax_backend.JaxBackend()
  dummy_mv = lambda x: x
  N = 10

  b = jax.numpy.zeros((N,))
  x0 = jax.numpy.zeros((N+1),)
  diff = "If x0 is supplied, its shape"
  with pytest.raises(ValueError, match=diff): # x0, b have different sizes
    backend.gmres(dummy_mv, b, x0=x0)

  x0 = jax.numpy.zeros((N,), dtype=jax.numpy.float32)
  b = jax.numpy.zeros((N,), dtype=jax.numpy.float64)
  diff = (f"If x0 is supplied, its dtype, {x0.dtype}, must match b's"
          f", {b.dtype}.")
  with pytest.raises(ValueError, match=diff): # x0, b have different dtypes
    backend.gmres(dummy_mv, b, x0=x0)

  x0 = jax.numpy.zeros((N,))
  b = jax.numpy.zeros((N,)).reshape(2, N//2)
  diff = "If x0 is supplied, its shape"
  with pytest.raises(ValueError, match=diff): # x0, b have different shapes
    backend.gmres(dummy_mv, b, x0=x0)

  num_krylov_vectors = 0
  diff = (f"num_krylov_vectors must be in "
          f"0 < {num_krylov_vectors} <= {b.size}")
  with pytest.raises(ValueError, match=diff): # num_krylov_vectors <= 0
    backend.gmres(dummy_mv, b, num_krylov_vectors=num_krylov_vectors)
  num_krylov_vectors = N+1
  diff = (f"num_krylov_vectors must be in "
          f"0 < {num_krylov_vectors} <= {b.size}")
  with pytest.raises(ValueError, match=diff): # num_krylov_vectors > b.size
    backend.gmres(dummy_mv, b, num_krylov_vectors=num_krylov_vectors)

  tol = -1.
  diff = (f"tol = {tol} must be positive.")
  with pytest.raises(ValueError, match=diff): # tol < 0
    backend.gmres(dummy_mv, b, tol=tol)

  atol = -1
  diff = (f"atol = {atol} must be positive.")
  with pytest.raises(ValueError, match=diff): # atol < 0
    backend.gmres(dummy_mv, b, atol=atol)

  M = lambda x: x
  diff = "M is not supported by the Jax backend."
  with pytest.raises(NotImplementedError, match=diff):
    backend.gmres(dummy_mv, b, M=M)

  A_kwargs = {"bee": "honey"}
  diff = "A_kwargs is not supported by the Jax backend."
  with pytest.raises(NotImplementedError, match=diff):
    backend.gmres(dummy_mv, b, A_kwargs=A_kwargs)


jax_qr_dtypes = [np.float32, np.float64, np.complex64, np.complex128]
@pytest.mark.parametrize("dtype", jax_qr_dtypes)
def test_gmres_on_small_known_problem(dtype):
  dummy = jax.numpy.zeros(1, dtype=dtype)
  dtype = dummy.dtype

  backend = jax_backend.JaxBackend()
  A = jax.numpy.array(([[1, 1], [3, -4]]), dtype=dtype)
  b = jax.numpy.array([3, 2], dtype=dtype)
  x0 = jax.numpy.ones(2, dtype=dtype)
  n_kry = 2

  def A_mv(x):
    return A @ x
  tol = 100*jax.numpy.finfo(dtype).eps
  x, _ = backend.gmres(A_mv, b, x0=x0, num_krylov_vectors=n_kry, tol=tol)
  solution = jax.numpy.array([2., 1.], dtype=dtype)
  eps = jax.numpy.linalg.norm(jax.numpy.abs(solution) - jax.numpy.abs(x))
  assert eps < tol


@pytest.mark.parametrize("dtype", jax_qr_dtypes)
def test_gmres_on_larger_random_problem(dtype):
  dummy = jax.numpy.zeros(1, dtype=dtype)
  dtype = dummy.dtype
  backend = jax_backend.JaxBackend()
  matshape = (100, 100)
  vecshape = (100,)
  A = backend.randn(matshape, seed=10, dtype=dtype)
  solution = backend.randn(vecshape, seed=10, dtype=dtype)
  def A_mv(x):
    return A @ x
  b = A_mv(solution)
  tol = b.size * jax.numpy.finfo(dtype).eps
  x, _ = backend.gmres(A_mv, b, tol=tol) # atol = tol by default
  err = jax.numpy.linalg.norm(jax.numpy.abs(x)-jax.numpy.abs(solution))
  rtol = tol*jax.numpy.linalg.norm(b)
  atol = tol
  assert err < max(rtol, atol)


@pytest.mark.parametrize("dtype", np_not_half)
def test_gmres_not_matrix(dtype):
  dummy = jax.numpy.zeros(1, dtype=dtype)
  dtype = dummy.dtype
  backend = jax_backend.JaxBackend()
  matshape = (100, 100)
  vecshape = (100,)
  A = backend.randn(matshape, dtype=dtype, seed=10)
  A = backend.reshape(A, (2, 50, 2, 50))
  solution = backend.randn(vecshape, dtype=dtype, seed=10)
  solution = backend.reshape(solution, (2, 50))
  def A_mv(x):
    return backend.einsum('ijkl,kl', A, x)
  b = A_mv(solution)
  tol = b.size * np.finfo(dtype).eps
  x, _ = backend.gmres(A_mv, b, tol=tol) # atol = tol by default
  err = jax.numpy.linalg.norm(jax.numpy.abs(x)-jax.numpy.abs(solution))
  rtol = tol*jax.numpy.linalg.norm(b)
  atol = tol
  assert err < max(rtol, atol)


@pytest.mark.parametrize("dtype", np_dtypes)
@pytest.mark.parametrize("offset", range(-2, 2))
@pytest.mark.parametrize("axis1", range(0, 3))
@pytest.mark.parametrize("axis2", range(0, 3))
def test_diagonal(dtype, offset, axis1, axis2):
  shape = (5, 5, 5, 5)
  backend = jax_backend.JaxBackend()
  array = backend.randn(shape, dtype=dtype, seed=10)
  if axis1 == axis2:
    with pytest.raises(ValueError):
      actual = backend.diagonal(array, offset=offset, axis1=axis1, axis2=axis2)
  else:
    actual = backend.diagonal(array, offset=offset, axis1=axis1, axis2=axis2)
    expected = jax.numpy.diagonal(array, offset=offset, axis1=axis1,
                                  axis2=axis2)
    np.testing.assert_allclose(actual, expected)


@pytest.mark.parametrize("dtype", np_dtypes)
@pytest.mark.parametrize("offset", range(-2, 2))
def test_diagflat(dtype, offset):
  shape = (5, 5, 5, 5)
  backend = jax_backend.JaxBackend()
  array = backend.randn(shape, dtype=dtype, seed=10)
  actual = backend.diagflat(array, k=offset)
  expected = jax.numpy.diag(jax.numpy.ravel(array), k=offset)
  np.testing.assert_allclose(actual, expected)


@pytest.mark.parametrize("dtype", np_dtypes)
@pytest.mark.parametrize("offset", range(-2, 2))
@pytest.mark.parametrize("axis1", range(0, 3))
@pytest.mark.parametrize("axis2", range(0, 3))
def test_trace(dtype, offset, axis1, axis2):
  shape = (5, 5, 5, 5)
  backend = jax_backend.JaxBackend()
  array = backend.randn(shape, dtype=dtype, seed=10)
  if axis1 == axis2:
    with pytest.raises(ValueError):
      actual = backend.trace(array, offset=offset, axis1=axis1, axis2=axis2)
  else:
    actual = backend.trace(array, offset=offset, axis1=axis1, axis2=axis2)
    expected = jax.numpy.trace(array, offset=offset, axis1=axis1, axis2=axis2)
    np.testing.assert_allclose(actual, expected)


@pytest.mark.parametrize("dtype", np_dtypes)
def test_abs(dtype):
  shape = (4, 3, 2)
  backend = jax_backend.JaxBackend()
  tensor = backend.randn(shape, dtype=dtype, seed=10)
  actual = backend.abs(tensor)
  expected = jax.numpy.abs(tensor)
  np.testing.assert_allclose(expected, actual)


@pytest.mark.parametrize("dtype", np_dtypes)
def test_sign(dtype):
  shape = (4, 3, 2)
  backend = jax_backend.JaxBackend()
  tensor = backend.randn(shape, dtype=dtype, seed=10)
  actual = backend.sign(tensor)
  expected = jax.numpy.sign(tensor)
  np.testing.assert_allclose(expected, actual)


@pytest.mark.parametrize("pivot_axis", [-1, 1, 2])
@pytest.mark.parametrize("dtype", np_dtypes)
def test_pivot(dtype, pivot_axis):
  shape = (4, 3, 2, 8)
  pivot_shape = (np.prod(shape[:pivot_axis]), np.prod(shape[pivot_axis:]))
  backend = jax_backend.JaxBackend()
  tensor = backend.randn(shape, dtype=dtype, seed=10)
  expected = tensor.reshape(pivot_shape)
  actual = backend.pivot(tensor, pivot_axis=pivot_axis)
  np.testing.assert_allclose(expected, actual)


@pytest.mark.parametrize("dtype, atol", [(np.float32, 1E-6),
                                         (np.float64, 1E-10),
                                         (np.complex64, 1E-6),
                                         (np.complex128, 1E-10)])
def test_inv(dtype, atol):
  shape = (10, 10)
  backend = jax_backend.JaxBackend()
  matrix = backend.randn(shape, dtype=dtype, seed=10)
  inv = backend.inv(matrix)
  np.testing.assert_allclose(inv @ matrix, np.eye(10), atol=atol)
  np.testing.assert_allclose(matrix @ inv, np.eye(10), atol=atol)
  tensor = backend.randn((10, 10, 10), dtype=dtype, seed=10)
  with pytest.raises(ValueError, match="input to"):
    backend.inv(tensor)

numpy.testing.assert_allclose numpy.random.rand numpy.finfo numpy.imag numpy.linalg.eigh numpy.eye numpy.prod numpy.array numpy.reshape numpy.matmul numpy.zeros numpy.testing.assert_almost_equal numpy.stack numpy.random.seed numpy.sum numpy.ones numpy.linalg.eig numpy.abs