"""
Backend magic inherited from tensornetwork: numpy backend
"""

# pylint: disable=invalid-name

import logging
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import numpy as np

import tensornetwork
from scipy.linalg import expm, solve, schur
from scipy.special import softmax, expit, jv
from scipy.sparse import coo_matrix, issparse
from scipy.sparse.linalg import lobpcg, LinearOperator
from tensornetwork.backends.numpy import numpy_backend
from .abstract_backend import ExtendedBackend

logger = logging.getLogger(__name__)

dtypestr: str
rdtypestr: str
Tensor = Any


def _sum_numpy(
    self: Any, a: Tensor, axis: Optional[Sequence[int]] = None, keepdims: bool = False
) -> Tensor:
    return np.sum(a, axis=axis, keepdims=keepdims)  # type: ignore
    # see https://github.com/google/TensorNetwork/issues/952


def _convert_to_tensor_numpy(
    self: Any, a: Tensor, dtype: Optional[str] = None
) -> Tensor:
    if not isinstance(a, np.ndarray) and not np.isscalar(a):
        a = np.array(a)
    a = np.asarray(a)
    if dtype is not None:
        if isinstance(dtype, str) and dtype == "bool":
            a = a.astype(bool)
        else:
            a = a.astype(getattr(np, dtype))
    return a


tensornetwork.backends.numpy.numpy_backend.NumPyBackend.sum = _sum_numpy
tensornetwork.backends.numpy.numpy_backend.NumPyBackend.convert_to_tensor = (
    _convert_to_tensor_numpy
)


