from qiskit import QuantumCircuit, transpile, assemble
from qiskit_aer import AerSimulator 
from qiskit.visualization import plot_histogram


# Define the function to create the oracle
def create_oracle(n, marked_states):
    oracle = QuantumCircuit(n)
    
    # Apply Z gate to the qubits corresponding to the marked states
    for state in marked_states:
        oracle.z(state)
    
    # Apply controlled-Z gates for diffusion
    if n > 2:
        oracle.h(range(n))
        oracle.x(range(n))
        oracle.h(n-1)
        oracle.mct(list(range(n-1)), n-1)
        oracle.h(n-1)
        oracle.x(range(n))
        oracle.h(range(n))
    
    return oracle

# Define the function to create the Grover's diffusion operator
def create_diffusion(n):
    diffusion = QuantumCircuit(n)
    
    diffusion.h(range(n))
    diffusion.x(range(n))
    diffusion.h(n-1)
    diffusion.mct(list(range(n-1)), n-1)
    diffusion.h(n-1)
    diffusion.x(range(n))
    diffusion.h(range(n))
    
    return diffusion

# Define the function to apply Grover's algorithm
def grover_algorithm(oracle, diffusion, n, marked_states, num_iterations):
    grover_circuit = QuantumCircuit(n)

    # Apply Hadamard gates to all qubits
    grover_circuit.h(range(n))
    
    # Apply the Grover iterations
    for _ in range(num_iterations):
        grover_circuit.append(oracle, range(n))
        grover_circuit.append(diffusion, range(n))

    # Measure all qubits
    grover_circuit.measure_all()

    return grover_circuit

# Define the marked states and the number of qubits
n = 3
marked_states = [4, 7]  # Corresponding to |100⟩ and |111⟩
num_iterations = 1

# Create the oracle and the diffusion operator
oracle = create_oracle(n, marked_states)
diffusion = create_diffusion(n)

# Apply Grover's algorithm
grover_circuit = grover_algorithm(oracle, diffusion, n, marked_states, num_iterations)

# Draw the circuit
print(grover_circuit.draw())

# Simulate the circuit
simulator = AerSimulator()
grover_transpiled = transpile(grover_circuit, simulator)
qobj = assemble(grover_transpiled)
result = simulator.run(qobj).result()

# Plot the histogram of the result
counts = result.get_counts()
print(counts)
plot_histogram(counts)








