#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

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")

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")

def cccz(circuit,qr):
    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])

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

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]
    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 inversion_about_average(circuit, register, n):
    """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')