
import numpy as np

class ClassicalRegister:
    """
    Stores measurement results as classical bits.

    Usage:
        cr = ClassicalRegister(2, name="c")
        cr[0] = 1
        print(cr.to_string())  # "01"
    """
    def __init__(self, size, name="c"):
        self.size = size
        self.name = name
        self.bits  = [0] * size

    def __setitem__(self, idx, val):
        self.bits[idx] = int(val)

    def __getitem__(self, idx):
        return self.bits[idx]

    def to_string(self):
        return "".join(str(b) for b in self.bits)

    def reset(self):
        self.bits = [0] * self.size

    def __repr__(self):
        return f"ClassicalRegister({self.name}={self.to_string()})"


def measure_qubit(state, qubit, n):
    """
    Measure a single qubit. Collapses the state.
    Returns (result_bit, new_state).
    """
    psi  = state.reshape([2]*n)
    # Probability of measuring |0> on this qubit
    idx0 = [slice(None)]*n
    idx0[qubit] = 0
    prob0 = float(np.sum(np.abs(psi[tuple(idx0)])**2))
    prob0 = max(0.0, min(1.0, prob0))

    result = 0 if np.random.random() < prob0 else 1

    # Collapse
    new_psi = psi.copy()
    kill_idx = [slice(None)]*n
    kill_idx[qubit] = 1 - result
    new_psi[tuple(kill_idx)] = 0

    norm = np.linalg.norm(new_psi)
    if norm > 1e-12:
        new_psi /= norm

    return result, new_psi.reshape(-1)


def measure_all(state, n, shots=1):
    """
    Sample from the statevector distribution.
    Returns a counts dict: {"00": 512, "11": 512}
    """
    probs = np.abs(state)**2
    probs /= probs.sum()
    outcomes = np.random.choice(2**n, size=shots, p=probs)
    counts = {}
    for o in outcomes:
        key = format(o, f"0{n}b")
        counts[key] = counts.get(key, 0) + 1
    return counts


def measure_partial(state, qubits, n):
    """
    Measure only a subset of qubits.
    Returns (result_bits_dict, collapsed_state).

    Example:
        bits, new_state = measure_partial(state, qubits=[0,1], n=3)
        print(bits)  # {0: 1, 1: 0}
    """
    bits = {}
    current_state = state.copy()
    for q in qubits:
        bit, current_state = measure_qubit(current_state, q, n)
        bits[q] = bit
    return bits, current_state
