import torch


def get_arcs(start_point, primitives, ids):
    """
    Args:
        start_point: tensor [batch_size, 2]
        primitives: tensor [num_primitives, 5]
            primitives[i] = [dx, dy, theta, sharpness, width]
        ids: long tensor [batch_size, num_arcs]

    Returns:
        arcs: tensor [batch_size, num_arcs, 7]
            arcs[i] = [x_start, y_start, dx, dy, theta, sharpness, width]
    """
    arcs = primitives[ids]  # [batch_size, num_arcs, 5]
    dxdy = arcs[..., :2]  # [batch_size, num_arcs, 5]
    start_and_dxdy = torch.cat(
        [start_point.unsqueeze(-2), dxdy], dim=-2
    )  # [batch_size, num_arcs + 1, 2]
    # [batch_size, num_arcs, 2]
    start_points = torch.cumsum(start_and_dxdy, dim=-2)[:, :-1]
    arcs = torch.cat([start_points, arcs], dim=-1)
    return arcs