import torch


def _legacy_to_opengl(face_vertices_image, face_vertices_z, valid_faces=None):
    """Transform Kaolin legacy rasterization coordinate system to OpenGL format
    (for nvdiffrast)

    For more details see: https://nvlabs.github.io/nvdiffrast/#coordinate-systems

    Args:
        face_vertices_image (torch.Tensor):
            of shape (batch_size, num_faces, 3, 3) or (num_faces, 3, 3)
        face_vertices_z (torch.Tensor):
            of shape (batch_size, num_faces, 3) or (num_faces, 3)
        valid_faces (torch.BoolTensor):
            of shape (batch_size, num_faces) or (num_faces)

    Returns:
        (torch.Tensor, torch.IntTensor):

        - position of vertices, of shape (batch_size, num_faces * 3, 4) or (num_faces * 3, 4)
        - face indices, of shape (num_faces, 3)
    """
    z = -face_vertices_z / (abs(face_vertices_z).max() + 1e-6)
    _face_vertices_image = face_vertices_image.reshape(*face_vertices_image.shape[:-3], -1, 2)
    pos = torch.stack([
        _face_vertices_image[..., 0],
        -_face_vertices_image[..., 1],
        z.reshape(*z.shape[:-2], -1)
    ], dim=-1)
    if valid_faces is None:
        pos = torch.nn.functional.pad(pos, (0, 1), value=1.)
    else:
        pad = (valid_faces.unsqueeze(-1) * 2. - 1.).expand(*valid_faces.shape, 3)
        pad = pad.reshape(*valid_faces.shape[:-1], -1, 1)
        pos = torch.cat([pos, pad], dim=-1)

    tri = torch.arange(pos.shape[-2], device=pos.device, dtype=torch.int).reshape(-1, 3)
   
    return pos, tri