import torch


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

    Args:
        verts_per_mesh: Tensor of shape (N,) giving the number of vertices
            in each mesh in the batch where N is the batch size.
        edges_per_mesh: Tensor of shape (N,) giving the number of edges
            in each mesh in the batch

    Returns:
        verts_idx: A tensor with vert indices for each mesh ordered sequentially
            by mesh index.
    """
    # e.g. verts_per_mesh = (4, 5, 6)
    # e.g. edges_per_mesh = (5, 7, 9)

    V = verts_per_mesh.sum()  # e.g. 15
    E = edges_per_mesh.sum()  # e.g. 21

    verts_per_mesh_cumsum = verts_per_mesh.cumsum(dim=0)  # (N,) e.g. (4, 9, 15)
    edges_per_mesh_cumsum = edges_per_mesh.cumsum(dim=0)  # (N,) e.g. (5, 12, 21)

    v_to_e_idx = verts_per_mesh_cumsum.clone()

    # vertex to edge index.
    v_to_e_idx[1:] += edges_per_mesh_cumsum[
        :-1
    ]  # e.g. (4, 9, 15) + (0, 5, 12) = (4, 14, 27)

    # vertex to edge offset.
    v_to_e_offset = V - verts_per_mesh_cumsum  # e.g. 15 - (4, 9, 15) = (11, 6, 0)
    v_to_e_offset[1:] += edges_per_mesh_cumsum[
        :-1
    ]  # e.g. (11, 6, 0) + (0, 5, 12) = (11, 11, 12)
    e_to_v_idx = (
        verts_per_mesh_cumsum[:-1] + edges_per_mesh_cumsum[:-1]
    )  # (4, 9) + (5, 12) = (9, 21)
    e_to_v_offset = (
        verts_per_mesh_cumsum[:-1] - edges_per_mesh_cumsum[:-1] - V
    )  # (4, 9) - (5, 12) - 15 = (-16, -18)

    # Add one new vertex per edge.
    idx_diffs = torch.ones(V + E, device=device, dtype=torch.int64)  # (36,)
    idx_diffs[v_to_e_idx] += v_to_e_offset
    idx_diffs[e_to_v_idx] += e_to_v_offset

    # e.g.
    # [
    #  1, 1, 1, 1, 12, 1, 1, 1, 1,
    #  -15, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1,
    #  -17, 1, 1, 1, 1, 1, 13, 1, 1, 1, 1, 1, 1, 1
    # ]

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

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

    return verts_idx