from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from qiskit import IBMQ, Aer, execute
from qiskit.tools.visualization import plot_histogram

## distance_black_box

distances = {
    "32": 3,
    "31": 2,
    "30": 4,
    "21": 7,
    "20": 6,
    "10": 5,
}

def dist_single():
    qr = QuantumRegister(2)
    qr_target = QuantumRegister(5)
    qc = QuantumCircuit(qr, qr_target, name='dist_single')
    
    for edge in distances:
        if edge[0] == '3':
            node = format(int(edge[1]), '02b')[::-1]
            d_bin = format(distances[edge], '02b')[::-1]
            
            for idx in range(len(node)):
                if node[idx] == '0':
                    qc.x(qr[idx])
            
            for idx in range(len(d_bin)):
                if d_bin[idx] == '1':
                    qc.ccx(qr[0], qr[1], qr_target[idx])
            
            for idx in range(len(node)):
                if node[idx] == '0':
                    qc.x(qr[idx])
    
    return qc

def dist():
    qr1 = QuantumRegister(2)
    qr2 = QuantumRegister(2)
    qr_target = QuantumRegister(5)
    qr_anc = QuantumRegister(2)
    qc = QuantumCircuit(qr1, qr2, qr_target, qr_anc, name='dist')
    
    for edge in distances:
        if edge[0] != '3':
            # convert to binaries
            node1 = format(int(edge[0]), '02b')[::-1]
            node2 = format(int(edge[1]), '02b')[::-1]
            d_bin = format(distances[edge], '02b')[::-1]

            for idx in range(len(node1)): # assume node1 and node2 have the same length
                if node1[idx] == '0':
                    qc.x(qr1[idx])
                    
            for idx in range(len(node2)):
                if node2[idx] == '0':
                    qc.x(qr2[idx])

            for idx in range(len(d_bin)):
                if d_bin[idx] == '1':
                    qc.mct(qr1[:]+qr2[:], qr_target[idx], qr_anc)
                    
            for idx in range(len(node2)): # invert back
                if node2[idx] == '0':
                    qc.x(qr2[idx])

            for idx in range(len(node1)):
                if node1[idx] == '0':
                    qc.x(qr1[idx])
                    
    return qc

qr1 = QuantumRegister(2) # node reg.
qr2 = QuantumRegister(2) # node reg.
qr_target = QuantumRegister(5) # distance reg.
qr_anc = QuantumRegister(2)
c = ClassicalRegister(5)
qc = QuantumCircuit(qr1, qr2, qr_target, qr_anc, c)

qc.x(qr1[0])
# qc.x(qr2[0])
qc.append(dist(), qr1[:] + qr2[:] + qr_target[:] + qr_anc[:])
qc.measure(qr_target, c)

# qc.draw(output='mpl')

backend = Aer.get_backend('qasm_simulator')
job = execute(qc, backend, shots=1024)
counts = job.result().get_counts()

plot_histogram(counts)

format(1, '02b')