# Copyright (c) 2026 The neuraLQX 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.


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

"""
NOTE: part(s) of, or the entire content, of this file is obtained from NetKet's source code
      the original copyright mentioned above applies.
"""

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

from netket.operator import PauliStrings

_pauli_basis = None
pauli_basis_str = list("IXYZ")


def _get_pauli_basis():
    global _pauli_basis
    if _pauli_basis is None:
        # Pauli Matrices: shape (2, 2)
        I = jnp.eye(2)
        X = jnp.array([[0, 1], [1, 0]])
        Y = jnp.array([[0, -1j], [1j, 0]])
        Z = jnp.array([[1, 0], [0, -1]])

        # stacked pauli matrices: shape (4, 2, 2)
        _pauli_basis = jnp.stack([I, X, Y, Z], axis=0)
    return _pauli_basis


@jax.jit
def _tensor_product(mats1, mats2):
    """Get all combinations of Pauli operators between two sets
    via tensorproduct.

    Args:
        mats1: either a (2ⁿ,2ⁿ) representing a single operator or a (N,2ⁿ,2ⁿ)
            tensor encoding N different pauli operators
        mats2: either a (2ᵐ,2ᵐ) representing a single operator or a (M,2ᵐ,2ᵐ)
            tensor encoding M different pauli operators
    Returns
        A (2ⁿ⁺ᵐ, 2ⁿ⁺ᵐ) tensor if both mats1 and mats2 are single paulis. Otherwise
        an (N,M, 2ⁿ⁺ᵐ, 2ⁿ⁺ᵐ) tensor (where N or M dimension might be dropped if
        not there in the input).
    """
    if mats1.ndim == 2 and mats2.ndim == 2:
        return jnp.kron(mats1, mats2)
    elif mats1.ndim == 3:
        return jax.vmap(_tensor_product, in_axes=(0, None))(mats1, mats2)
    elif mats2.ndim == 3:
        return jax.vmap(_tensor_product, in_axes=(None, 0))(mats1, mats2)
    else:
        raise ValueError("Not supposed to get more than 3 dimensions")


def _get_basis_till_n(n=1):
    """Get the basis set till the basis set with at most n non
    identities in the Pauli Strings, which will contain
    4ⁿ elements of size (2ⁿ, 2ⁿ).

    Returns:
        A dictionary with keys [1, 2, ..., n] where every value is
        a basis set encoded into a tensor of shape (4ⁿ, 2ⁿ, 2ⁿ)
        where the first dimension (4ⁿ) is the number of basis elements
        for that number of non-identity paulis, and 2ⁿ is the size
        of those matrices.

        The key represents the size of those matrices.
    """
    if n == 1:
        return {n: _get_pauli_basis()}
    else:
        lower_bases = _get_basis_till_n(n=n - 1)

        # extract the basis for n-1, which is (4ⁿ⁻¹, 2ⁿ, 2ⁿ) tensor
        # compute the (4,4ⁿ⁻¹, 4, 4)
        tp = _tensor_product(_get_pauli_basis(), lower_bases[n - 1])

        # reshape everything to a single set of bases (4ⁿ, ...)
        tp = tp.reshape(-1, *tp.shape[-2:])
        return {**lower_bases, n: tp}


def _make_basis_strings(n_acting):
    """Make the strings corresponding to the tensorproduct of Pauli operators"""
    bs = list(itertools.product(pauli_basis_str, repeat=n_acting))
    bs = list(map(lambda b: "".join(b), bs))
    return np.array(bs)


def _get_paulistrings_till_n(n=1):
    """
    Get the
    """
    return {n_paulis: _make_basis_strings(n_paulis) for n_paulis in range(1, n + 1)}


@jax.jit
def _hilbert_schmidt(A, basis_element):
    """Hilbert schmidt product where the basis elements are hermitian"""
    return jnp.trace(basis_element @ A)


@jax.jit
def _project_matrix(A, basis_set):
    """
    Project matrix onto basis of tensorproduct of Paulis
    according to hilbert-schmidt product
    """
    assert A.ndim == 2
    assert A.shape[0] == A.shape[1]
    size = A.shape[0]
    return jax.vmap(_hilbert_schmidt, in_axes=(None, 0))(A, basis_set) / size


def _create_pauli_string(substr, acting_on, hilbert_size):
    """Create a single PauliString from the operators and indices"""
    s = list("I" * hilbert_size)
    for i, si in zip(acting_on, substr):
        s[i] = si
    s = "".join(list(s))
    return s


def _local_operator_to_pauli_strings(
    op, acting_on, basis_set, basis_set_str, hilbert_size
):
    """Convert a matrix into a sum of Pauli strings"""
    acting_on = np.array(acting_on)
    all_weights = _project_matrix(op, basis_set)
    nz_idxs = np.nonzero(all_weights)[0]
    weights = all_weights[nz_idxs]

    # the next is not efficient
    operators = [basis_set_str[idx] for idx in nz_idxs]
    _str_maker = lambda s: _create_pauli_string(s, acting_on, hilbert_size)
    operators = list(map(_str_maker, operators))
    return list(operators), list(weights)


def _convert_to_dense(m):
    return m.todense() if issparse(m) else m


def local_operators_to_pauli_strings(
    hilbert, operators, acting_on, constant, dtype, cls=PauliStrings
):
    """Convert a LocalOperator into PauliStrings

    Args:
        local_op: LocalOperator

    Returns:
        PauliStrings
    """

    if any(d != 2 for d in hilbert.shape):
        raise TypeError(
            "Cannot convert to Pauli strings operators defined on hilbert spaces "
            "with local dimension != 2"
        )

    pauli_strings = []
    weights = []

    if len(operators) > 0:
        # maximum number of non-identity operators
        n_max = max(list(map(len, acting_on)))

        # make the basis and the string representation
        bases = _get_basis_till_n(n=n_max)
        bases_strs = _get_paulistrings_till_n(n=n_max)

        # loop trough operators and convert to Pauli strings by projecting
        for op, act in zip(operators, acting_on):
            # convert to dense-NumPy as we later use jax which does not support sparse.
            op = _convert_to_dense(op)

            n_paulis = len(act)

            op_list, w_list = _local_operator_to_pauli_strings(
                op, act, bases[n_paulis], bases_strs[n_paulis], hilbert.size
            )
            pauli_strings += op_list
            weights += w_list

    # add the constant if it's not zero
    if np.abs(constant) > 1e-13:
        pauli_strings.append("I" * hilbert.size)
        weights.append(constant)

    # the calculation above returns complex weights even for operators that are
    # purely real, so we discard their imaginary parts
    weights = [x.real if np.isreal(x) else x for x in weights]

    res = cls(hilbert, operators=pauli_strings, weights=weights, dtype=dtype)
    return res
