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

# importing Qiskit
from qiskit import IBMQ, BasicAer
from qiskit.providers.ibmq import least_busy
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, execute

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

# importing Qiskit
from qiskit import IBMQ, Aer, QuantumCircuit, ClassicalRegister, QuantumRegister, execute
from qiskit.providers.ibmq import least_busy
from qiskit.quantum_info import Statevector

def cz(circuit,qr):
    circuit.h(qr[1])
    circuit.cx(qr[0],qr[1])
    circuit.h(qr[1])
    
qr2 = QuantumRegister(2)
czCircuit = QuantumCircuit(qr2)
cz(czCircuit,qr2)
czCircuit.draw(output="mpl")

#Since Z=HXH and X=HZH
#This is identical to a CZ (controlled Z)

def ccz(circuit,qr):
    circuit.h(qr[2])
    # ccx is the Toffoli gate 
    circuit.ccx(qr[0], qr[1], qr[2])
    circuit.h(qr[2])
    
qr3 = QuantumRegister(3)
cczCircuit = QuantumCircuit(qr3)
ccz(cczCircuit,qr3)
cczCircuit.draw(output="mpl")

#We are going to check numerically whether this is the CCCZ gate
def cccz(circuit,qr):
    pi = math.pi
    circuit.x(qr[0]) #we change the inputvalue by applying an X gate at the beggining
    circuit.x(qr[1]) #second qubit to 1
    circuit.x(qr[2]) #third qubit to 1
    circuit.x(qr[3]) #fourth qubit to 1
    circuit.cu1(pi/4,qr[0],qr[3])
    circuit.cx(qr[0], qr[1])
    circuit.cu1(-pi/4, qr[1],qr[3])
    circuit.cx(qr[0], qr[1])
    circuit.cu1(pi/4, qr[1],qr[3])
    circuit.cx(qr[1], qr[2])
    circuit.cu1(-pi/4, qr[2],qr[3])
    circuit.cx(qr[0], qr[2])
    circuit.cu1(pi/4, qr[2],qr[3])
    circuit.cx(qr[1], qr[2])
    circuit.cu1(-pi/4, qr[2],qr[3])
    circuit.cx(qr[0], qr[2])
    circuit.cu1(pi/4, qr[2],qr[3])

qr4 = QuantumRegister(4)
ccczCircuit = QuantumCircuit(qr4)
cccz(ccczCircuit,qr4)
ccczCircuit.draw(output="mpl")

from qiskit.visualization import plot_state_qsphere
sv_sim = Aer.get_backend('statevector_simulator')
job_sim = execute(ccczCircuit, sv_sim)
statevec = job_sim.result().get_statevector()
from qiskit_textbook.tools import vector2latex
#vector2latex(statevec, pretext="|\\psi\\rangle =")
plot_state_qsphere(statevec)

def n_controlled_Z(circuit, qr):
    """Implement a Z gate with multiple controls"""
    if (len(qr) > 4):
        raise ValueError('The controlled Z with more than 3 controls is not implemented')
    # This is the case n = 2 (1 control + 1 target qubit)
    elif (len(qr) == 2):
        cz(circuit,qr)
    # This is the case n = 3 (2 control + 1 target qubit)
    elif (len(qr) == 3):
        ccz(circuit,qr)
        # This is the case n = 4 (3 control + 1 target qubit)
    elif (len(qr) == 4):
        cccz(circuit,qr)

def phase_oracle(circuit,qr,element):
    # element is an array that defines the searched element, for example, element =  [0,1,0,1]
    circuit.x(qr[0]) #we change the inputvalue by applying an X gate at the beggining
#    circuit.x(qr[1]) #for the second qubit
    circuit.x(qr[2]) 
#    circuit.x(qr[3]) 
    circuit.barrier() #in this part, we are changing the input to the oracle circuit
    n = len(element)
    for j,x in enumerate(element):
        if (x == 0):
            circuit.x(qr[j])
    n_controlled_Z(circuit,qr) 
    for j,x in enumerate(element):
        if (x == 0):
            circuit.x(qr[j])

element=[0,1,0,1]
qr4 = QuantumRegister(4)
oracleCircuit = QuantumCircuit(qr4)
phase_oracle(oracleCircuit,qr4, element)
oracleCircuit.draw(output="mpl")

oracleCircuit.measure_all()

qasm_simulator = Aer.get_backend('statevector_simulator')
job_sim = execute(oracleCircuit, sv_sim)
statevec = job_sim.result().get_statevector()
from qiskit_textbook.tools import vector2latex
vector2latex(statevec, pretext="|\\psi\\rangle =")

def inversion_about_average(circuit, register):
    """Apply inversion about the average step of Grover's algorithm."""
    circuit.h(register)
    circuit.x(register)
    n_controlled_Z(circuit, qr)
    circuit.x(register)
    circuit.h(register)

#Let's plot it!
qr = QuantumRegister(4)
qAverage = QuantumCircuit(qr)
inversion_about_average(qAverage, qr, 4)
qAverage.draw(output='mpl')

#implementation of the whole Grover Search Algorithm
def cccz(circuit,qr): #first I define again the oracle, since before it was personalized
    pi = math.pi
    circuit.cu1(pi/4,qr[0],qr[3])
    circuit.cx(qr[0], qr[1])
    circuit.cu1(-pi/4, qr[1],qr[3])
    circuit.cx(qr[0], qr[1])
    circuit.cu1(pi/4, qr[1],qr[3])
    circuit.cx(qr[1], qr[2])
    circuit.cu1(-pi/4, qr[2],qr[3])
    circuit.cx(qr[0], qr[2])
    circuit.cu1(pi/4, qr[2],qr[3])
    circuit.cx(qr[1], qr[2])
    circuit.cu1(-pi/4, qr[2],qr[3])
    circuit.cx(qr[0], qr[2])
    circuit.cu1(pi/4, qr[2],qr[3])
    
def phase_oracle(circuit,qr,element): #first I define again the oracle, since before it was personalized
    circuit.barrier() #in this part, we are changing the input to the oracle circuit
    n = len(element)
    for j,x in enumerate(element):
        if (x == 0):
            circuit.x(qr[j])
    n_controlled_Z(circuit,qr) 
    for j,x in enumerate(element):
        if (x == 0):
            circuit.x(qr[j])
            
def n_controlled_Z(circuit, qr):
    """Implement a Z gate with multiple controls"""
    if (len(qr) > 4):
        raise ValueError('The controlled Z with more than 3 controls is not implemented')
    # This is the case n = 2 (1 control + 1 target qubit)
    elif (len(qr) == 2):
        cz(circuit,qr)
    # This is the case n = 3 (2 control + 1 target qubit)
    elif (len(qr) == 3):
        ccz(circuit,qr)
        # This is the case n = 4 (3 control + 1 target qubit)
    elif (len(qr) == 4):
        cccz(circuit,qr)

def grover_algorithm(Circuit,R,qr,element):
    for i in range(1,R):
        phase_oracle(Circuit,qr,element)
        inversion_about_average(Circuit,qr)

#n=2 circuit
qr2=QuantumRegister(2)
print(len(qr2))
cr2=ClassicalRegister(2)
circuit2=QuantumCircuit(qr2,cr2)
circuit2.h(0)
circuit2.h(1)#quantum superposition of the initial states
element2 = [0,1]
grover_algorithm(circuit2,2,qr2,[0,1])