from qiskit import QuantumRegister, QuantumCircuit
from qiskit_aer import AerSimulator
from qiskit.visualization import plot_histogram

def simulate(qc: QuantumCircuit, shots):
    if shots is None:
        shots = 1

    simulator = AerSimulator()
    result = simulator.run(qc, shots).result()
    return result.get_counts(qc)


def entangle(register: QuantumRegister, circuit: QuantumCircuit):
    '''
    Accepts a register and a circuit, then entangles the first two qubits before returning the circuit.
    :param register: A register with at least two qubits
    :param circuit: The circuit.
    :return: The provided circuit and its entangled circuit.
    '''

    # Apply the Hadamard and CNOT gates to entangle the qubits
    circuit.h(register[0])
    circuit.cx(register[0], register[1])

    return circuit

def visualize(counts):
    plot_histogram(counts)
