#!/usr/bin/python3
# -*- coding: utf-8 -*-
#
# Copyright (c) 2022 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.

"""
This example demonstrates how to use the Measurement Error Mitigation method to
improve the accuracy of the outcomes generated by a noisy quantum computer.
"""
import copy

from QCompute import *
from qcompute_qep.utils import expval_z_from_counts
from qcompute_qep.utils.types import QProgram, QComputer, number_of_qubits
from qcompute_qep.utils.circuit import execute
from qcompute_qep.measurement.correction import LeastSquareCorrector
from qcompute_qep.measurement.correction import InverseCorrector
from qcompute_qep.measurement.correction import IBUCorrector
from qcompute_qep.measurement.correction import NeumannCorrector
from qcompute_qep.measurement.utils import plot_histograms
from qcompute_qep.measurement.correction import dict2vector


# Set the default number of shots
NUMBER_OF_SHOTS = 4096


def calculator(qp: QProgram = None, qc: QComputer = None):
    """
    Run the quantum program on the quantum computer and estimate the expectation value.
    This function must be specified by the user.

    :param qp: QProgram, describes the quantum program
    :param qc: QComputer, specifies the quantum computer
    :return: the evaluated expectation value
    """
    # Run the quantum program on the quantum computer and obtain the measurement outcomes
    counts = execute(qp, qc, shots=NUMBER_OF_SHOTS)
    # Compute the expectation value from counts
    expval = expval_z_from_counts(counts)
    return expval, counts


def corrected_calculator(qp: QProgram = None, qc: QComputer = None):
    """
    Run the quantum program on the quantum computer and estimate the expectation value.
    This function must be specified by the user.
    The measurement outcomes are corrected.

    :param qp: QProgram, describes the quantum program
    :param qc: QComputer, specifies the quantum computer
    :return: the evaluated expectation value
    """
    # Run the quantum program on the quantum computer and obtain the measurement outcomes
    counts = execute(qp, qc, shots=NUMBER_OF_SHOTS)
    # Number of qubits involved in the computation
    n = number_of_qubits(qp)

    # Use different combinations of "calibrator + corrector" to mitigate the noisy outcomes
    lsc = LeastSquareCorrector(qc=qc, calibrator='complete', qubits=range(n))
    lsc_tp = LeastSquareCorrector(qc=qc, calibrator='tp', qubits=range(n))
    inv = InverseCorrector(qc=qc, calibrator='complete', qubits=range(n))
    inv_tp = InverseCorrector(qc=qc, calibrator='tp', qubits=range(n))
    kwarg = {'tol': 1e-6, 'max_iter': 100}
    ibu = IBUCorrector(qc=qc, calibrator='complete', qubits=range(n), **kwarg)
    ibu_tp = IBUCorrector(qc=qc, calibrator='tp', qubits=range(n), **kwarg)
    neu = NeumannCorrector(qc=qc, calibrator='complete', qubits=range(n), **kwarg)
    neu_tp = NeumannCorrector(qc=qc, calibrator='tp', qubits=range(n), **kwarg)
    counts_corr_lsc = lsc.correct(counts)
    counts_corr_inv = inv.correct(counts)
    counts_corr_ibu = ibu.correct(counts)
    counts_corr_neu = neu.correct(counts)
    counts_corr_lsc_tp = lsc_tp.correct(counts)
    counts_corr_inv_tp = inv_tp.correct(counts)
    counts_corr_ibu_tp = ibu_tp.correct(counts)
    counts_corr_neu_tp = neu_tp.correct(counts)

    # Compute the expectation value from the mitigated counts
    expval_lsc = expval_z_from_counts(counts_corr_lsc)
    expval_inv = expval_z_from_counts(counts_corr_inv)
    expval_ibu = expval_z_from_counts(counts_corr_ibu)
    expval_neu = expval_z_from_counts(counts_corr_neu)
    expval_lsc_tp = expval_z_from_counts(counts_corr_lsc_tp)
    expval_inv_tp = expval_z_from_counts(counts_corr_inv_tp)
    expval_ibu_tp = expval_z_from_counts(counts_corr_ibu_tp)
    expval_neu_tp = expval_z_from_counts(counts_corr_neu_tp)

    return expval_lsc, expval_lsc_tp, expval_inv, expval_inv_tp, expval_ibu, expval_ibu_tp, expval_neu, expval_neu_tp,\
        counts_corr_lsc, counts_corr_lsc_tp, counts_corr_inv, counts_corr_inv_tp, counts_corr_ibu, counts_corr_ibu_tp,\
        counts_corr_neu, counts_corr_neu_tp


