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

# importing Qiskit
from qiskit import IBMQ, Aer, assemble, transpile
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit.providers.ibmq import least_busy

# import basic plot tools
from qiskit.visualization import plot_histogram
from qiskit_textbook.problems import grover_problem_oracle

def initialize_s(qc, qubits):
    """Apply a H-gate to 'qubits' in qc"""
    for q in qubits:
        qc.h(q)
    return qc

def diffuser(nqubits):
    qc = QuantumCircuit(nqubits)
    # Apply transformation |s> -> |00..0> (H-gates)
    for qubit in range(nqubits):
        qc.h(qubit)
    # Apply transformation |00..0> -> |11..1> (X-gates)
    for qubit in range(nqubits):
        qc.x(qubit)
    # Do multi-controlled-Z gate
    qc.h(nqubits-1)
    qc.mct(list(range(nqubits-1)), nqubits-1)  # multi-controlled-toffoli
    qc.h(nqubits-1)
    # Apply transformation |11..1> -> |00..0>
    for qubit in range(nqubits):
        qc.x(qubit)
    # Apply transformation |00..0> -> |s>
    for qubit in range(nqubits):
        qc.h(qubit)
    # We will return the diffuser as a gate
    U_s = qc.to_gate()
    U_s.name = "U$_s$"
    return U_s

n = 4
qc = QuantumCircuit(n)
oracle1 = grover_problem_oracle(4, variant=2, print_solutions=True)
qc = initialize_s(qc, [0,1,2,3])
# we run two iterations here
qc.append(oracle1, [0,1,2,3])
qc.append(diffuser(n), [0,1,2,3])
qc.append(oracle1, [0,1,2,3])
qc.append(diffuser(n), [0,1,2,3])
qc.measure_all()
qc.draw()

aer_sim = Aer.get_backend('aer_simulator')
transpiled_qc = transpile(qc, aer_sim)
qobj = assemble(transpiled_qc)
results = aer_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

# Verification of Solution
n = 4
oracle = grover_problem_oracle(n, variant=2, print_solutions = True)  # 0th variant of oracle, with n qubits
qc = QuantumCircuit(n)
qc.append(oracle, [0,1,2,3])


n = 4
qc = QuantumCircuit(n)
oracle = grover_problem_oracle(4, variant=1, print_solutions=True)
qc = initialize_s(qc, [0,1,2,3])
# we run two iterations here
qc.append(oracle, [0,1,2,3])
qc.append(diffuser(n), [0,1,2,3])
qc.measure_all()
qc.draw()

aer_sim = Aer.get_backend('aer_simulator')
transpiled_qc = transpile(qc, aer_sim)
qobj = assemble(transpiled_qc)
results = aer_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

n = 4
qc = QuantumCircuit(n)
oracle = grover_problem_oracle(4, variant=1, print_solutions=True)
qc = initialize_s(qc, [0,1,2,3])
# we run two iterations here
qc.append(oracle, [0,1,2,3])
qc.append(diffuser(n), [0,1,2,3])
qc.append(oracle, [0,1,2,3])
qc.append(diffuser(n), [0,1,2,3])
qc.measure_all()
qc.draw()

aer_sim = Aer.get_backend('aer_simulator')
transpiled_qc = transpile(qc, aer_sim)
qobj = assemble(transpiled_qc)
results = aer_sim.run(qobj).result()
counts = results.get_counts()
plot_histogram(counts)

# define the number of qubits needed n
# define your own oracle and the nummber of iterations you need
# this code makes use of just the diffuser function

def grover_oracle(oracle, iterations):
    # applying h gates 
    for i in range(n):
        qc.h(i)
    #applying the oracle specified
    a = []
    for i in range(n):
        a.append(i)
    
    for i in range(iterations):
        qc.append(oracle, a)
        qc.append(diffuser(n), a)
    
    qc.measure_all()
    return qc

n = 4
iterations = 2
oracle = grover_problem_oracle(n, variant=1, print_solutions=True)

qc = QuantumCircuit(n, n)
qc = grover_oracle(oracle, iterations)
qc.draw()

