import gradio as gr
from qiskit import QuantumCircuit, transpile
from qiskit_aer import Aer
from qiskit.visualization import plot_histogram
import matplotlib.pyplot as plt

def run_quantum_circuit(apply_x_gate: bool):
    # Setup Quantum Circuit
    qc = QuantumCircuit(2)

    if apply_x_gate:
        qc.x(0)  # Apply X gate if checkbox is checked

    qc.h(0)
    qc.cx(0, 1)
    qc.measure_all()

    # Simulation
    simulator = Aer.get_backend('qasm_simulator')
    transpiled_circuit = transpile(qc, simulator)
    result = simulator.run(transpiled_circuit, shots=1024).result()

    # Get results and plot histogram
    counts = result.get_counts(qc)
    fig, ax = plt.subplots()
    plot_histogram(counts, ax=ax)
    
    return fig

# Gradio Interface
interface = gr.Interface(
    fn=run_quantum_circuit,
    inputs=gr.Checkbox(label="Apply X Gate to Qubit 0 (Set to |1⟩)"),
    outputs=gr.Plot(),
    title="Quantum Circuit Simulator",
    description="Simulate a simple quantum circuit. Optionally apply an X gate before the Hadamard and CNOT gates."
)

if __name__ == "__main__":
    interface.launch()