import torch


def rgbd_to_world(p, depth, K, R_ex):
    """
    Given pixel location and depth, get world coordinates
    :param p: b x 2 tensor
    :param depth: b tensor
    :param k: 1 x 3 x 3 tensor
    :param r_ex: 1 x 3 x 3 tensor
    :return: p_world_right: b x 3 tensor in right hand coordinate
    """
    n = p.size(0)
    x_temp = (p[:, 0] + 1 - K[0, 0, 2]) / K[0, 0, 0]
    y_temp = (p[:, 1] + 1 - K[0, 1, 2]) / K[0, 1, 1]
    z_temp = 1
    x = x_temp / torch.sqrt(x_temp**2 + y_temp**2 + z_temp**2) * depth
    y = y_temp / torch.sqrt(x_temp**2 + y_temp**2 + z_temp**2) * depth
    z = z_temp / torch.sqrt(x_temp**2 + y_temp**2 + z_temp**2) * depth
    p_cam = torch.stack((x, y, z), 1).view(n, 3, 1) # n x 3
    p_world = torch.bmm(torch.transpose(R_ex, 1, 2).expand(n, -1, -1), p_cam)
    p_world_right = torch.stack((p_world[:, 0, 0], p_world[:, 2, 0], -p_world[:, 1, 0]), 1)
    return p_world_right