import torch


def get_z_counts(att, prem_mask, hypo_mask):
    """
    Compute z counts (number of 0, continious, 1 elements).
    :param att: similarity matrix [B, prem, hypo]
    :param prem_mask:
    :param hypo_mask:
    :return: z0, zc, z1
    """
    # mask out all invalid positions with -1
    att = torch.where(hypo_mask.unsqueeze(1), att, att.new_full([1], -1.))
    att = torch.where(prem_mask.unsqueeze(2), att, att.new_full([1], -1.))

    z0 = (att == 0.).sum().item()
    zc = ((0 < att) & (att < 1)).sum().item()
    z1 = (att == 1.).sum().item()

    assert (att > -1).sum().item() == z0 + zc + z1, "mismatch"

    return z0, zc, z1