import rustworkx as rx
from rustworkx.visualization import mpl_draw

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

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='tab:purple', font_color='white'
)

import numpy as np
from qiskit.quantum_info import SparsePauliOp

# Problem to Hamiltonian operator
hamiltonian = SparsePauliOp.from_list([("IIIZZ", 1), ("IIZIZ", 1), ("IZIIZ", 1), ("ZIIIZ", 1)])

print(hamiltonian)

from qiskit.circuit.library import QAOAAnsatz

# QAOA ansatz circuit
ansatz = QAOAAnsatz(hamiltonian, reps=2)

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

ansatz.decompose(reps=1).draw(output="mpl", style="clifford")

from qiskit.primitives import Estimator

estimator = Estimator()

# If we want to run it on the real device, use these lines instead:

#from qiskit_ibm_runtime import QiskitRuntimeService
#from qiskit_ibm_runtime import Estimator, Sampler, Session, Options

#service = QiskitRuntimeService(channel="ibm_quantum", token=<YOUR TOKEN>)
#backend = service.least_busy(operational=True, simulator=False)
#print(backend.name)

#options = Options()
#options.resilience_level = 1
#options.optimization_level = 3
 
#estimator = Estimator(backend, options=options)

def cost_func(params, ansatz, hamiltonian, estimator):
    
    """Return estimate of energy from estimator

    Parameters:
        params (ndarray): Array of ansatz parameters
        ansatz (QuantumCircuit): Parameterized ansatz circuit
        hamiltonian (SparsePauliOp): Operator representation of Hamiltonian
        estimator (Estimator): Estimator primitive instance

    Returns:
        float: Energy estimate
    """
    
    cost = estimator.run(ansatz, hamiltonian, parameter_values=params).result().values[0]
    return cost

from scipy.optimize import minimize

# Randomly initialize the parameters of the ansatz
x0 = 2 * np.pi * np.random.rand(ansatz.num_parameters)

# Run the variational algorithm
res = minimize(cost_func, x0, args=(ansatz, hamiltonian, estimator),
               method="COBYLA")

print(res)

from qiskit.primitives import Sampler

sampler = Sampler()

# If instead we want to run it on the device 
# we need the sampler from runtime
# sampler = Sampler(backend, options=options)

# Assign solution parameters to ansatz
qc = ansatz.assign_parameters(res.x)
qc.measure_all()

# Draw circuit with optimal parameters
#qc_ibm.draw(output="mpl", idle_wires=False, style="iqp")

from qiskit.visualization import plot_distribution

# Sample ansatz at optimal parameters
samp_dist = sampler.run(qc).result().quasi_dists[0]
plot_distribution(samp_dist.binary_probabilities(), figsize=(15, 5))