# Copyright 2023-2024 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.


# this file contains a more-or-less 1:1 port of the numba localoperator to jax, using padding where necessary

from functools import partial
from typing import TYPE_CHECKING

import jax
import jax.numpy as jnp
from jax.tree_util import register_pytree_node_class

from netket.errors import JaxOperatorNotConvertibleToNumba

from .base import LocalOperatorBase
from .compile_helpers import pack_internals_jax

from .._pauli_strings import PauliStringsJax
from .._discrete_operator_jax import DiscreteJaxOperator

if TYPE_CHECKING:
    from .numba import LocalOperatorNumba


@partial(jax.vmap, in_axes=(0, None))  # samples
@partial(jax.vmap, in_axes=(0, 0))  # operators
def _state_to_number(x_i, basis_i):
    # convert array of local states to number
    # in the hilbert space of all the sites the operator is acting on

    # number of sites these operators are acting on
    acting_size_i = basis_i.shape[-1]

    @partial(jax.vmap, in_axes=(None, None, None, 0))  # vmap over local sites (k)
    def _f(acting_size_i, x_i, basis_i, k):
        # special case for empty operator
        if len(basis_i) == 0:
            return jnp.zeros((), dtype=basis_i.dtype)
        return basis_i[k] * x_i[acting_size_i - k - 1]

    # vmap over sites we are acting on, computing the contribution to the total int,
    # then taking the sum
    return _f(acting_size_i, x_i, basis_i, jnp.arange(acting_size_i)).sum()


@partial(jax.vmap, in_axes=(0, 0, None))  # Ns
@partial(jax.vmap, in_axes=(None, 0, 0))  # rows
@partial(jax.vmap, in_axes=(None, 0, None))  # ncmax
def _set_at(x, new_x_ao, acting_on):
    return x.at[acting_on].set(new_x_ao)


@partial(jax.vmap, in_axes=(None, 0))  # Ns
@partial(jax.vmap, in_axes=(0, 0))  # rows
def _index_at(diag_mels, i):
    # special case for empty operator
    if len(diag_mels) == 0:
        return jnp.zeros((), dtype=diag_mels.dtype)
    return diag_mels[i]


@partial(jax.jit, static_argnums=(0, 1))
def _local_operator_kernel_jax(nonzero_diagonal, max_conn_size, mel_cutoff, op_args, x):
    assert x.ndim == 2

    # in the foollowing variables with a trailing underscore are per group
    # of operators with the same number of sites they act on

    (
        acting_on_,
        n_conns_,
        diag_mels_,
        x_prime_,
        mels_,
        basis_,
        constant,
    ) = op_args

    n_groups = len(acting_on_)
    ncmax_ = [m.shape[2] for m in mels_]

    ###
    # determine the row x corresponds to for each of the operators
    #
    i_row_ = []
    # iterate over all groups of operators acting on a certain number of sites
    for k in range(n_groups):
        # extract the local states of all sites the operators are acting on
        # and compute the corresponding number / row index
        i_row_.append(_state_to_number(x[:, acting_on_[k]], basis_[k]))
    ###
    # compute the number of connected elements
    #
    # extract the number of (nonzero) off-diagonal elements of the rows
    n_conn_offdiag_ = tuple(map(_index_at, n_conns_, i_row_))
    # sum for each group of operators acting on a certain number of sites
    # and sum over all gropus to get the total
    n_conn_offdiag = sum([n.sum(axis=-1) for n in n_conn_offdiag_])
    n_conn_diag = 1 if nonzero_diagonal else 0
    n_conn_total = n_conn_diag + n_conn_offdiag

    mels_offdiag_ = []
    xp_offdiag_ = []

    if nonzero_diagonal:
        xp_diag = x
        # extract diagonal mels of the rows
        mels_diag_ = tuple(map(_index_at, diag_mels_, i_row_))
        # sum over operators
        mels_diag = constant + sum([m.sum(axis=-1) for m in mels_diag_])
    else:
        # zero diagonal
        mels_diag = None
        xp_diag = None

    if mels_diag is not None:
        mels_offdiag_.append(mels_diag)
        xp_offdiag_.append(xp_diag)

    # iterate over all groups of operators acting on a certain number of sites
    for k in range(n_groups):
        # indices used to select rows
        i = i_row_[k]
        a = jnp.arange(i.shape[1])

        # mask for the connected elements
        # (not all operators have the same number nonzeros per row, and the
        # arrays we use here are padded to the max within each group)
        # this mask is False whenever it's just padding
        conn_maskall = jnp.arange(ncmax_[k])[None, None] < n_conns_[k][a, i][:, :, None]

        # set mels in the padding to 0
        # TODO check this is necessary (shouldn't it already be 0 ???)
        mels_offdiag = mels_[k][a, i] * conn_maskall
        #
        # compute xp
        #
        # we start from the xp for the sites the operator is acting on
        new = x_prime_[k][a, i].astype(x.dtype)  # (Ns, terms, ncmax, nsitesactiongon)
        acting_on = acting_on_[k]
        old = x[:, acting_on]  # (Ns, terms, nsitesactiongon)
        old = jnp.broadcast_to(old[:, :, None, :], new.shape)
        mask = jnp.broadcast_to(conn_maskall[:, :, :, None], new.shape)
        # select it only if we are not padding, otherwise keep old
        new_x_ao = jax.lax.select(mask, new, old)
        # now insert the local states into the full x
        xp_offdiag = _set_at(x, new_x_ao, acting_on)

        mels_offdiag_.append(mels_offdiag)
        xp_offdiag_.append(xp_offdiag)

    if max_conn_size is not None:
        # compute a mask which tells us where the actual mels are and
        # where there is just padding
        # (we could also just be lazy and check the mels for being 0,
        # but allows us to make the order of mels consistent with the numba op)
        mask_ = []
        if mels_diag is not None:
            mask_.append(jnp.full(mels_diag.shape, fill_value=True))
        # iterate over all groups of operators acting on a certain number of sites
        for k in range(n_groups):
            mask = jnp.arange(ncmax_[k])[None, None, :] < n_conn_offdiag_[k][:, :, None]
            mask_.append(mask)

        # pad with old state and mel 0
        xp_offdiag_.append(x[:, None][:, :1])
        mels_offdiag_.append(jnp.zeros(x.shape[:-1]))
        mask_.append(jnp.full((x.shape[0], 1), fill_value=False))

    mels = jnp.hstack([m.reshape(m.shape[0], -1) for m in mels_offdiag_])
    xp = jnp.hstack([x.reshape(x.shape[0], -1, x.shape[-1]) for x in xp_offdiag_])

    # TODO run unique on it, there might be repeated xps

    if max_conn_size is None:
        return xp, mels, n_conn_total
    else:
        if mel_cutoff is not None:
            mask = jnp.abs(mels) > mel_cutoff
            n_conn_total = mask.sum(axis=-1)
        else:
            mask = jnp.hstack([m.reshape(m.shape[0], -1) for m in mask_])
        # move nonzero mels to the front and keep exactly max_conn_size
        (ind,) = jax.vmap(partial(jnp.where, size=max_conn_size, fill_value=-1))(mask)

        xp = jax.vmap(partial(jnp.take, axis=0))(xp, ind)
        mels = jax.vmap(partial(jnp.take, axis=0))(mels, ind)

        return xp, mels, n_conn_total


