import torch


def depth_map_to_3d_torch(depth, cam_K, cam_W):
    """Derive 3D locations of each pixel of a depth map.

    Args:
        depth (torch.FloatTensor): tensor of size B x 1 x N x M
            with depth at every pixel
        cam_K (torch.FloatTensor): tensor of size B x 3 x 4 representing
            camera matrices
        cam_W (torch.FloatTensor): tensor of size B x 3 x 4 representing
            world matrices
    Returns:
        loc3d (torch.FloatTensor): tensor of size B x 3 x N x M
            representing color at given 3d locations
        mask (torch.FloatTensor):  tensor of size B x 1 x N x M with
            a binary mask if the given pixel is present or not
    """
    depth = torch.from_numpy(depth)
    cam_K = torch.from_numpy(cam_K)
    cam_W = torch.from_numpy(cam_W)

    N, M = depth.size()
    device = depth.device
    # Turn depth around. This also avoids problems with inplace operations
    depth = -depth.permute(1,0)

    zero_one_row = torch.tensor([[0., 0., 0., 1.]])
    zero_one_row = zero_one_row.expand(1, 4).to(device)
    # add row to world mat
    cam_W = torch.cat((cam_W, zero_one_row), dim=0)

    # clean depth image for mask
    # upperlimit =  1.e+10
    upperlimit = float("Inf")
    mask = (depth.abs() != upperlimit).float()
    depth[depth == upperlimit] = 0
    depth[depth == -1*upperlimit] = 0

    # 4d array to 2d array k=N*M
    d = depth.reshape(1,N * M)

    # create pixel location tensor
    px, py = torch.meshgrid([torch.arange(0, N), torch.arange(0, M)])
    px, py = px.to(device), py.to(device)

    p = torch.cat((
        px.expand(px.size(0), px.size(1)), 
        (M - py).expand(py.size(0), py.size(1))
    ), dim=0)
    p = p.reshape(2, py.size(0) * py.size(1))
    p = (p.float() / M * 2)
    
    # create terms of mapping equation x = P^-1 * d*(qp - b)
    P = cam_K[:2, :2].float().to(device)    
    q = cam_K[2:3, 2:3].float().to(device)   
    b = cam_K[:2, 2:3].expand(2, d.size(1)).to(device)
    Inv_P = torch.inverse(P).to(device)   

    rightside = (p.float() * q.float() - b.float()) * d.float()
    x_xy = torch.matmul(Inv_P, rightside)
    
    # add depth and ones to location in world coord system
    x_world = torch.cat((x_xy, d, torch.ones_like(d)), dim=0)

    # derive loactoion in object coord via loc3d = W^-1 * x_world
    Inv_W = torch.inverse(cam_W)
    Inv_W_exp = Inv_W.expand(4, 4)
    loc3d = torch.matmul(Inv_W_exp, x_world.double())
    loc3d = loc3d.reshape(4, N, M)

    loc3d = loc3d[:3,:,:].to(device)
    mask = mask.to(device)
    loc3d = loc3d.view(3, N * M)
    return loc3d, mask