import torch


def calculate_normal(p0, p1, p2, pad, return_mask=True):

    normals = torch.cross(p1 - p0, p2 - p0, 1)
    normals = normals / normals.norm(dim=1, keepdim=True)
    normals[torch.isnan(normals)] = 0.0
    normals = torch.nn.functional.pad(normals, pad, mode='constant', value=0.0)

    if return_mask:
        mask = calculate_normal_mask(p0, p1, p2, pad)
        return normals, mask
    else:
        return normals


def calculate_normal_mask(p0, p1, p2, pad):
    mask = (p0.sum(1, keepdim=True) != 0) & \
           (p1.sum(1, keepdim=True) != 0) & \
           (p2.sum(1, keepdim=True) != 0)
    return torch.nn.functional.pad(mask, pad, mode='constant', value=0)


def calculate_normals(depth, camera, to_world=False):

    points = camera.reconstruct_depth_map(
        depth, euclidean=False, to_world=to_world)

    p0 = points[:, :,  :-1 ,  :-1]
    p1 = points[:, :, 1:   ,  :-1]
    p2 = points[:, :,  :-1 , 1:  ]
    normals1, mask1 = calculate_normal(p0, p1, p2, [0, 1, 0, 1])

    p0 = points[:, :, 1:   , 1:  ]
    p1 = points[:, :,  :-1 , 1:  ]
    p2 = points[:, :, 1:   ,  :-1]
    normals2, mask2 = calculate_normal(p0, p1, p2, [1, 0, 1, 0])

    p0 = points[:, :, 1:   ,  :-1]
    p1 = points[:, :,  :-1 ,  :-1]
    p2 = points[:, :, 1:   , 1:  ]
    normals3, mask3 = calculate_normal(p0, p1, p2, [0, 1, 1, 0])

    p0 = points[:, :,  :-1 , 1:  ]
    p1 = points[:, :, 1:   , 1:  ]
    p2 = points[:, :,  :-1 ,  :-1]
    normals4, mask4 = calculate_normal(p0, p1, p2, [1, 0, 0, 1])

    normals = torch.stack([normals1, normals2, normals3, normals4], -1)
    mask = torch.stack([mask1, mask2, mask3, mask4], -1)
    return normals, mask