import argparse
from qiskit import QuantumCircuit, execute, Aer
from qiskit.visualization import plot_histogram
import numpy as np

def create_oracle(circuit, marked_item):
    """Creates an oracle that marks the specified item."""
    # For a 2-qubit search, the marked_item is a 2-bit string, e.g., '11'
    # The oracle flips the phase of the marked item
    if marked_item == '00':
        circuit.x([0, 1])
        circuit.cz(0, 1)
        circuit.x([0, 1])
    elif marked_item == '01':
        circuit.x(0)
        circuit.cz(0, 1)
        circuit.x(0)
    elif marked_item == '10':
        circuit.x(1)
        circuit.cz(0, 1)
        circuit.x(1)
    elif marked_item == '11':
        circuit.cz(0, 1)

def create_diffusion_operator(circuit, num_qubits):
    """Creates the Grover diffusion operator (amplitude amplification)."""
    circuit.h(range(num_qubits))
    circuit.z(range(num_qubits))
    circuit.cz(0, 1) # For 2 qubits
    circuit.h(range(num_qubits))

def create_grover_circuit(marked_item):
    """Creates a circuit for Grover's algorithm for a 2-qubit search."""
    num_qubits = 2
    circuit = QuantumCircuit(num_qubits, num_qubits)

    # 1. Initial superposition
    circuit.h(range(num_qubits))
    circuit.barrier()

    # Number of iterations
    num_iterations = int(np.floor(np.pi / 4 * np.sqrt(2**num_qubits)))

    for _ in range(num_iterations):
        # 2. Oracle
        create_oracle(circuit, marked_item)
        circuit.barrier()

        # 3. Diffusion operator
        create_diffusion_operator(circuit, num_qubits)
        circuit.barrier()

    # 4. Measurement
    circuit.measure(range(num_qubits), range(num_qubits))

    return circuit

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Grover's Algorithm for Database Search")
    parser.add_argument("--marked_item", type=str, default='11', help="The item to search for (e.g., '00', '01', '10', '11')")
    args = parser.parse_args()

    # Create the Grover's circuit
    grover_circuit = create_grover_circuit(args.marked_item)
    print(f"Grover's circuit for marked item: {args.marked_item}")
    print(grover_circuit)

    # Execute the circuit
    backend = Aer.get_backend('qasm_simulator')
    job = execute(grover_circuit, backend, shots=1024)
    result = job.result()
    counts = result.get_counts(grover_circuit)

    # Plot the results
    print("\nMeasurement Outcomes:")
    print(counts)
    plot_histogram(counts).show()