import io
import numpy as np 
from numpy import pi

import pydot
from PIL import Image
import retworkx as rx

from qiskit import IBMQ, QuantumRegister, ClassicalRegister, QuantumCircuit, transpile
from qiskit.test.mock import FakeManhattan

def create_qaoa_circuit() -> QuantumCircuit:
    """ taken from IBMQ Experience """
    qreg_q = QuantumRegister(3, 'q')
    circuit = QuantumCircuit(qreg_q)

    circuit.h(qreg_q[0])
    circuit.h(qreg_q[1])
    circuit.cx(qreg_q[1], qreg_q[0])
    circuit.rz(-0.06942544294642758, qreg_q[0])
    circuit.cx(qreg_q[1], qreg_q[0])
    circuit.rz(-0.20827632883928274, qreg_q[0])
    circuit.rx(0.7530122052855767, qreg_q[0])
    circuit.h(qreg_q[2])
    circuit.cx(qreg_q[2], qreg_q[1])
    circuit.rz(1.2496579730356965, qreg_q[1])
    circuit.cx(qreg_q[2], qreg_q[1])
    circuit.rz(-0.20827632883928274, qreg_q[1])
    circuit.rx(0.7530122052855767, qreg_q[1])
    circuit.rx(0.7530122052855767, qreg_q[2])
    
    return circuit

# transpile circuit for ibmq_manhattan
device = FakeManhattan() 
conf = device.configuration()
print(conf.basis_gates)

qc = create_qaoa_circuit()
circuit_transpiled = transpile(qc, backend=device, optimization_level=3)
circuit_transpiled.draw(output='mpl')

graph = rx.PyGraph(multigraph=False)
coupling_map = [tuple(edge) for edge in conf.coupling_map]
graph.add_nodes_from([i for i in range(65)])
graph.add_edges_from_no_data(coupling_map)
circuit_nodes = [19, 20, 21]
neighbor_nodes = [12, 18, 22, 25]

def node_attr(node):
    kwargs = dict(style='filled', shape="circle", fontsize="14", fixedsize="true", width="0.4", height="0.4")
    if node in circuit_nodes:
        kwargs.update(fillcolor='green')
    elif node in neighbor_nodes:
        kwargs.update(fillcolor='lightblue', label="")
    else:
        kwargs.update(fillcolor='darkblue', label="")
    return kwargs

dot_str = graph.to_dot(node_attr)
dot = pydot.graph_from_dot_data(dot_str)[0]
png = dot.create_png(prog='neato')
Image.open(io.BytesIO(png))

