import pytest
from qiskit import QuantumCircuit

from quantum_learning.utils import counts_to_probabilities, statevector_probabilities


def test_counts_to_probabilities_normalizes():
    probabilities = counts_to_probabilities({"11": 2, "00": 6})
    assert probabilities["00"] == pytest.approx(0.75)
    assert probabilities["11"] == pytest.approx(0.25)


def test_statevector_probabilities_ignores_final_measurements():
    circuit = QuantumCircuit(1, 1)
    circuit.h(0)
    circuit.measure(0, 0)

    probabilities = statevector_probabilities(circuit)
    assert probabilities["0"] == pytest.approx(0.5)
    assert probabilities["1"] == pytest.approx(0.5)