# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# 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.
"""Lie algebra optimizer"""
import warnings

import numpy as np
from scipy.sparse.linalg import expm
import pennylane as qml


@qml.qfunc_transform
def append_time_evolution(tape, riemannian_gradient, t, n, exact=False):
    r"""Append an approximate time evolution, corresponding to a Riemannian
    gradient on the Lie group, to an existing circuit.

    We want to implement the time evolution generated by an operator of the form

    .. math::

        \text{grad} f(U) = sum_i c_i O_i,

    where :math:`O_i` are Pauli words and :math:`c_t \in \mathbb{R}`.
    If ``exact`` is ``False``, we Trotterize this operator and apply the unitary

    .. math::

        U' = \prod_{n=1}^{N_{Trot.}} \left(\prod_i \exp{-it / N_{Trot.} O_i}\right),

    which is then appended to the current circuit.

    If ``exact`` is ``True``, we calculate the exact time evolution for the Riemannian gradient
     by way of the matrix exponential.

    .. math:

        U' = \exp{-it \text{grad} f(U)}

    and append this unitary.

    Args:
        tape (QuantumTape or .QNode): circuit to transform.
        riemannian_gradient (.Hamiltonian): Hamiltonian object representing the Riemannian gradient.
        t (float): time evolution parameter.
        n (int): number of Trotter steps.

    """
    for obj in tape.operations:
        qml.apply(obj)
    if exact:
        qml.QubitUnitary(
            expm(-1j * t * qml.utils.sparse_hamiltonian(riemannian_gradient).toarray()),
            wires=range(len(riemannian_gradient.wires)),
        )
    else:
        qml.templates.ApproxTimeEvolution(riemannian_gradient, t, n)

    for obj in tape.measurements:
        qml.apply(obj)


@qml.batch_transform
def algebra_commutator(tape, observables, lie_algebra_basis_names, nqubits):
    """Calculate the Riemannian gradient in the Lie algebra with the parameter shift rule
    (see :meth:`LieAlgebraOptimizer.get_omegas`).

    Args:
        tape (.QuantumTape or .QNode): input circuit
        observables (list[.Observable]): list of observables to be measured. Can be grouped.
        lie_algebra_basis_names (list[str]): List of strings corresponding to valid Pauli words.
        nqubits (int): the number of qubits.

    Returns:
        function or tuple[list[QuantumTape], function]:

        - If the input is a QNode, an object representing the Riemannian gradient function
          of the QNode that can be executed with the same arguments as the QNode to obtain
          the Lie algebra commutator.

        - If the input is a tape, a tuple containing a
          list of generated tapes, together with a post-processing
          function to be applied to the results of the evaluated tapes
          in order to obtain the Lie algebra commutator.
    """
    tapes_plus_total = []
    tapes_min_total = []
    for obs in observables:
        for o in obs:
            # create a list of tapes for the plus and minus shifted circuits
            queues_plus = [qml.queuing.AnnotatedQueue() for _ in lie_algebra_basis_names]
            queues_min = [qml.queuing.AnnotatedQueue() for _ in lie_algebra_basis_names]

            # loop through all operations on the input tape
            for op in tape.operations:
                for t in queues_plus + queues_min:
                    with t:
                        qml.apply(op)
            for i, t in enumerate(queues_plus):
                with t:
                    qml.PauliRot(
                        np.pi / 2,
                        lie_algebra_basis_names[i],
                        wires=list(range(nqubits)),
                    )
                    qml.expval(o)
            for i, t in enumerate(queues_min):
                with t:
                    qml.PauliRot(
                        -np.pi / 2,
                        lie_algebra_basis_names[i],
                        wires=list(range(nqubits)),
                    )
                    qml.expval(o)
            tapes_plus_total.extend(
                [
                    qml.tape.QuantumScript(*qml.queuing.process_queue(q), name=f"{p}_p")
                    for q, p in zip(queues_plus, lie_algebra_basis_names)
                ]
            )
            tapes_min_total.extend(
                [
                    qml.tape.QuantumScript(*qml.queuing.process_queue(q), name=f"{p}_m")
                    for q, p in zip(queues_min, lie_algebra_basis_names)
                ]
            )
    return tapes_plus_total + tapes_min_total, None


