import torch


def template_edge_loss(meshes, template_mesh):
    """
    Computes mesh edge length regularization loss averaged across all meshes
    in a batch. Each mesh contributes equally to the final loss, regardless of
    the number of edges per mesh in the batch by weighting each mesh with the
    inverse number of edges. For example, if mesh 3 (out of N) has only E=4
    edges, then the loss for each edge in mesh 3 should be multiplied by 1/E to
    contribute to the final loss.
    Args:
        meshes: Meshes object with a batch of meshes.
        target_length: Resting value for the edge length.
    Returns:
        loss: Average loss across the batch. Returns 0 if meshes contains
        no meshes or all empty meshes.
    """
    if meshes.isempty():
        return torch.tensor(
            [0.0], dtype=torch.float32, device=meshes.device, requires_grad=True
        )

    N = len(meshes)
    edges_packed = meshes.edges_packed()  # (sum(E_n), 3)
    temp_edges_packed = template_mesh.edges_packed()  # (sum(E_n), 3)
    verts_packed = meshes.verts_packed()  # (sum(V_n), 3)
    temp_verts_packed = template_mesh.verts_packed()  # (sum(V_n), 3)

    verts_edges = verts_packed[edges_packed]
    temp_verts_edges = temp_verts_packed[temp_edges_packed]
    v0, v1 = verts_edges.unbind(1)
    t_v0, t_v1 = temp_verts_edges.unbind(1)
    edge_distance = (v0 - v1).norm(dim=1, p=2) ** 2.0
    t_edge_distance = (t_v0 - t_v1).norm(dim=1, p=2) ** 2.0
    loss = (edge_distance - t_edge_distance).norm(p=2)

    return loss / N