import torch


def create_faces_index(faces_per_mesh, device=None):
    """
    Helper function to group the faces indices for each mesh. New faces are
    stacked at the end of the original faces tensor, so in order to have
    sequential packing, the faces tensor needs to be reordered to that faces
    corresponding to each mesh are grouped together.

    Args:
        faces_per_mesh: Tensor of shape (N,) giving the number of faces
            in each mesh in the batch where N is the batch size.

    Returns:
        faces_idx: A tensor with face indices for each mesh ordered sequentially
            by mesh index.
    """
    # e.g. faces_per_mesh = [2, 5, 3]

    F = faces_per_mesh.sum()  # e.g. 10
    faces_per_mesh_cumsum = faces_per_mesh.cumsum(dim=0)  # (N,) e.g. (2, 7, 10)

    switch1_idx = faces_per_mesh_cumsum.clone()
    switch1_idx[1:] += (
        3 * faces_per_mesh_cumsum[:-1]
    )  # e.g. (2, 7, 10) + (0, 6, 21) = (2, 13, 31)

    switch2_idx = 2 * faces_per_mesh_cumsum  # e.g. (4, 14, 20)
    switch2_idx[1:] += (
        2 * faces_per_mesh_cumsum[:-1]
    )  # e.g. (4, 14, 20) + (0, 4, 14) = (4, 18, 34)

    switch3_idx = 3 * faces_per_mesh_cumsum  # e.g. (6, 21, 30)
    switch3_idx[1:] += faces_per_mesh_cumsum[
        :-1
    ]  # e.g. (6, 21, 30) + (0, 2, 7) = (6, 23, 37)

    switch4_idx = 4 * faces_per_mesh_cumsum[:-1]  # e.g. (8, 28)

    switch123_offset = F - faces_per_mesh  # e.g. (8, 5, 7)

    idx_diffs = torch.ones(4 * F, device=device, dtype=torch.int64)
    idx_diffs[switch1_idx] += switch123_offset
    idx_diffs[switch2_idx] += switch123_offset
    idx_diffs[switch3_idx] += switch123_offset
    idx_diffs[switch4_idx] -= 3 * F

    # e.g
    # [
    #  1, 1, 9, 1, 9, 1, 9, 1,                                       -> mesh 0
    #  -29, 1, 1, 1, 1, 6, 1, 1, 1, 1, 6, 1, 1, 1, 1, 6, 1, 1, 1, 1, -> mesh 1
    #  -29, 1, 1, 8, 1, 1, 8, 1, 1, 8, 1, 1                          -> mesh 2
    # ]

    faces_idx = idx_diffs.cumsum(dim=0) - 1

    # e.g.
    # [
    #  0, 1, 10, 11, 20, 21, 30, 31,
    #  2, 3, 4, 5, 6, 12, 13, 14, 15, 16, 22, 23, 24, 25, 26, 32, 33, 34, 35, 36,
    #  7, 8, 9, 17, 18, 19, 27, 28, 29, 37, 38, 39
    # ]
    # where for mesh 0, [0, 1] are the indices of the existing faces, and
    # [10, 11, 20, 21, 30, 31] are the indices of the new faces after subdivision.

    return faces_idx