# Copyright 2019 The TensorNetwork Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#pyling: disable=line-too-long
from typing import Optional, Any, Sequence, Tuple, Callable, List, Text, Type
from typing import Union
from tensornetwork.backends import abstract_backend
from tensornetwork.backends.numpy import decompositions
import io
import numpy as np
import scipy as sp
import scipy.sparse.linalg
Tensor = Any

int_to_string = np.array(list(map(chr, list(range(65, 91)))))


class NumPyBackend(abstract_backend.AbstractBackend):
  """See base_backend.BaseBackend for documentation."""

  def __init__(self) -> None:
    super().__init__()
    self.name = "numpy"

  def tensordot(self, a: Tensor, b: Tensor,
                axes: Union[int, Sequence[Sequence[int]]]) -> Tensor:
    # use einsum for scalar-like products, its much faster
    if not isinstance(axes, int):
      if (len(axes[0]) == a.ndim) and (len(axes[1]) == b.ndim):
        if not len(axes[0]) == len(axes[1]):
          raise ValueError("shape-mismatch for sum")
        u, pos1, _ = np.intersect1d(axes[0],
                                    axes[1],
                                    return_indices=True,
                                    assume_unique=True)
        labels = int_to_string[0:len(u)]
        labels_1 = labels[pos1]
        labels_2 = np.array([''] * len(labels_1))

        labels_2[np.array(axes[1])] = labels
        einsum_label = ','.join([''.join(labels_1), ''.join(labels_2)])
        return np.array(np.einsum(einsum_label, a, b, optimize=True))
      return np.tensordot(a, b, axes)
    return np.tensordot(a, b, axes)

  def reshape(self, tensor: Tensor, shape: Tensor) -> Tensor:
    return np.reshape(tensor, np.asarray(shape).astype(np.int32))

  def transpose(self,
                tensor: Tensor,
                perm: Optional[Sequence] = None) -> Tensor:
    return np.transpose(tensor, perm)

  def slice(self, tensor: Tensor, start_indices: Tuple[int, ...],
            slice_sizes: Tuple[int, ...]) -> Tensor:
    if len(start_indices) != len(slice_sizes):
      raise ValueError("Lengths of start_indices and slice_sizes must be"
                       "identical.")
    obj = tuple(
        slice(start, start + size)
        for start, size in zip(start_indices, slice_sizes))
    return tensor[obj]

  def shape_concat(self, values: Tensor, axis: int) -> Tensor:
    return np.concatenate(values, axis)

  def shape_tensor(self, tensor: Tensor) -> Tensor:
    return tensor.shape

  def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]:
    return tensor.shape

  def sparse_shape(self, tensor: Tensor) -> Tuple[Optional[int], ...]:
    return self.shape_tuple(tensor)

  def shape_prod(self, values: Tensor) -> Tensor:
    return np.prod(values)

  def sqrt(self, tensor: Tensor) -> Tensor:
    return np.sqrt(tensor)

  def convert_to_tensor(self, tensor: Tensor) -> Tensor:
    if (not isinstance(tensor, np.ndarray) and not np.isscalar(tensor)):
      raise TypeError("Expected a `np.array` or scalar. Got {}".format(
          type(tensor)))
    result = np.asarray(tensor)
    return result

  def outer_product(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
    return np.tensordot(tensor1, tensor2, 0)

  def einsum(self,
             expression: str,
             *tensors: Tensor,
             optimize: bool = True) -> Tensor:
    return np.einsum(expression, *tensors, optimize=optimize)

  def norm(self, tensor: Tensor) -> Tensor:
    return np.linalg.norm(tensor)

  def eye(self,
          N,
          dtype: Optional[np.dtype] = None,
          M: Optional[int] = None) -> Tensor:
    dtype = dtype if dtype is not None else np.float64

    return np.eye(N, M=M, dtype=dtype)

  def ones(self,
           shape: Tuple[int, ...],
           dtype: Optional[np.dtype] = None) -> Tensor:
    dtype = dtype if dtype is not None else np.float64
    return np.ones(shape, dtype=dtype)

  def zeros(self,
            shape: Tuple[int, ...],
            dtype: Optional[np.dtype] = None) -> Tensor:
    dtype = dtype if dtype is not None else np.float64
    return np.zeros(shape, dtype=dtype)

  def randn(self,
            shape: Tuple[int, ...],
            dtype: Optional[np.dtype] = None,
            seed: Optional[int] = None) -> Tensor:

    if seed:
      np.random.seed(seed)
    dtype = dtype if dtype is not None else np.float64
    if ((np.dtype(dtype) is np.dtype(np.complex128)) or
        (np.dtype(dtype) is np.dtype(np.complex64))):
      return np.random.randn(
          *shape).astype(dtype) + 1j * np.random.randn(*shape).astype(dtype)
    return np.random.randn(*shape).astype(dtype)

  def random_uniform(self,
                     shape: Tuple[int, ...],
                     boundaries: Optional[Tuple[float, float]] = (0.0, 1.0),
                     dtype: Optional[np.dtype] = None,
                     seed: Optional[int] = None) -> Tensor:

    if seed: