# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/platform/qcircuit_util.ipynb.

# %% auto 0
__all__ = ['get_element_matching_indices', 'get_entanglement_bins']

# %% ../../src/platform/qcircuit_util.ipynb 2
from ..imports import *

# %% ../../src/platform/qcircuit_util.ipynb 4
def get_element_matching_indices(a, b):
    """Compares (2d) `a` with `b`. Returns the indices of `b`, where a element of `a` matches with `b`."""
    # Expand dimensions of a to match the shape of b for element-wise comparison
    expanded_a = a.unsqueeze(0).expand(b.shape[0], *a.shape)  # [b0, a0, a1]
    expanded_b = b.unsqueeze(1)                               # [b0,  1, b1]
    
    # Compare all vector entries of a with all vectors of b
    matches = torch.all(expanded_a == expanded_b, dim=-1)

    matching_indices = torch.nonzero(torch.any(matches, dim=1)).squeeze()
    
    if   matching_indices.dim() == 0: matching_indices = torch.tensor([matching_indices])

    return matching_indices

# %% ../../src/platform/qcircuit_util.ipynb 5
def get_entanglement_bins(num_of_qubits):
    """Returns all SRV sorted in entangle bins which correspond to a number of entangled qubits."""
    dist_srvs = [x for x in itertools.product(*([[1,2]]*num_of_qubits))]
    dist_srvs = np.array(dist_srvs, dtype=int)[np.sum(dist_srvs, axis=1)!=num_of_qubits+1].tolist()
    dist_srvs = sorted(dist_srvs, key=lambda x: sum(x))
    dist_srvs = np.array(dist_srvs)
    
    entangle = [1] + [scipy.special.comb(num_of_qubits, i, exact=True) for i in range(2, num_of_qubits)]
    
    entanglement_bins = np.split(dist_srvs, np.cumsum(entangle))
    
    ent_bits = [f"{sum(n[0])-num_of_qubits} qubit entangled" for n in entanglement_bins]
    
    return [x.tolist() for x in entanglement_bins], ent_bits
