# useful additional packages 
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

# useful math functions
from math import pi, cos, acos, sqrt

# importing the QISKit
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import available_backends, execute, register, get_backend

import sys, getpass
try:
    sys.path.append("../../") # go to parent dir
    import Qconfig
    qx_config = {
        "APItoken": Qconfig.APItoken,
        "url": Qconfig.config['url']}
    print('Qconfig loaded from %s.' % Qconfig.__file__)
except:
    APItoken = getpass.getpass('Please input your token and hit enter: ')
    qx_config = {
        "APItoken": APItoken,
        "url":"https://quantumexperience.ng.bluemix.net/api"}
    print('Qconfig.py not found in qiskit-tutorial directory; Qconfig loaded using user input.')
    
# import basic plot tools
from qiskit.tools.visualization import plot_histogram

M = 16                   # Maximum number of physical qubits available
numberOfCoins = 8        # This number should be up to M-1, where M is the number of qubits available
indexOfFalseCoin = 6     # This should be 0, 1, ..., numberOfCoins - 1, where we use Python indexing

if numberOfCoins < 4 or numberOfCoins >= M:
    raise Exception("Please use numberOfCoins between 4 and ", M-1)
if indexOfFalseCoin < 0 or indexOfFalseCoin >= numberOfCoins:
    raise Exception("indexOfFalseCoin must be between 0 and ", numberOfCoins-1)

# Creating registers
# numberOfCoins qubits for the binary query string and 1 qubit for working and recording the result of quantum balance
qr = QuantumRegister(numberOfCoins+1)
# for recording the measurement on qr
cr = ClassicalRegister(numberOfCoins+1)

circuitName = "QueryStateCircuit"
queryStateCircuit = QuantumCircuit(qr, cr)

N = numberOfCoins
# Create uniform superposition of all strings of length N
for i in range(N):
    queryStateCircuit.h(qr[i])

# Perform XOR(x) by applying CNOT gates sequentially from qr[0] to qr[N-1] and storing the result to qr[N]
for i in range(N):
    queryStateCircuit.cx(qr[i], qr[N])

# Measure qr[N] and store the result to cr[N]. We continue if cr[N] is zero, or repeat otherwise
queryStateCircuit.measure(qr[N], cr[N])

# we proceed to query the quantum beam balance if the value of cr[0]...cr[N] is all zero
# by preparing the Hadamard state of |1>, i.e., |0> - |1> at qr[N]
queryStateCircuit.x(qr[N]).c_if(cr, 0)
queryStateCircuit.h(qr[N]).c_if(cr, 0)

# we rewind the computation when cr[N] is not zero
for i in range(N):
    queryStateCircuit.h(qr[i]).c_if(cr, 2**N)


k = indexOfFalseCoin
# Apply the quantum beam balance on the desired superposition state (marked by cr equal to zero)
queryStateCircuit.cx(qr[k], qr[N]).c_if(cr, 0)

# Apply Hadamard transform on qr[0] ... qr[N-1]
for i in range(N):
    queryStateCircuit.h(qr[i]).c_if(cr, 0)

# Measure qr[0] ... qr[N-1]
for i in range(N):
    queryStateCircuit.measure(qr[i], cr[i])

backend = "local_qasm_simulator"

shots = 1 # We perform a one-shot experiment
results = execute(queryStateCircuit, backend=backend, shots=shots).result()
answer = results.get_counts()
for key in answer.keys():
    if key[0:1] == "1":
        raise Exception("Fail to create desired superposition of balanced query string. Please try again")
        
plot_histogram(answer)
from collections import Counter
for key in answer.keys():
    normalFlag, _ = Counter(key[1:]).most_common(1)[0] #get most common label
    for i in range(2,len(key)):
        if key[i] != normalFlag:
            print("False coin index is: ", len(key) - i - 1)