class LieAlgebraOptimizer:
    r"""Lie algebra optimizer.

    Riemannian gradient descent algorithms can be used to optimize a function directly on a Lie group
    as opposed to on a Euclidean parameter space. Consider the function
    :math:`f(U) = \text{Tr}(U \rho_0 U^\dagger H)`
    for a given Hamiltonian :math:`H`, unitary :math:`U\in \text{SU}(2^N)` and initial state
    :math:`\rho_0`. One can show that this function is minimized by the flow equation

    .. math::

        \dot{U} = \text{grad}f(U)

    where :math:`\text{grad}` is the Riemannian gradient operator on :math:`\text{SU}(2^N)`.
    By discretizing the flow above, we see that a step of the Lie algebra optimizer
    iterates the Riemannian gradient flow on :math:`\text{SU}(2^N)` as

    .. math::

        U^{(t+1)} = \exp\left\{\epsilon\: \text{grad}f(U^{(t)}) U^{(t)}\right\},


    where :math:`\epsilon` is a user-defined hyperparameter corresponding to the step size.

    The gradient in the Lie algebra is given by

    .. math::

        \text{grad}f(U^{(t)}) = -\left[U \rho U^\dagger, H\right] .


    Hence we see that subsequent steps of this optimizer will append the unitary generated by the Lie
    gradient and grow the circuit.

    The exact Riemannian gradient flow on :math:`\text{SU}(2^N)` has desirable optimization properties
    that can guarantee convergence to global minima under mild assumptions. However, this comes
    at a cost. Since :math:`\text{dim}(\text{SU}(2^N)) = 4^N-1`, we need an exponential number
    of parameters to calculate the gradient. This will not be problematic for small systems (:math:`N<5`),
    but will quickly get out of control as the number of qubits increases.

    To resolve this issue, we can restrict the Riemannian gradient to a subspace of the Lie algebra and calculate an
    approximate Riemannian gradient flow. The choice of restriction will affect the optimization behavior
    and quality of the final solution.

    For more information on Riemannian gradient flows on Lie groups see
    `T. Schulte-Herbrueggen et. al. (2008) <https://arxiv.org/abs/0802.4195>`_
    and the application to quantum circuits
    `Wiersema and Killoran (2022) <https://arxiv.org/abs/2202.06976>`_.

    Args:
        circuit (.QNode): a user defined circuit that does not take any arguments and returns
            the expectation value of a ``qml.Hamiltonian``.
        stepsize (float): the user-defined hyperparameter :math:`\epsilon`.
        restriction (.Hamiltonian): Restrict the Lie algebra to a corresponding subspace of
            the full Lie algebra. This restriction should be passed in the form of a
            ``qml.Hamiltonian`` that consists only of Pauli words.
        exact (bool): Flag that indicates wether we approximate the Riemannian gradient with a
            Trotterization or calculate the exact evolution via a matrix exponential. The latter is
            not hardware friendly and can only be done in simulation.

    **Examples**

    Define a Hamiltonian cost function to minimize:

    >>> coeffs = [-1., -1., -1.]
    >>> observables = [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0) @ qml.PauliX(1)]
    >>> hamiltonian = qml.Hamiltonian(coeffs, observables)

    Create an initial state and return the expectation value of the Hamiltonian:

    >>> @qml.qnode(qml.device("default.qubit", wires=2))
    ... def quant_fun():
    ...     qml.RX(0.1, wires=[0])
    ...     qml.RY(0.5, wires=[1])
    ...     qml.CNOT(wires=[0,1])
    ...     qml.RY(0.6, wires=[0])
    ...     return qml.expval(hamiltonian)

    Instantiate the optimizer with the initial circuit and the cost function and set the stepsize
    accordingly:

    >>> opt = qml.LieAlgebraOptimizer(circuit=quant_fun, stepsize=0.1)

    Applying 5 steps gets us close the ground state of :math:`E\approx-2.23`:

    >>> for step in range(6):
    ...    circuit, cost = opt.step_and_cost()
    ...    print(f"Step {step} - cost {cost}")
    Step 0 - cost -1.3351865007304005
    Step 1 - cost -1.9937887238935206
    Step 2 - cost -2.1524234485729834
    Step 3 - cost -2.1955105378898487
    Step 4 - cost -2.2137628169764256
    Step 5 - cost -2.2234364822091575

    The optimized circuit is returned at each step, and can be used as any other QNode:

    >>> circuit()
    -2.2283086057521713

    """
    # pylint: disable=too-many-arguments
    # pylint: disable=too-many-instance-attributes
    def __init__(self, circuit, stepsize=0.01, restriction=None, exact=False, trottersteps=1):

        if not isinstance(circuit, qml.QNode):
            raise TypeError(f"circuit must be a QNode, received {type(circuit)}")

        self.circuit = circuit
        self.circuit.construct([], {})
        self.hamiltonian = circuit.func().obs
        if not isinstance(self.hamiltonian, qml.Hamiltonian):
            raise TypeError(
                f"circuit must return the expectation value of a Hamiltonian,"
                f"received {type(circuit.func().obs)}"
            )
        self.nqubits = len(circuit.device.wires)

        if self.nqubits > 4:
            warnings.warn(
                "The exact Riemannian gradient is exponentially expensive in the number of qubits, "
                f"optimizing a {self.nqubits} qubit circuit may be slow.",
                UserWarning,
            )
        if restriction is not None and not isinstance(restriction, qml.Hamiltonian):
            raise TypeError(f"restriction must be a Hamiltonian, received {type(restriction)}")
        (
            self.lie_algebra_basis_ops,
            self.lie_algebra_basis_names,
        ) = self.get_su_n_operators(restriction)
        self.exact = exact
        self.trottersteps = trottersteps
        self.coeffs, self.observables = self.hamiltonian.terms()
        self.stepsize = stepsize

    def step(self):
        r"""Update the circuit with one step of the optimizer.

        Returns:
           float: the optimized circuit and the objective function output prior
           to the step.
        """
        return self.step_and_cost()[0]

    def step_and_cost(self):
        r"""Update the circuit with one step of the optimizer and return the corresponding
        objective function value prior to the step.

        Returns:
            tuple[.QNode, float]: the optimized circuit and the objective function output prior
            to the step.
        """
        # pylint: disable=not-callable

        cost = self.circuit()
        omegas = self.get_omegas()
        non_zero_omegas = -omegas[omegas != 0]

        nonzero_idx = np.nonzero(omegas)[0]
        non_zero_lie_algebra_elements = [self.lie_algebra_basis_names[i] for i in nonzero_idx]

        lie_gradient = qml.Hamiltonian(
            non_zero_omegas,
            [qml.pauli.string_to_pauli_word(ps) for ps in non_zero_lie_algebra_elements],
        )
        new_circuit = append_time_evolution(
            lie_gradient, self.stepsize, self.trottersteps, self.exact
        )(self.circuit.func)

        # we can set diff_method=None because the gradient of the QNode is computed
        # directly in this optimizer
        self.circuit = qml.QNode(new_circuit, self.circuit.device, diff_method=None)
        return self.circuit, cost

    def get_su_n_operators(self, restriction):
        r"""Get the SU(N) operators. The dimension of the group is :math:`N^2-1`.

        Args:
            restriction (.Hamiltonian): Restrict the Riemannian gradient to a subspace.

        Returns:
            tuple[list[array[complex]], list[str]]: list of :math:`N^2 \times N^2` NumPy complex arrays and corresponding Pauli words.
        """

        operators = []
        names = []
        # construct the corresponding pennylane observables
        wire_map = dict(zip(range(self.nqubits), range(self.nqubits)))
        if restriction is None:
            for ps in qml.pauli.pauli_group(self.nqubits):
                operators.append(ps)
                names.append(qml.pauli.pauli_word_to_string(ps, wire_map=wire_map))
        else:
            for ps in set(restriction.ops):
                operators.append(ps)
                names.append(qml.pauli.pauli_word_to_string(ps, wire_map=wire_map))

        return operators, names

    def get_omegas(self):
        r"""Measure the coefficients of the Riemannian gradient with respect to a Pauli word basis.
        We want to calculate the components of the Riemannian gradient in the Lie algebra
        with respect to a Pauli word basis. For a Hamiltonian of the form :math:`H = \sum_i c_i O_i`,
        where :math:`c_i\in\mathbb{R}`, this can be achieved by calculating

        .. math::

            \omega_{i,j} = \text{Tr}(c_i[\rho, O_i] P_j)

        where :math:`P_j` is a Pauli word in the set of Pauli monomials on :math:`N` qubits.

        Via the parameter shift rule, the commutator can be calculated as

        .. math::

            [\rho, O_i] = \frac{1}{2}(V(\pi/2) \rho V^\dagger(\pi/2) - V(-\pi/2) \rho V^\dagger(-\pi/2))

        where :math:`V` is the unitary generated by the Pauli word :math:`V(\theta) = \exp\{-i\theta P_j\}`.

        Returns:
            array: array of omegas for each direction in the Lie algebra.
        """

        obs_groupings, _ = qml.pauli.group_observables(self.observables, self.coeffs)
        # get all circuits we need to calculate the coefficients
        circuits = algebra_commutator(
            self.circuit.qtape,
            obs_groupings,
            self.lie_algebra_basis_names,
            self.nqubits,
        )[0]
        circuits = qml.execute(circuits, self.circuit.device, gradient_fn=None)
        circuits_plus = np.array(circuits[: len(circuits) // 2]).reshape(
            len(self.coeffs), len(self.lie_algebra_basis_names)
        )
        circuits_min = np.array(circuits[len(circuits) // 2 :]).reshape(
            len(self.coeffs), len(self.lie_algebra_basis_names)
        )

        # For each observable O_i in the Hamiltonian, we have to calculate all Lie coefficients
        omegas = circuits_plus - circuits_min

        return np.dot(self.coeffs, omegas)