class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
    """
    see the original backend API at `numpy backend
    <https://github.com/google/TensorNetwork/blob/master/tensornetwork/backends/numpy/numpy_backend.py>`_
    """

    def eye(
        self, N: int, dtype: Optional[str] = None, M: Optional[int] = None
    ) -> Tensor:
        if dtype is None:
            dtype = dtypestr
        r = np.eye(N, M=M)
        return self.cast(r, dtype)

    def ones(self, shape: Sequence[int], dtype: Optional[str] = None) -> Tensor:
        if dtype is None:
            dtype = dtypestr
        r = np.ones(shape)
        return self.cast(r, dtype)

    def zeros(self, shape: Sequence[int], dtype: Optional[str] = None) -> Tensor:
        if dtype is None:
            dtype = dtypestr
        r = np.zeros(shape)
        return self.cast(r, dtype)

    def zeros_like(self, a: Tensor, dtype: Optional[str] = None) -> Tensor:
        if dtype is None:
            dtype = self.dtype(a)
        r = np.zeros_like(a)
        return self.cast(r, dtype)

    def ones_like(self, a: Tensor, dtype: Optional[str] = None) -> Tensor:
        if dtype is None:
            dtype = self.dtype(a)
        r = np.ones_like(a)
        return self.cast(r, dtype)

    def copy(self, a: Tensor) -> Tensor:
        return a.copy()

    def expm(self, a: Tensor) -> Tensor:
        return expm(a)

    def abs(self, a: Tensor) -> Tensor:
        return np.abs(a)

    def sin(self, a: Tensor) -> Tensor:
        return np.sin(a)

    def cos(self, a: Tensor) -> Tensor:
        return np.cos(a)

    # acos acosh asin asinh atan atan2 atanh cosh (cos) tan tanh sinh (sin)
    def acos(self, a: Tensor) -> Tensor:
        return np.arccos(a)

    def acosh(self, a: Tensor) -> Tensor:
        return np.arccosh(a)

    def asin(self, a: Tensor) -> Tensor:
        return np.arcsin(a)

    def asinh(self, a: Tensor) -> Tensor:
        return np.arcsinh(a)

    def atan(self, a: Tensor) -> Tensor:
        return np.arctan(a)

    def atan2(self, y: Tensor, x: Tensor) -> Tensor:
        return np.arctan2(y, x)

    def atanh(self, a: Tensor) -> Tensor:
        return np.arctanh(a)

    def cosh(self, a: Tensor) -> Tensor:
        return np.cosh(a)

    def tan(self, a: Tensor) -> Tensor:
        return np.tan(a)

    def tanh(self, a: Tensor) -> Tensor:
        return np.tanh(a)

    def sinh(self, a: Tensor) -> Tensor:
        return np.sinh(a)

    def size(self, a: Tensor) -> Tensor:
        return a.size

    def eigvalsh(self, a: Tensor) -> Tensor:
        return np.linalg.eigvalsh(a)

    def lobpcg_standard(
        self,
        a: Union[Tensor, Callable[[Tensor], Tensor]],
        x0: Tensor,
        m: int = 100,
        tol: Optional[Union[Tensor, float]] = None,
    ) -> Tuple[Tensor, Tensor, int]:
        if callable(a):
            a = LinearOperator(
                shape=(x0.shape[0], x0.shape[0]), matvec=a, matmat=a, dtype=x0.dtype
            )

        theta, x = lobpcg(a, x0, maxiter=m, tol=tol, largest=True)
        return theta, x, m

    def kron(self, a: Tensor, b: Tensor) -> Tensor:
        return np.kron(a, b)

    def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
        return np.meshgrid(*args, **kwargs)

    def dtype(self, a: Tensor) -> str:
        return a.dtype.__str__()  # type: ignore

    def numpy(self, a: Tensor) -> Tensor:
        return a

    def det(self, a: Tensor) -> Tensor:
        return np.linalg.det(a)

    def schur(self, a: Tensor, output: str = "real") -> Tuple[Tensor, Tensor]:
        return schur(a, output=output)  # type: ignore

    def i(self, dtype: Any = None) -> Tensor:
        if not dtype:
            dtype = npdtype  # type: ignore
        if isinstance(dtype, str):
            if dtype == "bool":
                dtype = bool
            else:
                dtype = getattr(np, dtype)
        return np.array(1j, dtype=dtype)

    def expand_dims(self, a: Tensor, axis: int) -> Tensor:
        return np.expand_dims(a, axis)

    def stack(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
        return np.stack(a, axis=axis)

    def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
        return np.concatenate(a, axis=axis)

    def tile(self, a: Tensor, rep: Tensor) -> Tensor:
        return np.tile(a, rep)

    def mean(
        self,
        a: Tensor,
        axis: Optional[Sequence[int]] = None,
        keepdims: bool = False,
    ) -> Tensor:
        return np.mean(a, axis=axis, keepdims=keepdims)

    def std(
        self, a: Tensor, axis: Optional[Sequence[int]] = None, keepdims: bool = False
    ) -> Tensor:
        return np.std(a, axis=axis, keepdims=keepdims)

    def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
        return np.unique(a, return_counts=True)  # type: ignore

    def min(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
        return np.min(a, axis=axis)

    def max(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
        return np.max(a, axis=axis)

    def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
        return np.argmax(a, axis=axis)

    def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
        return np.argmin(a, axis=axis)

    def sort(self, a: Tensor, axis: int = -1) -> Tensor:
        return np.sort(a, axis=axis)

    def top_k(self, a: Tensor, k: int) -> Tuple[Tensor, Tensor]:
        # Use np.argpartition for O(n) top-k
        a = np.asarray(a)
        if a.ndim == 1:
            idx = np.argpartition(a, -k)[-k:]
            # partition doesn't guarantee order, so sort them
            idx = idx[np.argsort(a[idx])[::-1]]
            return a[idx], idx
        else:
            # Handle multi-dimensional if needed, but 1D is enough for us now
            raise NotImplementedError("Numpy top_k only implemented for 1D")

    def lexsort(self, keys: Any, axis: int = -1) -> Any:
        return np.lexsort(keys, axis=axis)

    def repeat(self, a: Any, repeats: Any, axis: Optional[int] = None) -> Any:
        return np.repeat(a, repeats, axis=axis)

    def popc(self, a: Any) -> Any:
        if hasattr(np, "bitcount"):  # NumPy 1.25+
            return np.bitcount(a)
        # Fallback for older numpy
        return np.vectorize(lambda x: bin(x).count("1"))(a)

    def sigmoid(self, a: Tensor) -> Tensor:
        return expit(a)

    def relu(self, a: Tensor) -> Tensor:
        return (abs(a) + a) / 2
        # this impl seems to be the fastest
        # see https://stackoverflow.com/questions/32109319/how-to-implement-the-relu-function-in-numpy

    def softmax(self, a: Sequence[Tensor], axis: Optional[int] = None) -> Tensor:
        return softmax(a, axis=axis)

    def onehot(self, a: Tensor, num: int) -> Tensor:
        a = np.asarray(a)
        res = np.eye(num)[a.reshape([-1])]
        return res.reshape(list(a.shape) + [num])
        # https://stackoverflow.com/questions/38592324/one-hot-encoding-using-numpy

    def cumsum(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
        return np.cumsum(a, axis)

    def is_tensor(self, a: Any) -> bool:
        if isinstance(a, np.ndarray):
            return True
        return False

    def real(self, a: Tensor) -> Tensor:
        return np.real(a)

    def imag(self, a: Tensor) -> Tensor:
        return np.imag(a)

    def cast(self, a: Tensor, dtype: Union[str, Any]) -> Tensor:
        if isinstance(dtype, str):
            if dtype == "bool":
                np_dtype = bool
            else:
                np_dtype = getattr(np, dtype)
        else:
            np_dtype = dtype
        if np.iscomplexobj(a) and not np.issubdtype(np_dtype, np.complexfloating):
            a = np.real(a)
        return a.astype(np_dtype)

    def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tensor:
        if stop is None:
            return np.arange(0, start, step)
        return np.arange(start, stop, step)

    def mod(self, x: Tensor, y: Tensor) -> Tensor:
        return np.mod(x, y)

    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
        return np.floor_divide(x, y)

    def floor(self, a: Tensor) -> Tensor:
        return np.floor(a)

    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
        return np.clip(a, a_min, a_max)

    def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
        return np.right_shift(x, y)

    def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
        return np.left_shift(x, y)

    def bitwise_and(self, x: Tensor, y: Tensor) -> Tensor:
        return np.bitwise_and(x, y)

    def bitwise_xor(self, x: Tensor, y: Tensor) -> Tensor:
        return np.bitwise_xor(x, y)

    def bitwise_or(self, x: Tensor, y: Tensor) -> Tensor:
        return np.bitwise_or(x, y)

    def any(self, a: Tensor) -> Any:
        return np.any(a)

    def all(self, a: Tensor, axis: Optional[Sequence[int]] = None) -> Any:
        return np.all(a, axis=axis)  # type: ignore

    def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type: ignore
        # gen, sym, her, pos
        # https://stackoverflow.com/questions/44672029/difference-between-numpy-linalg-solve-and-numpy-linalg-lu-solve/44710451
        return solve(A, b, assume_a=assume_a)

    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
        return jv(np.arange(v), z)

    def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
        return np.searchsorted(a, v, side=side)  # type: ignore

    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
        return np.argsort(a, axis=axis)

    def set_random_state(
        self, seed: Optional[int] = None, get_only: bool = False
    ) -> Any:
        g = np.random.default_rng(seed)  # None auto supported
        if get_only is False:
            self.g = g
        return g

    def stateful_randn(
        self,
        g: np.random.Generator,
        shape: Union[int, Sequence[int]] = 1,
        mean: float = 0,
        stddev: float = 1,
        dtype: Optional[str] = None,
    ) -> Tensor:
        if dtype is None:
            dtype = rdtypestr
        if isinstance(dtype, str):
            dtype = dtype[-2:]
        if isinstance(shape, int):
            shape = (shape,)
        r = g.normal(loc=mean, scale=stddev, size=shape)
        if dtype == "32":
            r = r.astype(np.float32)
        elif dtype == "64":
            r = r.astype(np.float64)
        elif not isinstance(dtype, str):
            r = r.astype(dtype)
        else:
            raise ValueError("unspported `dtype` %s" % dtype)
        return r

    def stateful_randu(
        self,
        g: np.random.Generator,
        shape: Union[int, Sequence[int]] = 1,
        low: float = 0,
        high: float = 1,
        dtype: Optional[str] = None,
    ) -> Tensor:
        if dtype is None:
            dtype = rdtypestr
        if isinstance(dtype, str):
            dtype = dtype[-2:]
        if isinstance(shape, int):
            shape = (shape,)
        r = g.uniform(low=low, high=high, size=shape)
        if dtype == "32":
            r = r.astype(np.float32)
        elif dtype == "64":
            r = r.astype(np.float64)
        elif not isinstance(dtype, str):
            r = r.astype(dtype)
        else:
            raise ValueError("unspported `dtype` %s" % dtype)
        return r

    def stateful_randc(
        self,
        g: np.random.Generator,
        a: Union[int, Sequence[int], Tensor],
        shape: Union[int, Sequence[int]],
        p: Optional[Union[Sequence[float], Tensor]] = None,
    ) -> Tensor:
        if isinstance(shape, int):
            shape = (shape,)
        return g.choice(a, size=shape, replace=True, p=p)

    def scatter(
        self, operand: Tensor, indices: Tensor, updates: Tensor, mode: str = "update"
    ) -> Tensor:
        operand_new = np.copy(operand)
        indices = np.asarray(indices)
        updates = np.asarray(updates)
        index_depth = indices.shape[-1]
        indices_reshaped = indices.reshape(-1, index_depth)
        updates_reshaped = updates.reshape((-1,) + operand.shape[index_depth:])

        # Create a tuple of arrays for each dimension being indexed
        idx_tuple = tuple(indices_reshaped[:, i] for i in range(index_depth))

        if mode == "update":
            operand_new[idx_tuple] = updates_reshaped
        elif mode == "add":
            np.add.at(operand_new, idx_tuple, updates_reshaped)
        elif mode == "sub":
            np.subtract.at(operand_new, idx_tuple, updates_reshaped)
        else:
            raise ValueError(f"Unsupported scatter mode: {mode}")

        return operand_new

    def coo_sparse_matrix(
        self, indices: Tensor, values: Tensor, shape: Tensor
    ) -> Tensor:
        return coo_matrix((values, (indices[:, 0], indices[:, 1])), shape=shape)

    def sparse_dense_matmul(
        self,
        sp_a: Tensor,
        b: Tensor,
    ) -> Tensor:
        return sp_a @ b

    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
        return coo.tocsr()

    def to_dense(self, sp_a: Tensor) -> Tensor:
        return sp_a.todense()

    def is_sparse(self, a: Tensor) -> bool:
        return issparse(a)  # type: ignore

    def where(
        self,
        condition: Tensor,
        x: Optional[Tensor] = None,
        y: Optional[Tensor] = None,
    ) -> Tensor:
        if x is None and y is None:
            return np.where(condition)
        assert x is not None and y is not None
        return np.where(condition, x, y)

    def cond(
        self,
        pred: bool,
        true_fun: Callable[[], Tensor],
        false_fun: Callable[[], Tensor],
    ) -> Tensor:
        if pred:
            return true_fun()
        return false_fun()

    def switch(self, index: Tensor, branches: Sequence[Callable[[], Tensor]]) -> Tensor:
        return branches[index]()

    def device(self, a: Tensor) -> str:
        return "cpu"

    def device_move(self, a: Tensor, dev: Any) -> Tensor:
        if dev == "cpu":
            return a
        raise ValueError("NumPy backend only support CPU device")

    def _dev2str(self, dev: Any) -> str:
        if dev == "cpu":
            return "cpu"
        raise ValueError("NumPy backend only support CPU device")

    def _str2dev(self, str_: str) -> Any:
        if str_ == "cpu":
            return "cpu"
        raise ValueError("NumPy backend only support CPU device")

    def stop_gradient(self, a: Tensor) -> Tensor:
        raise NotImplementedError("numpy backend doesn't support AD")

    def grad(
        self,
        f: Callable[..., Any],
        argnums: Union[int, Sequence[int]] = 0,
        has_aux: bool = False,
    ) -> Callable[..., Any]:
        raise NotImplementedError("numpy backend doesn't support AD")

    def value_and_grad(
        self,
        f: Callable[..., Any],
        argnums: Union[int, Sequence[int]] = 0,
        has_aux: bool = False,
    ) -> Callable[..., Tuple[Any, Any]]:
        raise NotImplementedError("numpy backend doesn't support AD")

    def jit(
        self,
        f: Callable[..., Any],
        static_argnums: Optional[Union[int, Sequence[int]]] = None,
        jit_compile: Optional[bool] = None,
        **kws: Any,
    ) -> Callable[..., Any]:
        logger.info("numpy backend has no jit interface, just do nothing")
        return f
        # raise NotImplementedError("numpy backend doesn't support jit compiling")

    def vmap(
        self, f: Callable[..., Any], vectorized_argnums: Union[int, Sequence[int]] = 0
    ) -> Any:
        logger.info(
            "numpy backend has no intrinsic vmap like interface"
            ", use vectorize instead (plain for loop)"
        )
        if isinstance(vectorized_argnums, int):
            vectorized_argnums = (vectorized_argnums,)

        def wrapper(*args: Any, **kws: Any) -> Tensor:
            results = []
            for barg in zip(*[args[i] for i in vectorized_argnums]):  # type: ignore
                narg = []
                j = 0
                for k in range(len(args)):
                    if k in vectorized_argnums:  # type: ignore
                        narg.append(barg[j])
                        j += 1
                    else:
                        narg.append(args[k])
                results.append(f(*narg, **kws))
            return np.array(results)

        return wrapper

    def vectorized_value_and_grad(
        self,
        f: Callable[..., Any],
        argnums: Union[int, Sequence[int]] = 0,
        vectorized_argnums: Union[int, Sequence[int]] = 0,
        has_aux: bool = False,
    ) -> Callable[..., Tuple[Any, Any]]:
        raise NotImplementedError("numpy backend doesn't support AD")

    vvag = vectorized_value_and_grad

    def vjp(
        self,
        f: Callable[..., Any],
        inputs: Union[Tensor, Sequence[Tensor]],
        v: Union[Tensor, Sequence[Tensor]],
    ) -> Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]:
        raise NotImplementedError("numpy backend doesn't support AD")

    def jvp(
        self,
        f: Callable[..., Any],
        inputs: Union[Tensor, Sequence[Tensor]],
        v: Union[Tensor, Sequence[Tensor]],
    ) -> Tuple[Union[Tensor, Sequence[Tensor]], Union[Tensor, Sequence[Tensor]]]:
        raise NotImplementedError("numpy backend doesn't support AD")
