from qiskit import QuantumCircuit

def prepare_bell_state(state_number,qc):
    if state_number == 1:  # |Φ+⟩
        qc.h(0)
        qc.cx(0, 1)
    elif state_number == 2:  # |Φ-⟩
        qc.h(0)
        qc.cx(0, 1)
        qc.z(0)
    elif state_number == 3:  # |Ψ+⟩
        qc.h(0)
        qc.cx(0, 1)
        qc.x(1)
    elif state_number == 4:  # |Ψ-⟩
        qc.h(0)
        qc.cx(0, 1)
        qc.x(1)
        qc.z(0)
    else:
        raise ValueError("State number must be 1, 2, 3, or 4.")

    return qc

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

def cxHI_simul(bell_state):
    # Create a 2-qubit quantum circuit
    qc = QuantumCircuit(2)

    prepare_bell_state(bell_state,qc)
    
    

    # Apply CNOT gate with the first qubit as control and the second as target
    qc.cx(0, 1)
    
    # Apply H gate on the first qubit
    qc.h(0)
    
    # Draw the circuit
    qc.draw(output='mpl')

    # (Optional) Simulate the circuit
    simulator = Aer.get_backend('statevector_simulator')
    result = execute(qc, simulator).result()
    statevector = result.get_statevector()
    statevector = np.asarray(result.get_statevector())
    rounded_statevector = [round(x.real, 3) + round(x.imag, 3) * 1j for x in statevector]
    print("Rounded Statevector:", rounded_statevector)

    qc.measure_all()
    # szimulálás
    job = execute(qc, simulator, shots=1024)
    result = job.result()
    counts = result.get_counts(qc)
    print(counts)

# bell állapotok előkészítése
for i in range(1,5):
    cxHI_simul(i)

