#!/usr/bin/python3
# -*- coding: utf8 -*-

# Copyright (c) 2020 Baidu, Inc. 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.

"""
Zero-Noise Extrapolation (ZNE) is one of the most powerful techniques for mitigating quantum gate error
in quantum computing. Notice that ZNE does not directly reduce the inherent noise in the quantum
computing process, but rather helps us to obtain the expected computation result by repeating
the same quantum computing process various times with different level of noise. The
advantage of ZNE is that we do not need to know the exact form of the noise as well as how
to control the noise source.

The ZNE method is composed of two steps: rescaling noise and extrapolating.
Among many of noise rescaling techniques, time-variant rescaling is one of the most robust
and promising. This technique stretches the system Hamiltonian in time according to some
rescaling coefficient in order to obtain an equivalently noise-rescaled final quantum state.
For simplicity, we use the Richardson extrapolation in our Quanlse implementation, which is
a mature numeric algorithm that can eliminate error of any order in principle. We remark
that there are many other extrapolation methods such as polynomial and exponential extrapolation.

For more details, please visit https://quanlse.baidu.com/#/doc/tutorial-ZNE.
"""

from Quanlse.QHamiltonian import QHamiltonian as QHam
from Quanlse.QOperation import CircuitLine
import numpy as np

from Quanlse.QRpc import rpcCall

from Quanlse.QPlatform.Utilities import (
    numpyMatrixToDictMatrix,
    circuitLineDump,
)

from typing import List, Tuple


def remoteZNEMitigation(rho: np.ndarray,
                        circuit: List[CircuitLine] = None,
                        A: np.ndarray = None,
                        ham: QHam = None,
                        order: int = 1) -> Tuple[float, List[float], List[float]]:
    r"""
    Use the extrapolation method to mitigate the expectation value:

    :math:`{\rm Tr}[A C_n\cdots C_1 \rho C_1^\dagger \cdots C_n^\dagger]`

    under the assumption that the implementation of :math:`C_i` suffer from noise.
    More specifically, the hamiltonian designed to implement the
    quantum circuit :math:`C_n ... C_1` is not perfect.

    If `order` is not set, it will set to 1 by default.
    The Richardson extrapolation method will be applied.
    The return value will be of the form:

    `mitigatedValue, [coefficients], [noisy values]`

    where the `[coefficients]` and `[noisy values]` are the data points
    used to accomplish the Richardson extrapolation, each of size `order+1`.

    :param rho: the input quantum state
    :param circuit: the quantum circuit in Quantum Leaf syntax
    :param A: the quantum observable
    :param ham: the system Hamiltonian designed to implement the quantum circuit.
                It may be generated by Scheduler or optimizers.
    :param order: the order of the error to be extrapolated.
                Must be a positive integer.
    :return: The error mitigated expectation value and possibly
                a list of coefficients and noisy expectation values.
    """
    args = [
        numpyMatrixToDictMatrix(rho),
        circuitLineDump(circuit),
        numpyMatrixToDictMatrix(A),
        ham.dump(),
        order
    ]
    kwargs = {}

    origin = rpcCall("zneMitigation", args, kwargs)

    return origin["mitigatedValue"], origin["infidelities"], origin["noisyValues"]
