import qiskit
qiskit.__qiskit_version__

#initialization
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# importing Qiskit
from qiskit import BasicAer, IBMQ
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, execute
from qiskit.compiler import transpile
from qiskit.tools.monitor import job_monitor

# import basic plot tools
from qiskit.tools.visualization import plot_histogram

# Load the saved IBMQ accounts
IBMQ.load_account()

s = "010101" # the hidden bitstring

assert 1 < len(s) < 20, "The length of s must be between 2 and 19"

for c in s:
    assert c == "0" or c == "1", "s must be a bitstring of '0' and '1'"
    
n = len(s) #the length of the bitstring

# Step 1
# Creating registers
# qubits for querying the oracle and recording its output
qr = QuantumRegister(2*n)
# for recording the measurement on the first register of qr
cr = ClassicalRegister(n)

circuitName = "Simon"
simonCircuit = QuantumCircuit(qr, cr)

# Step 2
# Apply Hadamard gates before querying the oracle
for i in range(n):
    simonCircuit.h(qr[i])

# Apply barrier to mark the beginning of the blackbox function
simonCircuit.barrier()

# Step 3 query the blackbox function
# copy the content of the first register to the second register
for i in range(n):
    simonCircuit.cx(qr[i], qr[n+i])

# get the least index j such that s_j is "1"
j = -1
for i, c in enumerate(s):
    if c == "1":
        j = i
        break

# Creating 1-to-1 or 2-to-1 mapping with the j-th qubit of x as control to XOR the second register with s
for i, c in enumerate(s):
    if c == "1" and j >= 0:
        simonCircuit.cx(qr[j], qr[n+i]) #the i-th qubit is flipped if s_i is 1

# get random permutation of n qubits
perm = list(np.random.permutation(n))

#initial position
init = list(range(n))
i = 0
while i < n:
    if init[i] != perm[i]:
        k = perm.index(init[i])
        simonCircuit.swap(qr[n+i], qr[n+k]) #swap qubits
        init[i], init[k] = init[k], init[i] #marked swapped qubits
    else:
        i += 1
        
# randomly flip the qubit
for i in range(n):
    if np.random.random() > 0.5:
        simonCircuit.x(qr[n+i])

# Apply the barrier to mark the end of the blackbox function
simonCircuit.barrier()

# Step 4 apply Hadamard gates to the first register
for i in range(n):
    simonCircuit.h(qr[i])
    
# Step 5 perform measurement on the first register
for i in range(n):
    simonCircuit.measure(qr[i], cr[i])

#draw the circuit
simonCircuit.draw(output='mpl')

# use local simulator
backend = BasicAer.get_backend("qasm_simulator")
# the number of shots is twice the length of the bitstring
shots = 2*n
job = execute(simonCircuit, backend=backend, shots=shots)
answer = job.result().get_counts()

plot_histogram(answer)

# Post-processing step
# Constructing the system of linear equations Y s = 0
# By k[::-1], we reverse the order of the bitstring
lAnswer = [ (k[::-1],v) for k,v in answer.items() if k != "0"*n  ] #excluding the trivial all-zero
#Sort the basis by their probabilities
lAnswer.sort(key = lambda x: x[1], reverse=True)

Y = []
for k, v in lAnswer:
    Y.append( [ int(c) for c in k ] )


#import tools from sympy
from sympy import Matrix, pprint, MatrixSymbol, expand, mod_inverse


Y = Matrix(Y)

#pprint(Y)

#Perform Gaussian elimination on Y
Y_transformed = Y.rref(iszerofunc=lambda x: x % 2==0) # linear algebra on GF(2) 

#to convert rational and negatives in rref of linear algebra on GF(2)
def mod(x,modulus):
    numer, denom = x.as_numer_denom()
    return numer*mod_inverse(denom,modulus) % modulus

Y_new = Y_transformed[0].applyfunc(lambda x: mod(x,2)) #must takecare of negatives and fractional values
#pprint(Y_new)

print("The hidden bistring s[ 0 ], s[ 1 ]....s[",n-1,"] is the one satisfying the following system of linear equations:")
rows, cols = Y_new.shape
for r in range(rows):
    Yr = [ "s[ "+str(i)+" ]" for i, v in enumerate(list(Y_new[r,:])) if v == 1 ]
    if len(Yr) > 0:
        tStr = " + ".join(Yr)
        print(tStr, "= 0")


#Use one of the available backends 
backend = IBMQ.get_backend("ibmq_16_melbourne")

# show the status of the backend
print("Status of", backend, "is", backend.status())

shots = 10*n #run more experiments to be certain
max_credits = 3          # Maximum number of credits to spend on executions. 

simonCompiled = transpile(simonCircuit, backend=backend, optimization_level=1)
job_exp = execute(simonCompiled, backend=backend, shots=shots, max_credits=max_credits)
job_monitor(job_exp)

results = job_exp.result()
answer = results.get_counts(simonCircuit)

plot_histogram(answer)

# Post-processing step
# Constructing the system of linear equations Y s = 0
# By k[::-1], we reverse the order of the bitstring
lAnswer = [ (k[::-1][:n],v) for k,v in answer.items() ] #excluding the qubits that are not part of the inputs
#Sort the basis by their probabilities
lAnswer.sort(key = lambda x: x[1], reverse=True)

Y = []
for k, v in lAnswer:
    Y.append( [ int(c) for c in k ] )

Y = Matrix(Y)

#Perform Gaussian elimination on Y
Y_transformed = Y.rref(iszerofunc=lambda x: x % 2==0) # linear algebra on GF(2)
Y_new = Y_transformed[0].applyfunc(lambda x: mod(x,2)) #must takecare of negatives and fractional values
#pprint(Y_new)

print("The hidden bistring s[ 0 ], s[ 1 ]....s[",n-1,"] is the one satisfying the following system of linear equations:")
rows, cols = Y_new.shape
for r in range(rows):
    Yr = [ "s[ "+str(i)+" ]" for i, v in enumerate(list(Y_new[r,:])) if v == 1 ]
    if len(Yr) > 0:
        tStr = " + ".join(Yr)
        print(tStr, "= 0")