import numpy as np
import time
import networkx as nx
import matplotlib.pyplot as plt
import random
from qiskit import (QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer)
%matplotlib qt

steps = 5
n=4

def counts_to_prob_1d(counts):
    # Convert histogram counts to probability vector of size 1 x 2^n
    states = list(counts.keys())
    state_counts = list(counts.values())
    nshots = sum(state_counts)
    # Convert binary to integer index, reversing order as consistent w qiskit convention
    states_x = [int(s[::-1],2) for s in states]
    # Create array of probability values
    probabilities = np.zeros(2**n)
    probabilities[states_x] = state_counts
    probabilities /= nshots
    return probabilities

def increment_gate(circuit, qpos, qcoin):
    n = len(qpos)
    for i in range(n):
        circuit.mct([qcoin[0]]+qpos[i+1:], qpos[i], None, mode='noancilla')
        
def decrement_gate(circuit, qpos, qcoin):
    n = len(qpos)
    circuit.x(qcoin[0])
    for i in range(n):
        if i+1 < n: circuit.x(qpos[i+1:])
        circuit.mct([qcoin[0]]+qpos[i+1:], qpos[i], None, mode='noancilla')
        if i+1 < n: circuit.x(qpos[i+1:])
            
def step(circuit, qpos, qcoin):
    circuit.h(qcoin)
    circuit.barrier()
    increment_gate(circuit, qpos, qcoin)
    circuit.barrier()
    decrement_gate(circuit, qpos, qcoin)
    circuit.barrier()

def run(steps):
    simulator = Aer.get_backend('qasm_simulator')
    qpos = QuantumRegister(n,'qc')
    qcoin = QuantumRegister(1,'qanc')
    cpos = ClassicalRegister(n,'cr')
    circuit = QuantumCircuit(qpos, qcoin, cpos)
    for i in range(steps):
        step(circuit, qpos, qcoin)
    # # Map the quantum measurement to the classical bits
    circuit.measure(qpos,cpos)
    # # Execute the circuit on the qasm simulator
    job = execute(circuit, simulator, shots=1000)
    # # Grab results from the job
    result = job.result()
    # # Returns counts
    counts = result.get_counts(circuit)
    return counts

seed = 12
random.seed(seed)
np.random.seed(seed)
def plotCircleGraph(propabilities):
    G = nx.Graph()
    colorarray = []
    # generate array of colors
    numProp = len(propabilities)
    for idx in range(numProp):
        colorarray.append([1-prop[idx],1-prop[idx],1])
    # generate graph
    for idx in range(numProp-1):
        G.add_edge(idx, idx+1)
    # add last edge
    G.add_edge(0,numProp-1)
    nx.draw(G,pos=nx.circular_layout(G),node_color = colorarray,cmap=plt.cm.Blues)

for i in range(steps+1):
    counts = run(i)
    prop = counts_to_prob_1d(counts)
    plotCircleGraph(prop)
    plt.pause(0.3)
    plt.close()



