%load_ext autoreload
%autoreload 2

import numpy as np

from qiskit.circuit.library import RealAmplitudes
from qiskit.opflow import PauliSumOp

observable = PauliSumOp.from_list(
    [
        ("II", -1.052373245772859),
        ("IZ", 0.39793742484318045),
        ("ZI", -0.39793742484318045),
        ("ZZ", -0.01128010425623538),
        ("XX", 0.18093119978423156),
    ]
)
print("observable")
print(observable)

ansatz = RealAmplitudes(num_qubits=2, reps=2)
print("\nansatz")
ansatz.decompose().draw("mpl")

from qiskit.evaluators import PauliExpectationValue

from qiskit.providers.aer import AerSimulator

backend = AerSimulator()

expval = PauliExpectationValue(ansatz, observable, backend=backend)
expval.evaluate([0, 1, 1, 2, 3, 5], shots=1000)

# multi prameters
expval = PauliExpectationValue(ansatz, observable, backend=backend)
expval.evaluate([[0, 1, 1, 2, 3, 5], [1, 1, 2, 3, 5, 8]], shots=3000)

from qiskit.evaluators import ExactExpectationValue

expval = ExactExpectationValue(ansatz, observable, backend=backend)
expval.evaluate([0, 1, 1, 2, 3, 5])

expval = ExactExpectationValue(ansatz, observable, backend=backend)
expval.evaluate(np.array([[0, 1, 1, 2, 3, 5], [1, 1, 2, 3, 5, 8]]))

from qiskit.test.mock import FakeBogota

backend = AerSimulator.from_backend(FakeBogota())

expval = PauliExpectationValue(ansatz, observable, backend=backend)
expval.set_transpile_options(initial_layout=[3, 2])
print(expval.evaluate([0, 1, 1, 2, 3, 5]))
expval.transpiled_circuits[0].draw("mpl")

try:
    from qiskit import QuantumCircuit, QiskitError
    qc = QuantumCircuit(1, 1)
    qc.h(0)
    qc.measure(0, 0)
    from qiskit import IBMQ
    prov = IBMQ.load_account()
    ibmq_qasm_sim = prov.get_backend('ibmq_qasm_simulator')
    print(f"max_shots {backend.configuration().max_shots}")
    ibmq_qasm_sim.run(qc, shots=10000).result().get_counts()
except QiskitError as ex:
    print(ex.message)

expval = PauliExpectationValue(ansatz, observable, backend=ibmq_qasm_sim)
expval.evaluate([0, 1, 1, 2, 3, 5], shots=10000)

from qiskit.evaluators.backends import ReadoutErrorMitigation

backend = AerSimulator.from_backend(FakeBogota())
mit_tensored = ReadoutErrorMitigation(
    backend, mitigation="tensored", refresh=600, shots=2000, mit_pattern=[[0], [1]]
)
mit_mthree = ReadoutErrorMitigation(
    backend, mitigation="mthree", refresh=600, shots=2000, qubits=[0, 1]
)
expval_raw = PauliExpectationValue(ansatz, observable, backend=backend)
expval_tensored = PauliExpectationValue(ansatz, observable, backend=mit_tensored)
expval_mthree = PauliExpectationValue(ansatz, observable, backend=mit_mthree)
shots = 4000
print(f"w/o mitigation shots={shots}, result={expval_raw.evaluate([0, 1, 1, 2, 3, 5], shots=shots)}")
print(f"w/ tensored mitigation shots={shots}, result={expval_tensored.evaluate([0, 1, 1, 2, 3, 5], shots=shots)}")
print(f"w/ M3 mitigation shots={shots}, result={expval_mthree.evaluate([0, 1, 1, 2, 3, 5], shots=shots)}")

from qiskit.evaluators.expectation_value.expectation_value_gradient import (
    FiniteDiffGradient,
    ParameterShiftGradient,
)

parameters = [0, 1, 1, 2, 3, 5]

exact_expval = ExactExpectationValue(ansatz, observable, backend=AerSimulator())
exact_findiff = FiniteDiffGradient(exact_expval, 1e-8)
print(f"fin diff of exact {exact_findiff.evaluate(parameters).values}")

exact_expval = ExactExpectationValue(ansatz, observable, backend=AerSimulator())
exact_findiff = ParameterShiftGradient(exact_expval)
print(f"param shift of exact {exact_findiff.evaluate([0, 1, 1, 2, 3, 5]).values}")

shots = 2000
findiff = FiniteDiffGradient(expval_raw, 1e-1)
paramshift = ParameterShiftGradient(expval_raw)
print(f"fin diff w/o mit {findiff.evaluate(parameters, shots=shots).values}")
print(f"param shift w/o mit {paramshift.evaluate(parameters, shots=shots).values}")

findiff = FiniteDiffGradient(expval_mthree, 1e-1)
paramshift = ParameterShiftGradient(expval_mthree)
print(f"fin diff w/  mit {findiff.evaluate([0, 1, 1, 2, 3, 5], shots=shots).values}")
print(f"param shift w/  mit {paramshift.evaluate(parameters, shots=shots).values}")

from scipy.optimize import minimize

shots = 1000
expval = PauliExpectationValue(ansatz, observable, backend=AerSimulator(), append=True)
paramshift = ParameterShiftGradient(expval)
# this may take a long time...
result = minimize(
    lambda x: expval.evaluate(x, shots=shots, seed_simulator=123).value,
    np.zeros(6),
    jac=lambda x: paramshift.evaluate(x, shots=shots, seed_simulator=123).values,
)
print(result)

import matplotlib.pyplot as plt
from qiskit.evaluators.results import ExpectationValueResult

plt.figure(figsize=(12, 8))
y = [h.value for h in expval.history if isinstance(h, ExpectationValueResult)]
x = list(range(len(y)))
yerr = np.array(
    [tuple(abs(c - h.value) for c in h.confidence_interval) for h in expval.history if isinstance(h, ExpectationValueResult)]
).transpose()
plt.plot(y, color="blue")
plt.errorbar(x, y, yerr=yerr, capsize=5, fmt="o", markersize=8, ecolor="blue", markeredgecolor="blue", color="w")
plt.xlabel("Eval count")
plt.ylabel("Energy")
plt.show()