# General imports
import numpy as np

# SciPy minimizer routine
from scipy.optimize import minimize

# rustworkx graph library
import rustworkx as rx
from rustworkx.visualization import mpl_draw

# Pre-defined ansatz circuit, operator class and visualization tools
from qiskit import QuantumCircuit
from qiskit.circuit.library import QAOAAnsatz
from qiskit.quantum_info import SparsePauliOp
from qiskit.visualization import plot_distribution

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

num_nodes = 5
# The edge syntax is (start, end, weight)
edges = [(0, 1, 1), (0, 2, 1), (0, 3, 1), (0, 4, 1)]

G = rx.PyGraph()
G.add_nodes_from(range(num_nodes))
G.add_edges_from(edges)
mpl_draw(
    G, pos=rx.bipartite_layout(G, {0}), with_labels=True, node_color="#EE5396", font_color="#F4F4F4"
)

# Problem to Hamiltonian operator
hamiltonian = SparsePauliOp.from_list([("IIIZZ", 1), ("IIZIZ", 1), ("IZIIZ", 1), ("ZIIIZ", 1)])
# QAOA ansatz circuit
ansatz = QAOAAnsatz(hamiltonian, reps=2)

ansatz.decompose(reps=3).draw(output="mpl", style="iqp")

ansatz.decompose().draw(output="mpl", style="iqp")

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(ansatz_isa.layout)
hamiltonian_isa

def cost_func(
    params: np.ndarray,
    ansatz: QuantumCircuit,
    hamiltonian: SparsePauliOp,
    estimator: BaseEstimatorV2,
) -> 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.

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

    return cost

# Configure estimator
estimator = Estimator(backend=backend)
estimator.options.default_shots = 10_000

# Configure sampler
sampler = Sampler(backend=backend)
sampler.options.default_shots = 10_000

x0 = 2 * np.pi * np.random.rand(ansatz_isa.num_parameters)

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

res

# Assign solution parameters to ansatz
qc = ansatz.assign_parameters(res.x)
# Add measurements to our circuit
qc.measure_all()
qc_isa = pm.run(qc)
qc_isa.draw(output="mpl", idle_wires=False, style="iqp")

result = sampler.run([qc_isa]).result()
samp_dist = result[0].data.meas.get_counts()

plot_distribution(samp_dist, figsize=(15, 5))

# Reverse 11110 to convert quantum result to classical solution
solution = [0, 1, 1, 1, 1]

mpl_draw(
    G,
    pos=rx.bipartite_layout(G, {0}),
    with_labels=True,
    node_color=["#EE5396" if kk else "#0F62FE" for kk in solution],
    font_color="#F4F4F4",
)

import qiskit

qiskit.version.get_version_info()

import qiskit_aer

qiskit_aer.version.get_version_info()