# Copyright 2021-2022 The NetKet Authors - All rights reserved.
#
# 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.

from typing import Union, SupportsFloat, SupportsComplex, TYPE_CHECKING

import numbers
import warnings

import numpy as np
import jax.numpy as jnp
from scipy.sparse import issparse

from netket.hilbert import AbstractHilbert
from netket.utils.types import DType, Array
from netket.utils.numbers import dtype as _dtype, is_scalar
from netket.errors import OperatorMultiplicationDeprecationWarning

from .._discrete_operator import DiscreteOperator
from .._lazy import Transpose

from .helpers import (
    canonicalize_input,
    _multiply_operators,
    cast_operator_matrix_dtype,
)
from .convert import local_operators_to_pauli_strings

if TYPE_CHECKING:
    from .._pauli_strings import PauliStrings


def is_hermitian(a: np.ndarray, rtol=1e-05, atol=1e-08) -> bool:
    if issparse(a):
        return np.allclose(a.todense(), a.T.conj().todense(), rtol=rtol, atol=atol)
    else:
        return np.allclose(a, a.T.conj(), rtol=rtol, atol=atol)


def _is_sorted(a):
    for i in range(len(a) - 1):
        if a[i + 1] < a[i]:
            return False
    return True


class LocalOperatorBase(DiscreteOperator):
    """Base implementation of an operator composed of a sum of
    local terms, each of which acts on a small number of sites.

    .. warning::

        The complexity of the constructor of this operator scales
        exponentially with the number of sites on which each term
        acts, and linearly with the number of terms.

        The memory requirement scales the same, because this operator
        stores internally a lookup table as large as the local hilbert
        space size of each term.

        Refrain from using this for terms acting on more than
        6-sites. For Spin-1/2 systems, prefer instead
        :class:`netket.operator.PauliStrings`. For non spin-1/2
        there is nothing that will work efficiently out of the box,
        but you can easily roll your own.

    The runtime complexity is proportional to the number of
    connected entries
    :attr:`~netket.operator.DiscreteOperator.max_conn_size`
    """

    def __init__(
        self,
        hilbert: AbstractHilbert,
        operators: list[Array] | Array = [],
        acting_on: list[int] | list[list[int]] = [],
        constant: SupportsComplex | SupportsFloat = 0,
        dtype: DType | None = None,
        *,
        mel_cutoff: float = 1.0e-10,
    ):
        r"""
        Constructs a new ``LocalOperator`` given a hilbert space and (if
        specified) a constant level shift.

        Args:
           hilbert: Hilbert space the operator acts on.
           operators: A list of operators, in matrix form. Supports numpy dense or scipy
           sparse format
           acting_on: A list of list of sites, which the corresponding operators act on. This
                should be constructed such that :code:`operators[i]` acts on the sites :code:`acting_on[i]`.
                If operators is not a list of operators, acting_on should just be the list of
                corresponding sites.
           constant: Constant diagonal shift of the operator, equivalent to
                :math:`+\text{c}\hat{I}`. Default is 0.0.
           dtype: The datatype to use for the matrix elements. Defaults to double precision if
                available.
           mel_cutoff (float): a cutoff to remove small matrix elements (default = 1e-10)

        Examples:
           Constructs a ``LocalOperator`` without any operators.

           >>> from netket.hilbert import Spin
           >>> from netket.operator import LocalOperator
           >>> hi = Spin(0.5)**20
           >>> empty_hat = LocalOperator(hi)
           >>> print(len(empty_hat.acting_on))
           0
        """
        super().__init__(hilbert)
        self.mel_cutoff = mel_cutoff
        self._initialized = None
        self._is_hermitian = None

        # Canonicalize input. From now on input is guaranteed to be in canonical order
        operators, acting_on, dtype = canonicalize_input(
            self.hilbert, operators, acting_on, constant, dtype=dtype
        )
        self._dtype = dtype
        self._constant = np.array(constant, dtype=dtype)

        self._operators_dict = {}
        for op, aon in zip(operators, acting_on):
            self._add_operator(aon, op)

    def _add_operator(self, acting_on: tuple, operator: Array):
        """
        Adds an operator acting on a subset of sites.

        Does not modify in-place the operators themselves which are treated as
        immutables.
        """
        assert isinstance(acting_on, tuple)
        # acting_on_key = tuple(acting_on)
        if acting_on in self._operators_dict:
            operator = self._operators_dict[acting_on] + operator

        self._operators_dict[acting_on] = operator

    @property
    def operators(self) -> list[np.ndarray]:
        """List of the matrices of the operators encoded in this Local Operator.
        Returns a copy.
        """
        return list(self._operators_dict.values())

    @property
    def _operators(self) -> list[np.ndarray]:
        return self.operators

    @property
    def acting_on(self) -> list[list[int]]:
        """List containing the list of the sites on which every operator acts.

        Every operator `self.operators[i]` acts on the sites `self.acting_on[i]`
        """
        return list(self._operators_dict.keys())

    @property
    def n_operators(self) -> int:
        """The total number of operators that were summed upon to build this operator.

        This excludes an optional identity term that is tracked by
        :attr:`netket.operator.LocalOperator.constant`.
        """
        return len(self._operators_dict)

    @property
    def dtype(self) -> DType:
        """DType of the matrix elements of this operator."""
        return self._dtype

    @property
    def size(self) -> int:
        return self._size

    @property
    def is_hermitian(self) -> bool:
        """Returns true if this operator is hermitian."""
        # TODO: (VolodyaCO) I guess that if we have an operator with diagonal elements equal to 1j*C+Y, some complex constant, and
        # self._constant=-1j*C, then the actual diagonal would be Y. How do we check hermiticity taking into account the diagonal
        # elements as well as the self._constant? For the moment I just check hermiticity of the added constant, which must be real.
        if self._is_hermitian is None:
            self._is_hermitian = all(map(is_hermitian, self.operators)) and np.isreal(
                self._constant
            )

        return self._is_hermitian

    @property
    def mel_cutoff(self) -> float:
        r"""float: The cutoff for matrix elements.
        Only matrix elements such that abs(O(i,i))>mel_cutoff
        are considered"""
        return self._mel_cutoff

    @mel_cutoff.setter
    def mel_cutoff(self, mel_cutoff):
        self._mel_cutoff = mel_cutoff
        assert self.mel_cutoff >= 0

    @property
    def constant(self) -> numbers.Number:
        """
        A constant multiplying the identity added to the diagonal of the operator.
        """
        return self._constant

    def to_pauli_strings(self, **kwargs) -> "PauliStrings":  # noqa: F821
        """Convert to PauliStrings object"""
        return local_operators_to_pauli_strings(
            self.hilbert,
            self.operators,
            self.acting_on,
            self.constant,
            self.dtype,
            **kwargs,
        )

    def copy(self, *, dtype: DType | None = None, _cls=None):
        """Returns a copy of the operator, while optionally changing the dtype
        of the operator.

        Args:
            dtype: optional dtype

        Internal args:
            _cls: used to specify the target class
        """

        if dtype is None:
            dtype = self.dtype

        if not np.can_cast(self.dtype, dtype, casting="same_kind"):
            raise ValueError(f"Cannot cast {self.dtype} to {dtype}")

        if _cls is None:
            _cls = type(self)

        new = _cls(self.hilbert, constant=self.constant, dtype=dtype)
        new.mel_cutoff = self.mel_cutoff

        if dtype == self.dtype:
            new._operators_dict = self._operators_dict.copy()
        else:
            new._operators_dict = {
                aon: cast_operator_matrix_dtype(op, dtype)
                for aon, op in self._operators_dict.items()
            }

        return new

    def transpose(self, *, concrete=False):
        r"""LocalOperator: Returns the transpose of this operator."""
        if concrete:
            new = self.copy()
            for aon in new._operators_dict.keys():
                new._operators_dict[aon] = new._operators_dict[aon].transpose()
            return new
        else:
            return Transpose(self)

    def conjugate(self, *, concrete=False):
        r"""LocalOperator: Returns the complex conjugate of this operator."""
        new = self.copy()
        for aon in new._operators_dict.keys():
            new._operators_dict[aon] = new._operators_dict[aon].copy().conjugate()
        return new

    def __radd__(self, other):
        if is_scalar(other):
            return self.__add__(other)
        return super().__radd__(other)

    def __sub__(self, other):
        return self + (-other)

    def __rsub__(self, other):
        return other + (-self)

    def __isub__(self, other):
        return self.__iadd__(-other)

    def __neg__(self):
        return -1 * self

    def __add__(self, other: Union["LocalOperatorBase", numbers.Number]):
        if isinstance(other, LocalOperatorBase) or is_scalar(other):
            op = self.copy(dtype=jnp.promote_types(self.dtype, _dtype(other)))
            op = op.__iadd__(other)
            return op
        return super().__add__(other)

    def __iadd__(self, other):
        if isinstance(other, LocalOperatorBase):
            if self.hilbert != other.hilbert:
                return NotImplemented

            if not np.can_cast(other.dtype, self.dtype, casting="same_kind"):
                raise ValueError(
                    f"Cannot add inplace operator with dtype {other.dtype} "
                    f"to operator with dtype {self.dtype}"
                )

            self.mel_cutoff = min(other.mel_cutoff, self.mel_cutoff)
            self._constant += other.constant.item()
            for aon, op in other._operators_dict.items():
                self._add_operator(aon, op)

            self._reset_caches()
            return self
        if is_scalar(other):
            if not np.can_cast(type(other), self.dtype, casting="same_kind"):
                raise ValueError(
                    f"Cannot add inplace operator with dtype {type(other)} "
                    f"to operator with dtype {self.dtype}"
                )

            self._reset_caches()
            self._constant += other
            return self
        return NotImplemented

    def __truediv__(self, other):
        if not is_scalar(other):
            raise TypeError("Only division by a scalar number is supported.")

        if other == 0:
            raise ValueError("Dividing by 0")
        return self.__mul__(1.0 / other)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __mul__(self, other):
        if isinstance(other, DiscreteOperator):
            # TODO: Deprecated in September 2025
            warnings.warn(OperatorMultiplicationDeprecationWarning())
            op = self.copy(dtype=jnp.promote_types(self.dtype, _dtype(other)))
            return op.__imatmul__(other)
        elif is_scalar(other):
            op = self.copy(dtype=jnp.promote_types(self.dtype, _dtype(other)))
            return op.__imul__(other)
        return NotImplemented

    def __imul__(self, other):
        if isinstance(other, DiscreteOperator):
            return self.__imatmul__(other)
        elif is_scalar(other):
            if not np.can_cast(_dtype(other), self.dtype, casting="same_kind"):
                raise ValueError(
                    f"Cannot add inplace operator of type {type(other)} and "
                    f"dtype {_dtype(other)} to operator with dtype {self.dtype}"
                )
            other = np.asarray(
                other, dtype=jnp.promote_types(self.dtype, _dtype(other))
            )

            self._constant *= other
            if np.abs(other) <= self.mel_cutoff:
                self._operators_dict = {}
            else:
                for key in self._operators_dict:
                    self._operators_dict[key] = other * self._operators_dict[key]

            self._reset_caches()
            return self

        return NotImplemented

    def __imatmul__(self, other):
        if not isinstance(other, LocalOperatorBase):
            return NotImplemented

        if not np.can_cast(other.dtype, self.dtype, casting="same_kind"):
            raise ValueError(
                f"Cannot add inplace operator with dtype {type(other)} to operator with dtype {self.dtype}"
            )

        return self._op_imatmul_(other)

    def _op__matmul__(self, other: "LocalOperatorBase") -> "LocalOperatorBase":
        if not isinstance(other, LocalOperatorBase):
            return NotImplemented
        op = self.copy(dtype=jnp.promote_types(self.dtype, _dtype(other)))
        return op._op_imatmul_(other)

    def _op_imatmul_(self, other: "LocalOperatorBase") -> "LocalOperatorBase":
        if not isinstance(other, LocalOperatorBase):
            return NotImplemented

        self.mel_cutoff = min(other.mel_cutoff, self.mel_cutoff)

        # (α + ∑ᵢAᵢ)(β + ∑ᵢBᵢ) =
        # = αβ + α ∑ᵢBᵢ + β ∑ᵢAᵢ + ∑ᵢⱼAᵢBⱼ
        # = β(α + ∑ᵢAᵢ) + α ∑ᵢBᵢ + ∑ᵢⱼAᵢBⱼ

        α = self.constant.item()
        β = other.constant.item()
        # copy A dict because it is modified inplace in __imul__(β) and add_operators
        A_op_dict = self._operators_dict.copy()
        B_op_dict = other._operators_dict

        # αβ + β ∑ᵢAᵢ
        self.__imul__(β)

        # α ∑ᵢBᵢ
        if np.abs(α) > self.mel_cutoff:
            for aon, op in B_op_dict.items():
                self._add_operator(aon, α * op)

        # ∑ᵢⱼAᵢBⱼ
        for supp_A_i, A_i in A_op_dict.items():
            for supp_B_j, B_j in B_op_dict.items():
                self._add_operator(
                    *_multiply_operators(
                        self.hilbert, supp_A_i, A_i, supp_B_j, B_j, dtype=self.dtype
                    )
                )

        is_hermitian_A = self._is_hermitian
        is_hermitian_B = other._is_hermitian

        self._reset_caches()

        if is_hermitian_A and is_hermitian_B:
            self._is_hermitian = is_hermitian_A

        return self

    def _reset_caches(self):
        """
        Cleans the internal caches built on the operator.
        """
        self._initialized = False
        self._is_hermitian = None

    @property
    def max_conn_size(self) -> int:
        """The maximum number of non zero ⟨x|O|x'⟩ for every x."""
        self._setup()
        return self._max_conn_size

    def __repr__(self):
        ao = self.acting_on

        acting_str = f"acting_on={ao}"
        if len(acting_str) > 55:
            acting_str = f"#acting_on={len(ao)} locations"
        return f"{type(self).__name__}(dim={self.hilbert.size}, {acting_str}, constant={self.constant}, dtype={self.dtype})"