@register_pytree_node_class
class LocalOperatorJax(LocalOperatorBase, DiscreteJaxOperator):
    """Jax 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.DiscreteJaxOperator.max_conn_size`.

    For the numba-based implementation, look at
    :class:`netket.operator.LocalOperatorNumba`.
    """

    _convertible: bool = True
    """Internal flag. True if it can be converted to numba, false
    otherwise."""

    def _setup(self, force=False):
        if force or not self._initialized:
            data = pack_internals_jax(
                self.hilbert,
                self._operators_dict,
                self.constant,
                self.dtype,
                self.mel_cutoff,
            )

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

            self._initialized = True

    def _get_conn_padded(self, x):
        self._setup()

        shape = x.shape
        x = x.reshape(-1, x.shape[-1])
        x_ids = self.hilbert.states_to_local_indices(x)

        xp_ids, mels, n_conn = _local_operator_kernel_jax(
            self._nonzero_diagonal,
            self._max_conn_size,
            self._mel_cutoff,
            (
                self._acting_on,
                self._n_conns,
                self._diag_mels,
                self._x_prime,
                self._mels,
                self._basis,
                self._constant,
            ),
            x_ids,
        )

        xp = self.hilbert.local_indices_to_states(xp_ids, dtype=x.dtype)
        xp = xp.reshape(shape[:-1] + xp.shape[-2:])
        mels = mels.reshape(shape[:-1] + mels.shape[-1:])
        n_conn = n_conn.reshape(shape[:-1])

        return xp, mels, n_conn

    def get_conn_padded(self, x):
        xp, mels, _ = self._get_conn_padded(x)
        return xp, mels

    def n_conn(self, x, out=None):
        if out is not None:
            raise NotImplementedError()
        _, _, n_conn = self._get_conn_padded(x)
        return n_conn

    def tree_flatten(self):
        self._setup()
        data = (
            self._acting_on,
            self._n_conns,
            self._diag_mels,
            self._x_prime,
            self._mels,
            self._basis,
            self._constant,
        )
        metadata = {
            "hilbert": self.hilbert,
            "dtype": self.dtype,
            "nonzero_diagonal": self._nonzero_diagonal,
            "max_conn_size": self._max_conn_size,
        }
        return data, metadata

    @classmethod
    def tree_unflatten(cls, metadata, data):
        hi = metadata["hilbert"]
        dtype = metadata["dtype"]

        op = cls(hi, dtype=dtype)

        op._nonzero_diagonal = metadata["nonzero_diagonal"]
        op._max_conn_size = metadata["max_conn_size"]
        (
            op._acting_on,
            op._n_conns,
            op._diag_mels,
            op._x_prime,
            op._mels,
            op._basis,
            op._constant,
        ) = data

        op._initialized = True
        op._convertible = False
        return op

    def to_numba_operator(self) -> "LocalOperatorNumba":  # noqa: F821
        """
        Returns the standard numba version of this operator, which is an
        instance of :class:`netket.operator.LocalOperator`.
        """
        from .numba import LocalOperatorNumba

        if self._convertible is False:
            raise JaxOperatorNotConvertibleToNumba(self)

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

    def to_pauli_strings(self) -> "PauliStringsJax":  # noqa: F821
        """Convert to PauliStrings object"""
        return super().to_pauli_strings(
            cls=PauliStringsJax,
        )
