import torch
import math


def cube_face_spherical_meshgrid(cube_dim, polar=False):
    """
    Returns a mesh grid of a face of a cube in spherical coordinates, assuming that the center pixel of the image is (0,0). This is appropriate for cube maps as each face uses the same equatorial aspect of the gnomonic projection.
    """
    if not polar:
        lat = torch.linspace(
            -math.pi / 4, math.pi / 4, steps=cube_dim).view(-1, 1).expand(
                -1, cube_dim)
        lon = torch.linspace(
            -math.pi / 4, math.pi / 4, steps=cube_dim).view(1, -1).expand(
                cube_dim, -1)
    else:
        lat = torch.linspace(
            math.pi / 4, 3 * math.pi / 4, steps=cube_dim).view(-1, 1).expand(
                -1, cube_dim)
        lon = torch.linspace(
            0, math.pi, steps=cube_dim).view(1, -1).expand(cube_dim, -1)
    return lon, lat