from qiskit import QuantumCircuit
from qiskit.circuit import QuantumRegister, ClassicalRegister
import math


# N-fold cx using ccx gate.
# len(ancs) = len(qs) -2
# len(qs) > 2
def nfold_cx(qs, q, ancs, circ):
    circ.ccx(qs[0] , qs[1], ancs[0])
    k = len(ancs)
    for i in range(1, k):
        circ.ccx(qs[i+1], ancs[i-1], ancs[i])
    circ.ccx(ancs[k-1],qs[k+1], q)
    for j in reversed(range(1, k)):
        circ.ccx(qs[j+1], ancs[j-1], ancs[j])
    circ.ccx(qs[0] , qs[1], ancs[0])        
    return circ

# Testing nfold_cx
# qs = QuantumRegister(5)
# q = QuantumRegister(1)
# ancs = QuantumRegister(3)
# circ = QuantumCircuit(qs,q,ancs)

# print(nfold_cx(qs, q, ancs, circ))

# Note that len(ancs) = len(bs)-3
def reflection(bs, ancs, circ):
    circ.h(bs)
    circ.x(bs)
    n = len(bs)
    circ.barrier()
    # note that alternatively, bs[-1] can be used 
    circ.h(bs[n-1])
    circ = nfold_cx(bs[0:n-1], bs[n-1], ancs, circ)
    circ.h(bs[n-1])
    circ.barrier()
    circ.x(bs)
    circ.h(bs)
    return circ

# Testing reflection. 
# bs = QuantumRegister(5)
# ancs = QuantumRegister(2)
# circ = QuantumCircuit(bs,ancs)

# print(reflection(bs, ancs, circ))

# An example of U_f that f(00111) = 1, and f(otherwise) = 0. 
# len(qs) = 5
# len(ancs) = 3
def u_f(qs, q, ancs, circ):
    circ.x(qs[0])
    circ.x(qs[1])
    nfold_cx(qs, q, ancs, circ)
    circ.x(qs[0])
    circ.x(qs[1])
    return circ

# Testing U_f
# qs = QuantumRegister(5)
# q = QuantumRegister(1)
# ancs = QuantumRegister(3)
# circ = QuantumCircuit(qs,q, ancs)

# print(u_f(qs, q, ancs, circ))



def grover():
    qs = QuantumRegister(5)
    q = QuantumRegister(1)
    # Ancillas for U_f, and we want to reuse
    # these ancillas for all the subsequent calls of U_f
    ancs_uf = QuantumRegister(3)

    # Ancillas for reflection, and we want to reuse
    # these ancillas for all the subsequent calls of reflection
    ancs_ref = QuantumRegister(2)

    # Classical register for storing the measurement result
    cls = ClassicalRegister(5)
    circ = QuantumCircuit(qs,q, ancs_uf, ancs_ref, cls)    

    # Create the initial 2^n search space
    circ.h(qs)

    # Create a minus state for phase kickback
    circ.x(q)
    circ.h(q)

    # Grover iteration, u_f followed by reflection about mean.
    # In this case we iterate it sqrt(2^5) ~ 5/6 times. 
    for i in range(5):
        circ.barrier()
        u_f(qs, q, ancs_uf, circ)
        circ.barrier()
        reflection(qs, ancs_ref, circ)
    circ.measure(qs, cls)
    return circ

qc = grover()

print(qc)

# Code for running on simulator

'''
from qiskit_ibm_runtime import QiskitRuntimeService, Sampler


service = QiskitRuntimeService()
backend = service.get_backend("simulator_mps")
grover = backend.run(qc)
grover_result = grover.result()
print(grover_result.get_counts())
'''