from qiskit.quantum_info import Operator
from qiskit import QuantumCircuit
import numpy as np

def phase_oracle(n, solutions, name = 'Oracle'):
    qc = QuantumCircuit(n, name=name)

    oracle_matrix = np.identity(2**n)
    for solution in solutions:
        oracle_matrix[solution, solution] = -1
    qc.unitary(Operator(oracle_matrix), range(n))
    
    return qc

def diffuser(n):
    
    qc = QuantumCircuit(n, name='Diffuser')
    qc.h(range(n))
    qc.append(phase_oracle(n, [0]), range(n))
    qc.h(range(n))
    
    return qc

def Grover(n, solutions):
    
    qc = QuantumCircuit(n, n)
    
    r = int(np.floor(np.pi/4*np.sqrt(2**n/len(solutions))))
    print(f'{n} qubits, basis states {solutions} marked, {r} rounds')
    
    qc.h(range(n))
    for _ in range(r):
        qc.append(phase_oracle(n, solutions), range(n))
        qc.append(diffuser(n), range(n))
        
    qc.measure(range(n), range(n))
    return qc

circuit = Grover(8, [1, 42, 120, 199, 250])
circuit.draw(output='text')

from qiskit import Aer, execute
simulator = Aer.get_backend('qasm_simulator')
counts = execute(circuit, backend=simulator).result().get_counts(circuit)
from qiskit.visualization import plot_histogram
plot_histogram(counts)

