import torch


def convert_map2world(pose, map_shape, map_scale, angles):
    """
    Inputs:
        pose - (bs, 3) Tensor agent pose in map coordinates (x, y, theta) 
             - Note: theta is discrete angle idx
        map_shape - (f, h, w) tuple
        map_scale - scalar
        angles - (nangles, ) Tensor

    Conventions:
        world positions - X upward, Y rightward, origin at center
        map positions - X rightward, Y downward, origin at top-left
    """
    x = pose[:, 0].float()
    y = pose[:, 1].float()
    angle_idx = pose[:, 2].long()
    mh, mw = map_shape[1:]

    x_world = ((mh-1)/2 - y) * map_scale
    y_world = (x - (mw-1)/2) * map_scale
    theta_world = angles[angle_idx]

    return torch.stack([x_world, y_world, theta_world], dim=1)