"""
********************************************************************************
* Copyright (c) 2024 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************
"""

from qrisp.operators.hamiltonian import Hamiltonian
from qrisp.operators.qubit.bound_qubit_term import BoundQubitTerm
from qrisp.operators.qubit.measurement import get_measurement
from qrisp import h, sx, IterationEnvironment, conjugate, merge

import sympy as sp

threshold = 1e-9

#
# BoundQubitOperator
#


class BoundQubitOperator(Hamiltonian):
    r"""
    This class provides an efficient implementation of Pauli Hamiltonians acting on QuantumVariables, i.e.,
    Hamiltonians of the form

    .. math::

        H=\sum\limits_{j}\alpha_jP_j

    where $P_j=\prod_i\sigma_i^j$ is a Pauli product,
    and $\sigma_i^j\in\{I,X,Y,Z\}$ is the Pauli operator acting on qubit $i$.

    Parameters
    ----------
    terms_dict : dict, optional
        A dictionary representing a BoundQubitOperator.

    Examples
    --------

    A BoundQubitOperator can be specified conveniently in terms of ``X``, ``Y``, ``Z`` operators:

    ::

        from qrisp import QuantumVariable
        from qrisp.operators import BoundQubitOperator, X,Y,Z

        qv = QuantumVariable(2)
        H = 1+2*X(qv[0])+3*X(qv[0])*Y(qv[1])

    Yields $1+2X(qv.0)+3X(qv.0)Y(qv.1)$.

    """

    def __init__(self, terms_dict={}):
        self.terms_dict = terms_dict
        pass

    def len(self):
        return len(self.terms_dict)

    #
    # Printing
    #

    def _repr_latex_(self):
        # Convert the sympy expression to LaTeX and return it
        expr = self.to_expr()
        return f"${sp.latex(expr)}$"

    def __str__(self):
        # Convert the sympy expression to a string and return it
        expr = self.to_expr()
        return str(expr)

    def to_expr(self):
        """
        Returns a SymPy expression representing the operator.

        Returns
        -------
        expr : sympy.expr
            A SymPy expression representing the operator.

        """

        expr = 0
        for pauli, coeff in self.terms_dict.items():
            expr += coeff * pauli.to_expr()
        return expr

    #
    # Arithmetic
    #

    def __pow__(self, e):
        if self.len() == 1:
            if isinstance(e, int) and e >= 0:
                if e % 2 == 0:
                    return BoundQubitOperator({BoundQubitTerm(): 1})
                else:
                    return self
            else:
                raise TypeError(
                    "Unsupported operand type(s) for ** or pow(): "
                    + str(type(self))
                    + " and "
                    + str(type(e))
                )
        else:
            raise TypeError(
                "Unsupported operand type(s) for ** or pow(): "
                + str(type(self))
                + " and "
                + str(type(e))
            )

    def __add__(self, other):
        """
        Returns the sum of the operator self and other.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to add to the operator self.

        Returns
        -------
        result : BoundQubitOperator
            The sum of the operator self and other.

        """

        if isinstance(other, (int, float, complex)):
            other = BoundQubitOperator({BoundQubitTerm(): other})
        if not isinstance(other, BoundQubitOperator):
            raise TypeError("Cannot add BoundQubitOperator and " + str(type(other)))

        res_terms_dict = {}

        for pauli, coeff in self.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) + coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        for pauli, coeff in other.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) + coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        result = BoundQubitOperator(res_terms_dict)
        return result

    def __sub__(self, other):
        """
        Returns the difference of the operator self and other.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to substract from the operator self.

        Returns
        -------
        result : BoundQubitOperator
            The difference of the operator self and other.

        """

        if isinstance(other, (int, float, complex)):
            other = BoundQubitOperator({BoundQubitTerm(): other})
        if not isinstance(other, BoundQubitOperator):
            raise TypeError(
                "Cannot substract BoundQubitOperator and " + str(type(other))
            )

        res_terms_dict = {}

        for pauli, coeff in self.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) + coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        for pauli, coeff in other.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) - coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        result = BoundQubitOperator(res_terms_dict)
        return result

    def __rsub__(self, other):
        """
        Returns the difference of the operator other and self.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to substract the operator self from.

        Returns
        -------
        result : BoundQubitOperator
            The difference of the operator other and self.

        """

        if isinstance(other, (int, float, complex)):
            other = BoundQubitOperator({BoundQubitTerm(): other})
        if not isinstance(other, BoundQubitOperator):
            raise TypeError(
                "Cannot substract BoundQubitOperator and " + str(type(other))
            )

        res_terms_dict = {}

        for pauli, coeff in self.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) - coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        for pauli, coeff in other.terms_dict.items():
            res_terms_dict[pauli] = res_terms_dict.get(pauli, 0) + coeff
            if abs(res_terms_dict[pauli]) < threshold:
                del res_terms_dict[pauli]

        result = BoundQubitOperator(res_terms_dict)
        return result

    def __mul__(self, other):
        """
        Returns the product of the operator self and other.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to multiply with the operator self.

        Returns
        -------
        result : BoundQubitOperator
            The product of the operator self and other.

        """

        if isinstance(other, (int, float, complex)):
            other = BoundQubitOperator({BoundQubitTerm(): other})
        if not isinstance(other, BoundQubitOperator):
            raise TypeError(
                "Cannot multipliy BoundQubitOperator and " + str(type(other))
            )

        res_terms_dict = {}

        for pauli1, coeff1 in self.terms_dict.items():
            for pauli2, coeff2 in other.terms_dict.items():
                curr_pauli, curr_coeff = pauli1 * pauli2
                res_terms_dict[curr_pauli] = (
                    res_terms_dict.get(curr_pauli, 0) + curr_coeff * coeff1 * coeff2
                )

        result = BoundQubitOperator(res_terms_dict)
        return result

    __radd__ = __add__
    __rmul__ = __mul__

    #
    # Inplace arithmetic
    #

    def __iadd__(self, other):
        """
        Adds other to the operator self.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to add to the operator self.

        """

        if isinstance(other, (int, float, complex)):
            self.terms_dict[BoundQubitTerm()] = (
                self.terms_dict.get(BoundQubitTerm(), 0) + other
            )
            return self
        if not isinstance(other, BoundQubitOperator):
            raise TypeError("Cannot add BoundQubitOperator and " + str(type(other)))

        for pauli, coeff in other.terms_dict.items():
            self.terms_dict[pauli] = self.terms_dict.get(pauli, 0) + coeff
            if abs(self.terms_dict[pauli]) < threshold:
                del self.terms_dict[pauli]
        return self

    def __isub__(self, other):
        """
        Substracts other from the operator self.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to substract from the operator self.

        """

        if isinstance(other, (int, float, complex)):
            self.terms_dict[BoundQubitTerm()] = (
                self.terms_dict.get(BoundQubitTerm(), 0) - other
            )
            return self
        if not isinstance(other, BoundQubitOperator):
            raise TypeError("Cannot add BoundQubitOperator and " + str(type(other)))

        for pauli, coeff in other.terms_dict.items():
            self.terms_dict[pauli] = self.terms_dict.get(pauli, 0) - coeff
            if abs(self.terms_dict[pauli]) < threshold:
                del self.terms_dict[pauli]
        return self

    def __imul__(self, other):
        """
        Multiplys other to the operator self.

        Parameters
        ----------
        other : int, float, complex or BoundQubitOperator
            A scalar or a BoundQubitOperator to multiply with the operator self.

        """

        if isinstance(other, (int, float, complex)):
            # other = BoundQubitOperator({BoundQubitTerm():other})
            for term in self.terms_dict:
                self.terms_dict[term] *= other
            return self

        if not isinstance(other, BoundQubitOperator):
            raise TypeError(
                "Cannot multipliy BoundQubitOperator and " + str(type(other))
            )

        res_terms_dict = {}

        for pauli1, coeff1 in self.terms_dict.items():
            for pauli2, coeff2 in other.terms_dict.items():
                curr_pauli, curr_coeff = pauli1 * pauli2
                res_terms_dict[curr_pauli] = (
                    res_terms_dict.get(curr_pauli, 0) + curr_coeff * coeff1 * coeff2
                )

        self.terms_dict = res_terms_dict
        return self

    #
    # Substitution
    #

    def subs(self, subs_dict):
        """

        Parameters
        ----------
        subs_dict : dict
            A dictionary with indices (int) as keys and numbers (int, float, complex) as values.

        Returns
        -------
        result : BoundQubitOperator
            The resulting BoundQubitOperator.

        """

        res_terms_dict = {}

        for pauli, coeff in self.terms_dict.items():
            curr_pauli, curr_coeff = pauli.subs(subs_dict)
            res_terms_dict[curr_pauli] = (
                res_terms_dict.get(curr_pauli, 0) + curr_coeff * coeff
            )

        result = BoundQubitOperator(res_terms_dict)
        return result

    #
    # Miscellaneous
    #

    def apply_threshold(self, threshold):
        """
        Removes all Pauli terms with coefficient absolute value below the specified threshold.

        Parameters
        ----------
        threshold : float
            The threshold for the coefficients of the Pauli terms.

        """

        delete_list = []
        for pauli, coeff in self.terms_dict.items():
            if abs(coeff) < threshold:
                delete_list.append(pauli)
        for pauli in delete_list:
            del self.terms_dict[pauli]

    def to_sparse_matrix(self):
        """
        Returns a matrix representing the operator.

        Returns
        -------
        M : scipy.sparse.csr_matrix
            A sparse matrix representing the operator.

        """

        import scipy.sparse as sp
        from scipy.sparse import kron as TP, csr_matrix

        I = csr_matrix([[1, 0], [0, 1]])

        def get_matrix(P):
            if P == "I":
                return csr_matrix([[1, 0], [0, 1]])
            if P == "X":
                return csr_matrix([[0, 1], [1, 0]])
            if P == "Y":
                return csr_matrix([[0, -1j], [1j, 0]])
            else:
                return csr_matrix([[1, 0], [0, -1]])

        def recursive_TP(keys, factor_dict):
            if len(keys) == 1:
                return get_matrix(factor_dict.get(keys[0], "I"))
            return TP(
                get_matrix(factor_dict.get(keys.pop(0), "I")),
                recursive_TP(keys, factor_dict),
            )

        factor_dicts = []
        coeffs = []

        keys = set()
        for pauli, coeff in self.terms_dict.items():
            curr_dict = pauli.factor_dict
            keys.update(set(curr_dict.keys()))
            factor_dicts.append(curr_dict)
            coeffs.append(coeff)

        keys = set()
        for item in factor_dicts:
            keys.update(set(item.keys()))
        keys = sorted(keys)
        dim = len(keys)

        m = len(coeffs)
        M = sp.csr_matrix((2**dim, 2**dim))
        for k in range(m):
            M += complex(coeffs[k]) * recursive_TP(keys.copy(), factor_dicts[k])

        return M

    def ground_state_energy(self):
        """
        Calculates the ground state energy (i.e., the minimum eigenvalue) of the operator classically.

        Returns
        -------
        E : float
            The ground state energy.

        """

        from scipy.sparse.linalg import eigsh

        M = self.to_sparse_matrix()
        # Compute the smallest eigenvalue
        eigenvalues, _ = eigsh(M, k=1, which="SA")  # 'SA' stands for smallest algebraic
        E = eigenvalues[0]

        return E

    #
    # Partitions
    #

    # Commutativity: Partitions the BoundQubitOperator into BoundQubitOperators with pairwise commuting BoundQubitTerms
    def commuting_groups(self):
        r"""
        Partitions the QubitOperator into QubitOperators with pairwise commuting terms. That is,

        .. math::

            H = \sum_{i=1}^mH_i

        where the terms in each $H_i$ are mutually commuting.

        Returns
        -------
        groups : list[QubitOperator]
            The partition of the Hamiltonian.

        """

        groups = []  # Groups of commuting BoundQubitTerms

        # Sorted insertion heuristic https://quantum-journal.org/papers/q-2021-01-20-385/pdf/
        sorted_terms = sorted(
            self.terms_dict.items(), key=lambda item: abs(item[1]), reverse=True
        )

        for pauli, coeff in sorted_terms:

            commute_bool = False
            if len(groups) > 0:
                for group in groups:
                    for pauli_, coeff_ in group.terms_dict.items():
                        commute_bool = pauli_.commute(pauli)
                        if not commute_bool:
                            break
                    if commute_bool:
                        group.terms_dict[pauli] = coeff
                        break
            if len(groups) == 0 or not commute_bool:
                groups.append(BoundQubitOperator({pauli: coeff}))

        return groups

    # Qubit-wise commutativity: Partitions the BoundQubitOperator into BoundQubitOperators with pairwise qubit-wise commuting BoundQubitTerms
    def commuting_qw_groups(self, show_bases=False):
        r"""
        Partitions the BoundQubitOperator into BoundQubitOperators with pairwise qubit-wise commuting terms. That is,

        .. math::

            H = \sum_{i=1}^mH_i

        where the terms in each $H_i$ are mutually qubit-wise commuting.

        Returns
        -------
        groups : list[BoundQubitOperator]
            The partition of the Hamiltonian.

        """

        groups = []  # Groups of qubit-wise commuting BoundQubitTerms
        bases = []  # Bases as BoundQubitTerms

        # Sorted insertion heuristic https://quantum-journal.org/papers/q-2021-01-20-385/pdf/
        sorted_terms = sorted(
            self.terms_dict.items(), key=lambda item: abs(item[1]), reverse=True
        )

        for pauli, coeff in sorted_terms:

            commute_bool = False
            if len(groups) > 0:
                n = len(groups)
                for i in range(n):
                    commute_bool = bases[i].commute_qw(pauli)
                    if commute_bool:
                        bases[i].update(pauli.factor_dict)
                        groups[i].terms_dict[pauli] = coeff
                        break
            if len(groups) == 0 or not commute_bool:
                groups.append(BoundQubitOperator({pauli: coeff}))
                bases.append(pauli.copy())

        if show_bases:
            return groups, bases
        else:
            return groups

    #
    # Measurement settings
    #

    def get_measurement(
        self,
        qarg,
        precision=0.01,
        backend=None,
        shots=1000000,
        compile=True,
        compilation_kwargs={},
        subs_dic={},
        precompiled_qc=None,
        measurement_data=None,  # measurement settings
    ):
        r"""
        This method returns the expected value of a Hamiltonian for the state of a quantum argument.

        Parameters
        ----------
        qarg : QuantumVariable, QuantumArray or list[QuantumVariable]
            The quantum argument to evaluate the Hamiltonian on.
        precision: float, optional
            The precision with which the expectation of the Hamiltonian is to be evaluated.
            The default is 0.01. The number of shots scales quadratically with the inverse precision.
        backend : BackendClient, optional
            The backend on which to evaluate the quantum circuit. The default can be
            specified in the file default_backend.py.
        shots : integer, optional
            The maximum amount of shots to evaluate the expectation of the Hamiltonian.
            The default is 1000000.
        compile : bool, optional
            Boolean indicating if the .compile method of the underlying QuantumSession
            should be called before. The default is True.
        compilation_kwargs  : dict, optional
            Keyword arguments for the compile method. For more details check
            :meth:`QuantumSession.compile <qrisp.QuantumSession.compile>`. The default
            is ``{}``.
        subs_dic : dict, optional
            A dictionary of Sympy symbols and floats to specify parameters in the case
            of a circuit with unspecified, :ref:`abstract parameters<QuantumCircuit>`.
            The default is {}.
        circuit_preprocessor : Python function, optional
            A function which recieves a QuantumCircuit and returns one, which is applied
            after compilation and parameter substitution. The default is None.
        precompiled_qc : QuantumCircuit, optional
            A precompiled quantum circuit.

        Raises
        ------
        Exception
            If the containing QuantumSession is in a quantum environment, it is not
            possible to execute measurements.

        Returns
        -------
        float
            The expected value of the Hamiltonian.

        Examples
        --------

        We define a Hamiltonian, and measure its expected value for the state of a :ref:`QuantumVariable`.

        ::

            from qrisp import QuantumVariable, h
            from qrisp.operators import X,Y,Z
            qv = QuantumVariable(2)
            h(qv)
            H = Z(0)*Z(1)
            res = H.get_measurement(qv)
            print(res)
            #Yields 0.0

        We define a Hamiltonian, and measure its expected value for the state of a :ref:`QuantumArray`.

        ::

            from qrisp import QuantumVariable, QuantumArray, h
            from qrisp.operators import X,Y,Z
            qtype = QuantumVariable(2)
            q_array = QuantumArray(qtype, shape=(2))
            h(q_array)
            H = Z(0)*Z(1) + X(2)*X(3)
            res = H.get_measurement(q_array)
            print(res)
            #Yields 1.0

        """

        unbound_operator, qarg = self.unbind()

        return unbound_operator.get_measurement(
            qarg,
            precision=precision,
            backend=backend,
            compile=compile,
            compilation_kwargs=compilation_kwargs,
            subs_dic=subs_dic,
            precompiled_qc=precompiled_qc,
            measurement_data=measurement_data,
        )

    def unbind(self):
        from qrisp.operators import QubitTerm, QubitOperator

        participating_qubits = []
        for term, coeff in self.terms_dict.items():
            for qb in term.factor_dict.keys():
                if qb not in participating_qubits:
                    participating_qubits.append(qb)

        index_inv = {
            participating_qubits[i]: i for i in range(len(participating_qubits))
        }

        unbound_operator = 0

        for term, coeff in self.terms_dict.items():
            new_term_factor_dict = {}
            for qb, factor in term.factor_dict.items():
                new_term_factor_dict[index_inv[qb]] = factor

            unbound_operator += QubitOperator({QubitTerm(new_term_factor_dict): coeff})

        return unbound_operator, participating_qubits

    #
    # Trotterization
    #

    def trotterization(self):
        r"""
        Returns a function for performing Hamiltonian simulation, i.e., approximately implementing the unitary operator $e^{itH}$ via Trotterization.

        Returns
        -------
        U : function
            A Python function that implements the first order Suzuki-Trotter formula.
            Given a Hamiltonian $H=H_1+\dotsb +H_m$ the unitary evolution $e^{itH}$ is
            approximated by

            .. math::

                e^{itH}\approx U_1(t,N)=\left(e^{iH_1t/N}\dotsb e^{iH_mt/N}\right)^N

            This function receives the following arguments:

            * qarg : QuantumVariable
                The quantum argument.
            * t : float, optional
                The evolution time $t$. The default is 1.
            * steps : int, optional
                The number of Trotter steps $N$. The default is 1.
            * iter : int, optional
                The number of iterations the unitary $U_1(t,N)$ is applied. The default is 1.

        """

        def change_of_basis(factor_dict):
            for qubit, axis in factor_dict.items():
                if axis == "X":
                    h(qubit)
                if axis == "Y":
                    sx(qubit)

        groups, bases = self.commuting_qw_groups(show_bases=True)

        def trotter_step(qarg, t, steps):

            # qubit to apply gphase for identity term
            if isinstance(qarg, list):
                qubit = qarg[0][0]
            else:
                qubit = qarg[0]

            for index, basis in enumerate(bases):
                with conjugate(change_of_basis)(basis.factor_dict):
                    for term, coeff in groups[index].terms_dict.items():
                        term.simulate(coeff * t / steps, qubit)

        def U(qarg, t=1, steps=1, iter=1):
            merge([qarg])
            with IterationEnvironment(qarg.qs, iter * steps):
                trotter_step(qarg, t, steps)

        return U