# Copyright (C) Unitary Fund
#
# This source code is licensed under the GPL license (v3) found in the
# LICENSE file in the root directory of this source tree.

"""High-level digital dynamical decoupling (DDD) tools."""

from typing import (
    Optional,
    Callable,
    Union,
    Tuple,
    Dict,
    Any,
    List,
)

from functools import wraps, partial
import numpy as np

from mitiq import Executor, Observable, QPROGRAM, QuantumResult

from mitiq.ddd.insertion import insert_ddd_sequences


def execute_with_ddd(
    circuit: QPROGRAM,
    executor: Union[Executor, Callable[[QPROGRAM], QuantumResult]],
    observable: Optional[Observable] = None,
    *,
    rule: Callable[[int], QPROGRAM],
    rule_args: Dict[str, Any] = {},
    num_trials: int = 1,
    full_output: bool = False,
) -> Union[float, Tuple[float, Dict[str, Any]]]:
    r"""Estimates the error-mitigated expectation value associated to the
    input circuit, via the application of digital dynamical decoupling (DDD).

    Args:
        circuit: The input circuit to execute with DDD.
        executor: A Mitiq executor that executes a circuit and returns the
            unmitigated ``QuantumResult`` (e.g. an expectation value).
        observable: Observable to compute the expectation value of. If None,
            the ``executor`` must return an expectation value. Otherwise,
            the ``QuantumResult`` returned by ``executor`` is used to compute
            the expectation of the observable.
        rule: A function that takes as main argument a slack length (i.e. the
            number of idle moments) of a slack window (i.e. a single-qubit idle
            window in a circuit) and returns the DDD sequence of gates to be
            applied in that window. Mitiq provides standard built-in rules
            that can be directly imported from ``mitiq.ddd.rules``.
        rule_args: An optional dictionary of keyword arguments for ``rule``.
        num_trials: The number of independent experiments to average over.
            A number larger than 1 can be useful to average over multiple
            applications of a rule returning non-deterministic DDD sequences.
        full_output: If ``False`` only the mitigated expectation value is
            returned. If ``True`` a dictionary containing all DDD data is
            returned too.

    Returns:
        The tuple ``(ddd_value, ddd_data)`` where ``ddd_value`` is the
        expectation value estimated with DDD and ``ddd_data`` is a dictionary
        containing all the raw data involved in the DDD process (e.g. the
        circuit filled with DDD sequences). If ``full_output`` is false,
        only ``ddd_value`` is returned.
    """
    # Initialize executor
    if not isinstance(executor, Executor):
        executor = Executor(executor)

    rule_partial: Callable[[int], QPROGRAM]
    rule_partial = partial(rule, **rule_args)

    # Insert DDD sequences in (a copy of) the input circuit
    circuits_with_ddd = [
        insert_ddd_sequences(circuit, rule_partial) for _ in range(num_trials)
    ]
    results = executor.evaluate(
        circuits_with_ddd,
        observable,
        force_run_all=True,
    )

    assert len(results) == num_trials

    ddd_value = np.sum(results) / num_trials
    if not full_output:
        return ddd_value

    ddd_data = {
        "ddd_value": ddd_value,
        "ddd_trials": results,
        "circuits_with_ddd": circuits_with_ddd,
    }
    return ddd_value, ddd_data


def mitigate_executor(
    executor: Callable[[QPROGRAM], QuantumResult],
    observable: Optional[Observable] = None,
    *,
    rule: Callable[[int], QPROGRAM],
    rule_args: Dict[str, Any] = {},
    num_trials: int = 1,
    full_output: bool = False,
) -> Callable[[QPROGRAM], Union[float, Tuple[float, Dict[str, Any]]]]:
    """Returns a modified version of the input 'executor' which is
    error-mitigated with digital dynamical decoupling (DDD).

    Args:
        executor: A function that executes a circuit and returns the
            unmitigated `QuantumResult` (e.g. an expectation value).
        observable: Observable to compute the expectation value of. If None,
            the `executor` must return an expectation value. Otherwise,
            the `QuantumResult` returned by `executor` is used to compute the
            expectation of the observable.
        rule: A function that takes as main argument a slack length (i.e. the
            number of idle moments) of a slack window (i.e. a single-qubit idle
            window in a circuit) and returns the DDD sequence of gates to be
            applied in that window. Mitiq provides standard built-in rules
            that can be directly imported from `mitiq.ddd.rules`.
        rule_args: An optional dictionary of keyword arguments for `rule`.
        num_trials: The number of independent experiments to average over.
            A number larger than 1 can be useful to average over multiple
            applications of a rule returning non-deterministic DDD sequences.
        full_output: If False only the mitigated expectation value is returned.
            If True a dictionary containing all DDD data is returned too.

    Returns:
        The error-mitigated version of the input executor.
    """
    executor_obj = Executor(executor)
    if not executor_obj.can_batch:

        @wraps(executor)
        def new_executor(
            circuit: QPROGRAM,
        ) -> Union[float, Tuple[float, Dict[str, Any]]]:
            return execute_with_ddd(
                circuit,
                executor,
                observable,
                rule=rule,
                rule_args=rule_args,
                num_trials=num_trials,
                full_output=full_output,
            )

    else:

        @wraps(executor)
        def new_executor(
            circuits: List[QPROGRAM],
        ) -> List[Union[float, Tuple[float, Dict[str, Any]]]]:
            return [
                execute_with_ddd(
                    circuit,
                    executor,
                    observable,
                    rule=rule,
                    rule_args=rule_args,
                    num_trials=num_trials,
                    full_output=full_output,
                )
                for circuit in circuits
            ]

    return new_executor


def ddd_decorator(
    observable: Optional[Observable] = None,
    *,
    rule: Callable[[int], QPROGRAM],
    rule_args: Dict[str, Any] = {},
    num_trials: int = 1,
    full_output: bool = False,
) -> Callable[
    [Callable[[QPROGRAM], QuantumResult]],
    Callable[[QPROGRAM], Union[float, Tuple[float, Dict[str, Any]]]],
]:
    """Decorator which adds an error-mitigation layer based on digital
    dynamical decoupling (DDD) to an executor function, i.e., a function which
    executes a quantum circuit with an arbitrary backend and returns a
    ``QuantumResult`` (e.g. an expectation value).

    Args:
        observable: Observable to compute the expectation value of. If None,
            the `executor` must return an expectation value. Otherwise,
            the `QuantumResult` returned by `executor` is used to compute the
            expectation of the observable.
        rule: A function that takes as main argument a slack length (i.e. the
            number of idle moments) of a slack window (i.e. a single-qubit idle
            window in a circuit) and returns the DDD sequence of gates to be
            applied in that window. Mitiq provides standard built-in rules
            that can be directly imported from `mitiq.ddd.rules`.
        rule_args: An optional dictionary of keyword arguments for `rule`.
        num_trials: The number of independent experiments to average over.
            A number larger than 1 can be useful to average over multiple
            applications of a rule returning non-deterministic DDD sequences.
        full_output: If False only the mitigated expectation value is returned.
            If True a dictionary containing all DDD data is returned too.

    Returns:
        The error-mitigating decorator to be applied to an executor function.
    """

    def decorator(
        executor: Callable[[QPROGRAM], QuantumResult]
    ) -> Callable[[QPROGRAM], Union[float, Tuple[float, Dict[str, Any]]]]:
        return mitigate_executor(
            executor,
            observable,
            rule=rule,
            rule_args=rule_args,
            num_trials=num_trials,
            full_output=full_output,
        )

    return decorator
