from __future__ import annotations

from collections.abc import Mapping

from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector


def _bitstring_key(label: str) -> tuple[int, int | str]:
    if label and set(label) <= {"0", "1"}:
        return (len(label), int(label, 2))
    return (len(label), label)


def counts_to_probabilities(counts: Mapping[str, int]) -> dict[str, float]:
    total = sum(counts.values())
    if total <= 0:
        raise ValueError("counts must contain a positive total")

    ordered = sorted(counts.items(), key=lambda item: _bitstring_key(item[0]))
    return {bitstring: value / total for bitstring, value in ordered}


def statevector_probabilities(circuit: QuantumCircuit) -> dict[str, float]:
    clean_circuit = circuit.remove_final_measurements(inplace=False)
    probabilities = Statevector.from_instruction(clean_circuit).probabilities_dict()
    ordered = sorted(probabilities.items(), key=lambda item: _bitstring_key(item[0]))
    return {bitstring: float(value) for bitstring, value in ordered}