# 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.
"""
This module contains the qml.evolve function.
"""
import warnings
from functools import singledispatch

from pennylane.operation import Operator
from pennylane.ops import Evolution
from pennylane.pulse import ParametrizedEvolution, ParametrizedHamiltonian


@singledispatch
def evolve(*args, **kwargs):  # pylint: disable=unused-argument
    r"""This method is dispatched and its functionality depends on the type of the input ``op``.

    .. raw:: html

        <html>
            <h3>Input: Operator</h3>
            <hr>
        </html>

    Returns a new operator that computes the evolution of ``op``.

    .. math::

        e^{-i x \bm{O}}

    Args:
        op (.Operator): operator to evolve
        coeff (float): coefficient multiplying the exponentiated operator

    Returns:
        .Evolution: evolution operator

    .. seealso::

        :class:`~.Evolution`

    **Examples**

    We can use ``qml.evolve`` to compute the evolution of any PennyLane operator:

    >>> op = qml.evolve(qml.PauliX(0), coeff=2)
    >>> op
    Exp(2j PauliX)

    .. raw:: html

        <html>
            <h3>Input: ParametrizedHamiltonian</h3>
            <hr>
        </html>

    Args:
        op (.ParametrizedHamiltonian): Hamiltonian to evolve

    Returns:
        .ParametrizedEvolution: time evolution :math:`U(t_0, t_1)` of the Hamiltonian


    The function takes a :class:`.ParametrizedHamiltonian` and solves the time-dependent Schrodinger equation

    .. math:: \frac{\partial}{\partial t} |\psi\rangle = -i H(t) |\psi\rangle

    It returns a :class:`~.ParametrizedEvolution`, :math:`U(t_0, t_1)`, which is the solution to the time-dependent
    Schrodinger equation for the :class:`~.ParametrizedHamiltonian`, such that

    .. math:: |\psi(t_1)\rangle = U(t_0, t_1) |\psi(t_0)\rangle

    The :class:`~.ParametrizedEvolution` class uses a numerical ordinary differential equation
    solver (`here <https://github.com/google/jax/blob/main/jax/experimental/ode.py>`_).

    .. seealso::

        :class:`~.ParametrizedEvolution`

    **Examples**

    When evolving a :class:`.ParametrizedHamiltonian`, a :class:`.ParametrizedEvolution`
    instance is returned:

    .. code-block:: python3

        coeffs = [lambda p, t: p * t for _ in range(4)]
        ops = [qml.PauliX(i) for i in range(4)]

        # ParametrizedHamiltonian
        H = qml.dot(coeffs, ops)

        # ParametrizedEvolution
        ev = qml.evolve(H)

    >>> ev
    ParametrizedEvolution(wires=[0, 1, 2, 3])

    The :class:`.ParametrizedEvolution` is an :class:`~.Operator`, but does not have a defined matrix unless it
    is evaluated at set parameters. This is done by calling the :class:`.ParametrizedEvolution`, which has the call
    signature ``(p, t)``:

    >>>  qml.matrix(ev([1., 2., 3., 4.], t=[0, 4]))
    Array([[ 0.04930558+0.j        ,  0.        -0.03259093j,
         0.        +0.1052632j ,  0.06957878+0.j        ,
         0.        -0.01482305j, -0.00979751+0.j        ,
         0.03164552+0.j        ,  0.        -0.0209179j ,
         0.        +0.33526757j,  0.22161038+0.j        ,
         ...
         ...
         ...
         0.        -0.03259093j,  0.04930566+0.j        ]],      dtype=complex64)

    Additional options regarding how the matrix is calculated can be passed to the :class:`.ParametrizedEvolution`
    along with the parameters, as keyword arguments. These options are:

    - ``atol (float, optional)``: Absolute error tolerance
    - ``rtol (float, optional)``: Relative error tolerance
    - ``mxstep (int, optional)``: maximum number of steps to take for each time point
    - ``hmax (float, optional)``: maximum step size

    If not specified, they will default to predetermined values.

    The :class:`~.ParametrizedEvolution` can be implemented in a QNode:

    .. code-block:: python

        import jax

        dev = qml.device("default.qubit.jax", wires=4)
        @jax.jit
        @qml.qnode(dev, interface="jax")
        def circuit(params):
            qml.evolve(H)(params, t=[0, 10])
            return qml.expval(qml.PauliZ(0))

    >>> params = [1., 2., 3., 4.]
    >>> circuit(params)
    Array(0.8627419, dtype=float32)

    >>> jax.grad(circuit)(params)
    [Array(50.690746, dtype=float32),
    Array(-6.296886e-05, dtype=float32),
    Array(-6.3341584e-05, dtype=float32),
    Array(-7.052516e-05, dtype=float32)]

    .. note::
        In the example above, the decorator ``@jax.jit`` is used to compile this execution just-in-time. This means
        the first execution will typically take a little longer with the benefit that all following executions
        will be significantly faster, see the jax docs on jitting. JIT-compiling is optional, and one can remove
        the decorator when only single executions are of interest.

    Please check the :class:`.ParametrizedEvolution` class for more usage details.
    """


# pylint: disable=missing-docstring
@evolve.register
def parametrized_evolution(op: ParametrizedHamiltonian):
    return ParametrizedEvolution(H=op)


# pylint: disable=missing-docstring
@evolve.register
def evolution(op: Operator, coeff: float = 1, num_steps: int = None):
    with warnings.catch_warnings():
        # Ignore the warning raised in `Evolution`
        warnings.simplefilter("ignore")
        ev = Evolution(op, -1 * coeff, num_steps)
    return ev
