# General imports
import numpy as np

# SciPy minimizer routine
from scipy.optimize import minimize

# Plotting functions
import matplotlib.pyplot as plt

# Pre-defined ansatz circuit and operator class for Hamiltonian
from qiskit import QuantumCircuit
from qiskit.circuit.library import EfficientSU2
from qiskit.quantum_info import SparsePauliOp

# Qiskit Primitive imports
from qiskit.primitives import BaseEstimatorV2, BackendEstimatorV2 as Estimator
from qiskit.providers.fake_provider import GenericBackendV2

hamiltonian = SparsePauliOp.from_list(
    [("YZ", 0.3980), ("ZI", -0.3980), ("ZZ", -0.0113), ("XX", 0.1810)]
)

ansatz = EfficientSU2(hamiltonian.num_qubits)
ansatz.decompose().draw("mpl", style="iqp")

num_params = ansatz.num_parameters
num_params

from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

backend = GenericBackendV2(hamiltonian.num_qubits)
target = backend.target
pm = generate_preset_pass_manager(target=target, optimization_level=3)

ansatz_isa = pm.run(ansatz)

ansatz_isa.draw(output="mpl", idle_wires=False, style="iqp")

hamiltonian_isa = hamiltonian.apply_layout(layout=ansatz_isa.layout)

def cost_func(
    params: np.ndarray,
    ansatz: QuantumCircuit,
    hamiltonian: SparsePauliOp,
    estimator: BaseEstimatorV2,
    history: dict,
) -> float:
    """Return estimate of energy from estimator

    Parameters:
        params: Array of ansatz parameters.
        ansatz: Parameterized ansatz circuit.
        hamiltonian: Operator representation of Hamiltonian.
        estimator: Estimator primitive instance.
        history: Dictionary for storing intermediate results.

    Returns:
        Energy estimate.
    """
    pub = (ansatz, [hamiltonian], [params])
    result = estimator.run(pubs=[pub]).result()
    energy = result[0].data.evs[0]

    history["iters"] += 1
    history["prev_vector"] = params
    history["cost_history"].append(energy)
    print(f"Iters. done: {history['iters']} [Current cost: {energy}]")

    return energy

history = {
    "prev_vector": None,
    "iters": 0,
    "cost_history": [],
}

x0 = 2 * np.pi * np.random.random(num_params)

x0

estimator = Estimator(backend=backend)
estimator.options.default_shots = 10000

res = minimize(
    cost_func,
    x0,
    args=(ansatz_isa, hamiltonian_isa, estimator, history),
    method="cobyla",
)

res

all(history["prev_vector"] == res.x)

history["iters"] == res.nfev

fig, ax = plt.subplots()
ax.plot(range(history["iters"]), history["cost_history"])
ax.set_xlabel("Iterations")
ax.set_ylabel("Cost")
plt.draw()

import qiskit

qiskit.version.get_version_info()

import qiskit_aer

qiskit_aer.version.get_version_info()