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

def apply_cnot_gradio(initial_state_str, visualize_statevector):
    initial_states = {
        "|00⟩": [1, 0, 0, 0],
        "|01⟩": [0, 1, 0, 0],
        "|10⟩": [0, 0, 1, 0],
        "|11⟩": [0, 0, 0, 1]
    }
    
    initial_state = initial_states[initial_state_str]
    
    qc = QuantumCircuit(2, 2)
    qc.initialize(initial_state, [0, 1])
    qc.cx(0, 1)
    
    if visualize_statevector:
        simulator = Aer.get_backend('statevector_simulator')
        transpiled_circuit = transpile(qc, simulator)
        result = simulator.run(transpiled_circuit).result()
        statevector = result.get_statevector()
        fig = plot_bloch_multivector(statevector, title=f"Statevector for {initial_state_str}")
    else:
        qc.measure([0, 1], [0, 1])
        simulator = Aer.get_backend('qasm_simulator')
        transpiled_circuit = transpile(qc, simulator)
        result = simulator.run(transpiled_circuit, shots=1024).result()
        counts = result.get_counts()
        fig, ax = plt.subplots()
        plot_histogram(counts, ax=ax, title=f"Measurement Counts for {initial_state_str}")
    
    return fig

interface = gr.Interface(
    fn=apply_cnot_gradio,
    inputs=[
        gr.Radio(["|00⟩", "|01⟩", "|10⟩", "|11⟩"], label="Initial State"),
        gr.Checkbox(label="Visualize Statevector (Bloch Spheres)")
    ],
    outputs=gr.Plot(),
    title="Quantum CNOT Operation and Visualization",
    description="Apply a CNOT gate on selected initial state and visualize results either as measurement counts or as statevector on Bloch spheres."
)

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