#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

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

def grover(n):
    qc = QuantumCircuit(3)
    qc.cz(0, 2)
    qc.cz(1, 2)
    oracle_ex3 = qc.to_gate()
    oracle_ex3.name = "U$_\omega$"
    grover_circuit = QuantumCircuit(n)
    grover_circuit = initialize_s(grover_circuit, [0,1,2])
    grover_circuit.append(oracle_ex3, [0,1,2])
    grover_circuit.append(diffuser(n), [0,1,2])
    grover_circuit.measure_all()
    return grover_circuit
    
def eval_counts_grover(counts):
    shots = 0
    p110 = 0
    p101 = 0
    for i in counts:
        shots += counts[i]
        if i == "110" or i == 6 or i =="6":
            p110 = counts[i]
        elif i == "101" or i == 5 or i == "5":
            p101 = counts[i]
    p110 /= shots
    p101 /= shots
    return p110 + p101 # probabillity of a correct state