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

import os
print(os.environ.get('QT_API'))

n=3
steps = 8

def counts_to_prob_2d(counts):
    states = list(counts.keys())
    state_counts = list(counts.values())
    nshots = sum(state_counts)
    n = int(len(states[0])/2)
    def sep_xy(states):
        # Separate x and y coordinates in state vector
        states_x = [s[:n] for s in states]
        states_y = [s[n:] for s in states]
        states_x = np.array([int(s[::-1],2) for s in states_x])
        states_y = np.array([int(s[::-1],2) for s in states_y])
        return states_x, states_y
    x,y = sep_xy(states)
    # Create array of probability values
    probabilities = np.zeros((2**n,2**n))
    probabilities[x,y] = state_counts
    probabilities /= nshots
    return probabilities

def increment_gate(circuit, qpos, qcoin):
    n = len(qpos)
    for i in range(n):
        circuit.mct(qcoin[:]+qpos[i+1:], qpos[i], None, mode='noancilla')

def decrement_gate(circuit, qpos, qcoin):
    n = len(qpos)
    for i in range(n):
        if i+1 < n: circuit.x(qpos[i+1:])
        circuit.mct(qcoin[:]+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()
    # y operations
    increment_gate(circuit, qpos[len(qpos)//2:], qcoin)
    circuit.x(qcoin[0])
    decrement_gate(circuit, qpos[len(qpos)//2:], qcoin)
    # x operations
    circuit.x(qcoin)
    increment_gate(circuit, qpos[:len(qpos)//2], qcoin)
    circuit.x(qcoin[0])
    decrement_gate(circuit, qpos[:len(qpos)//2:], qcoin)
    circuit.barrier()
    
def initialize_2D(circuit, n, pos):
    # convert position to binary
    x = "{:b}".format(pos[0])
    y = "{:b}".format(pos[1])
    for i in range(len(x)):
        if x[i]=='1': circuit.x((n-i)-1)
    for j in range(len(y)):
        if y[j]=='1': circuit.x((2*n-j)-1)
    return circuit
    
def run(steps):
    simulator = Aer.get_backend('qasm_simulator')
    qpos = QuantumRegister(2*n,'qc')
    qcoin = QuantumRegister(2,'qanc')
    cpos = ClassicalRegister(2*n,'cr')
    circuit = QuantumCircuit(qpos, qcoin, cpos)
    circuit = initialize_2D(circuit, n, (3,3))
    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 plotSquareGraph(propMat):
    colorarray = []
    # generate array of colors
    numProp = len(propMat)
    for idx in range(numProp):
        for jdx in range(numProp):
            colorarray.append([1-propMat[idx][jdx],1-propMat[idx][jdx],1])
    G = nx.grid_2d_graph(numProp, numProp)  # 5x5 grid
    # print the adjacency list
    for line in nx.generate_adjlist(G):
        pass
    # write edgelist to grid.edgelist
    nx.write_edgelist(G, path="grid.edgelist", delimiter=":")
    nx.draw(G,node_color = colorarray,cmap=plt.cm.Blues)
    plt.plot()

for i in range(steps+1):
    counts = run(i)
    prop = counts_to_prob_2d(counts)
    plt.close()
    plt.imshow(prop)
#     plotSquareGraph(prop)
    plt.pause(1/steps)



