# 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 TYPE_CHECKING

import numpy as np
import numba

from netket.errors import concrete_or_error, NumbaOperatorGetConnDuringTracingError


from .compile_helpers import pack_internals
from .base import LocalOperatorBase

if TYPE_CHECKING:
    from .jax import LocalOperatorJax


class LocalOperatorNumba(LocalOperatorBase):
    """Numba 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.

        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.

    .. warning::

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

    For the jax-based implementation, look at
    :class:`netket.operator.LocalOperatorJax`.
    """

    __module__ = "netket.operator"

    def _setup(self, force: bool = False):
        """Analyze the operator strings and precompute arrays for get_conn inference"""
        if force or not self._initialized:
            data = pack_internals(
                self.hilbert,
                self._operators_dict,
                self.constant,
                self.dtype,
                self.mel_cutoff,
            )

            self._acting_on = data["acting_on"]
            self._acting_size = data["acting_size"]
            self._diag_mels = data["diag_mels"]
            self._mels = data["mels"]
            self._x_prime = data["x_prime"]
            self._n_conns = data["n_conns"]
            self._basis = data["basis"]
            self._nonzero_diagonal = data["nonzero_diagonal"]
            self._max_conn_size = data["max_conn_size"]
            self._initialized = True

    def get_conn_flattened(self, x, sections, pad=False):
        r"""Finds the connected elements of the Operator. Starting
        from a given quantum number x, it finds all other quantum numbers x' such
        that the matrix element :math:`O(x,x')` is different from zero. In general there
        will be several different connected states x' satisfying this
        condition, and they are denoted here :math:`x'(k)`, for :math:`k=0,1...N_{\mathrm{connected}}`.

        This is a batched version, where x is a matrix of shape (batch_size,hilbert.size).

        Args:
            x (matrix): A matrix of shape (batch_size,hilbert.size) containing
                        the batch of quantum numbers x.
            sections (array): An array of size (batch_size) useful to unflatten
                        the output of this function.
                        See numpy.split for the meaning of sections.
            pad (bool): Whether to use zero-valued matrix elements in order to return all equal sections.

        Returns:
            matrix: The connected states x', flattened together in a single matrix.
            array: An array containing the matrix elements :math:`O(x,x')` associated to each x'.

        """
        self._setup()

        x_ids = self.hilbert.states_to_local_indices(x)

        x_ids = concrete_or_error(
            np.asarray,
            x_ids,
            NumbaOperatorGetConnDuringTracingError,
            self,
        )

        xp_ids, mels = self._get_conn_flattened_kernel(
            x_ids,
            sections,
            self._basis,
            self._constant,
            self._diag_mels,
            self._n_conns,
            self._mels,
            self._x_prime,
            self._acting_on,
            self._acting_size,
            self._nonzero_diagonal,
            pad,
        )

        xp = self.hilbert.local_indices_to_states(xp_ids, dtype=x.dtype)
        return xp, mels

    @staticmethod
    @numba.jit(nopython=True)
    def _get_conn_flattened_kernel(
        x,
        sections,
        basis,
        constant,
        diag_mels,
        n_conns,
        all_mels,
        all_x_prime,
        acting_on,
        acting_size,
        nonzero_diagonal,
        pad=False,
    ):
        batch_size = x.shape[0]
        n_sites = x.shape[1]
        dtype = all_mels.dtype

        # TODO remove this line when numba 0.53 is dropped 0.54 is minimum version
        # workaround a bug in Numba arising when NUMBA_BOUNDSCHECK=1
        constant = constant.item()

        assert sections.shape[0] == batch_size

        n_operators = n_conns.shape[0]
        # array to store the row index
        xs_n = np.empty((batch_size, n_operators), dtype=np.intp)

        tot_conn = 0
        max_conn = 0

        for b in range(batch_size):
            # diagonal element
            conn_b = 1 if nonzero_diagonal else 0

            # counting the off-diagonal elements
            for i in range(n_operators):
                acting_size_i = acting_size[i]
                # compute the number (row index) from the local states
                # (here we do the inverse of _number_to_state from
                # compile_helpers.py, so this is essentially _state_to_number)
                xs_n[b, i] = 0
                x_b = x[b]
                x_i = x_b[acting_on[i, :acting_size_i]]
                # iterate over sites the current operator is acting on
                for k in range(acting_size_i):
                    # compute
                    xs_n[b, i] += x_i[acting_size_i - k - 1] * basis[i, k]

                # sum the number of off-diagonal connected elements
                conn_b += n_conns[i, xs_n[b, i]]

            tot_conn += conn_b
            sections[b] = tot_conn

            if pad:
                max_conn = max(conn_b, max_conn)

        if pad:
            tot_conn = batch_size * max_conn

        x_prime = np.empty((tot_conn, n_sites), dtype=x.dtype)
        mels = np.empty(tot_conn, dtype=dtype)

        c = 0
        for b in range(batch_size):
            c_diag = c
            x_batch = x[b]

            if nonzero_diagonal:
                mels[c_diag] = constant
                x_prime[c_diag] = np.copy(x_batch)
                c += 1

            for i in range(n_operators):
                if nonzero_diagonal:
                    mels[c_diag] += diag_mels[i, xs_n[b, i]]

                # get the number of connected elements for the current operator
                # at the rows index corresponding to the state of x at the sites
                # the operator is acting on
                n_conn_i = n_conns[i, xs_n[b, i]]

                if n_conn_i > 0:
                    sites = acting_on[i]
                    acting_size_i = acting_size[i]

                    for cc in range(n_conn_i):  # iterate over compressed nonzero cols
                        # get the nonzero mels of the current row
                        mels[c + cc] = all_mels[i, xs_n[b, i], cc]
                        x_prime[c + cc] = np.copy(x_batch)
                        # set the changed local states of the sites the operator
                        # is acting on
                        # it is stored in all_x_prime, where we select the row
                        for k in range(acting_size_i):
                            x_prime[c + cc, sites[k]] = all_x_prime[
                                i, xs_n[b, i], cc, k
                            ]
                    c += n_conn_i

            if pad:
                delta_conn = max_conn - (c - c_diag)
                mels[c : c + delta_conn].fill(0)
                x_prime[c : c + delta_conn, :] = np.copy(x_batch)
                c += delta_conn
                sections[b] = c

        return x_prime, mels

    def get_conn_filtered(self, x, sections, filters):
        r"""Finds the connected elements of the Operator using only a subset of operators. Starting
        from a given quantum number x, it finds all other quantum numbers x' such
        that the matrix element :math:`O(x,x')` is different from zero. In general there
        will be several different connected states x' satisfying this
        condition, and they are denoted here :math:`x'(k)`, for :math:`k=0,1...N_{\mathrm{connected}}`.

        This is a batched version, where x is a matrix of shape (batch_size,hilbert.size).

        Args:
            x (matrix): A matrix of shape (batch_size,hilbert.size) containing
                        the batch of quantum numbers x.
            sections (array): An array of size (batch_size) useful to unflatten
                        the output of this function.
                        See numpy.split for the meaning of sections.
            filters (array): Only operators op(filters[i]) are used to find the connected elements of
                        x[i].

        Returns:
            matrix: The connected states x', flattened together in a single matrix.
            array: An array containing the matrix elements :math:`O(x,x')` associated to each x'.

        """
        self._setup()

        x_ids = self.hilbert.states_to_local_indices(x)
        x_ids = concrete_or_error(
            np.asarray,
            x_ids,
            NumbaOperatorGetConnDuringTracingError,
            self,
        )

        xp_ids, mels = self._get_conn_filtered_kernel(
            x_ids,
            sections,
            self._basis,
            self._constant,
            self._diag_mels,
            self._n_conns,
            self._mels,
            self._x_prime,
            self._acting_on,
            self._acting_size,
            filters,
        )
        xp = self.hilbert.local_indices_to_states(xp_ids, dtype=x.dtype)
        return xp, mels

    @staticmethod
    @numba.jit(nopython=True)
    def _get_conn_filtered_kernel(
        x,
        sections,
        basis,
        constant,
        diag_mels,
        n_conns,
        all_mels,
        all_x_prime,
        acting_on,
        acting_size,
        filters,
    ):
        batch_size = x.shape[0]
        n_sites = x.shape[1]
        dtype = all_mels.dtype

        assert filters.shape[0] == batch_size and sections.shape[0] == batch_size

        # TODO remove this line when numba 0.53 is dropped 0.54 is minimum version
        # workaround a bug in Numba arising when NUMBA_BOUNDSCHECK=1
        constant = constant.item()

        n_operators = n_conns.shape[0]
        xs_n = np.empty((batch_size, n_operators), dtype=np.intp)

        tot_conn = 0

        for b in range(batch_size):
            # diagonal element
            tot_conn += 1

            # counting the off-diagonal elements
            i = filters[b]

            assert i < n_operators and i >= 0
            acting_size_i = acting_size[i]

            xs_n[b, i] = 0
            x_b = x[b]
            x_i = x_b[acting_on[i, :acting_size_i]]
            for k in range(acting_size_i):
                xs_n[b, i] += x_i[acting_size_i - k - 1] * basis[i, k]

            tot_conn += n_conns[i, xs_n[b, i]]
            sections[b] = tot_conn

        x_prime = np.empty((tot_conn, n_sites))
        mels = np.empty(tot_conn, dtype=dtype)

        c = 0
        for b in range(batch_size):
            c_diag = c
            mels[c_diag] = constant
            x_batch = x[b]
            x_prime[c_diag] = np.copy(x_batch)
            c += 1

            i = filters[b]
            # Diagonal part
            mels[c_diag] += diag_mels[i, xs_n[b, i]]
            n_conn_i = n_conns[i, xs_n[b, i]]

            if n_conn_i > 0:
                sites = acting_on[i]
                acting_size_i = acting_size[i]

                for cc in range(n_conn_i):
                    mels[c + cc] = all_mels[i, xs_n[b, i], cc]
                    x_prime[c + cc] = np.copy(x_batch)

                    for k in range(acting_size_i):
                        x_prime[c + cc, sites[k]] = all_x_prime[i, xs_n[b, i], cc, k]
                c += n_conn_i

        return x_prime, mels

    def to_jax_operator(self) -> "LocalOperatorJax":  # noqa: F821
        """
        Returns the jax-compatible version of this operator, which is an
        instance of :class:`netket.operator.LocalOperatorJax`.
        """
        from .jax import LocalOperatorJax

        return LocalOperatorJax(
            self.hilbert,
            self.operators,
            self.acting_on,
            self.constant,
            dtype=self.dtype,
            mel_cutoff=self.mel_cutoff,
        )


# Keep the old name for backward compatibility
LocalOperator = LocalOperatorNumba