if __name__ == '__main__':
    ##########################################################################################
    # Setup the quantum program for estimating the expectation value of GHZ state
    ##########################################################################################
    qp = QEnv()
    n = 4
    qp.Q.createList(n)
    H(qp.Q[0])
    for i in range(1, n):
        CX(qp.Q[0], qp.Q[i])
    MeasureZ(*qp.Q.toListPair())

    ##########################################################################################
    # Set the quantum hardware for running the Bell circuit.
    ##########################################################################################
    # For numeric test on the ideal simulator, change qc to BackendName.LocalBaiduSim2
    qc = BackendName.LocalBaiduSim2

    # For experiment on the real quantum device, change qc to BackendName.CloudBaiduQPUQian.
    # You must set your VIP token first in order to access the Baidu hardware.
    # Define.hubToken = "Token"
    # qc = BackendName.CloudBaiduQPUQian

    # For numeric test on the noisy simulator, change qc to Qiskit's FakeSantiago simulator
    # from qiskit.providers.fake_provider import FakeSantiago
    # qc = FakeSantiago()

    ##########################################################################################
    # Compute the ideal and noisy expectation values.
    ##########################################################################################
    val_ideal = calculator(copy.deepcopy(qp), qc=BackendName.LocalBaiduSim2)

    val_noisy = calculator(copy.deepcopy(qp), qc=qc)
    val_corr = corrected_calculator(copy.deepcopy(qp), qc=qc)

    print('-' * 20)
    print("The ideal expectation value is: {}".format(val_ideal[0]))
    print("The noisy expectation value is: {}".format(val_noisy[0]))
    print("The corrected expectation value using complete LSC is: {}".format(val_corr[0]))
    print("The corrected expectation value using tensor LSC is: {}".format(val_corr[1]))
    print("The corrected expectation value using complete INV is: {}".format(val_corr[2]))
    print("The corrected expectation value using tensor INV is: {}".format(val_corr[3]))
    print("The corrected expectation value using complete IBU is: {}".format(val_corr[4]))
    print("The corrected expectation value using tensor IBU is: {}".format(val_corr[5]))
    print("The corrected expectation value using complete Neumann is: {}".format(val_corr[6]))
    print("The corrected expectation value using tensor Neumann is: {}".format(val_corr[7]))
    new_val_ideal = dict2vector(val_ideal[1])
    new_val_noisy = dict2vector(val_noisy[1])
    new_val_corr_lsc = dict2vector(val_corr[8])
    new_val_corr_lsc_tensor = dict2vector(val_corr[9])
    new_val_corr_inv = dict2vector(val_corr[10])
    new_val_corr_inv_tensor = dict2vector(val_corr[11])
    new_val_corr_ibu = dict2vector(val_corr[12])
    new_val_corr_ibu_tensor = dict2vector(val_corr[13])
    new_val_corr_neu = dict2vector(val_corr[14])
    new_val_corr_neu_tensor = dict2vector(val_corr[15])

    count_list = [new_val_ideal, new_val_noisy, new_val_corr_lsc, new_val_corr_lsc_tensor,
                  new_val_corr_inv, new_val_corr_inv_tensor, new_val_corr_ibu, new_val_corr_ibu_tensor,
                  new_val_corr_neu, new_val_corr_neu_tensor]
    legend = ['val ideal simulator', 'val noisy', 'val corr noisy LSC',
              'val corr noisy LSC tensor', 'val corr noisy INV', 'val corr noisy INV tensor',
              'val corr noisy IBU', 'val corr noisy IBU tensor', 'val corr noisy Neumann',
              'val corr noisy Neumann tensor']
    plot_histograms(count_list, legend)
