import torch

def expanded_pairwise_distances(x, y):
    '''
    Input: x is a bxNxd matrix
           y is an optional bxMxd matirx
    Output: dist is a bxNxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
            if y is not given then use 'y=x'.
    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
    '''
    differences = x.unsqueeze(2) - y.unsqueeze(1)
    distances = torch.sum(differences * differences, -1)
    return distances

def chamfer_distance(x, y):
    '''
    input x and y are bxNxM matrix, b: batch, N:number of point, M: point dim (ex. 2 for 2D or 3 for 3D)
    output is a bx1 Matrix with the value of the chamfer distance for each sample of the batch
    '''
    dist_vec = expanded_pairwise_distances(x, y)
    min_distances = torch.topk(dist_vec, k=1, dim=2, largest=False).values
    chamfer = torch.sum(min_distances, dim=1) / torch.tensor(x.shape[1])
    
    return chamfer

class ChamferLoss(torch.nn.Module):
   def forward(self, x, y):
        chamfer = chamfer_distance(x, y)

        return torch.sum(chamfer)

if __name__ == "__main__":
    x = torch.tensor([
        [
            [0., 0., 0.],
            [0., 1., 0.],
            [0., 1., 0.],
        ],
        [
            [1., 1., 0.],
            [1., 2., 0.],
            [0., 1., 0.],
        ]
    ])
    y = torch.tensor([
        [
            [0., 1., 0.],
            [0., 1., 0.],
            [0., 1., 0.],
        ],
        [
            [1., 1., 0.],
            [1., 2., 0.],
            [0., 1., 0.],
        ]
    ])
    chamfer = ChamferLoss()

    print('chamfer loss torch (cpu):', chamfer(x, y))
    print('chamfer loss torch (cuda):', chamfer(x.cuda(), y.cuda()))
    
    # import sys   
    # sys.path.append("../distance/chamfer/")
    # import dist_chamfer as cd
    # CD = cd.chamferDist()
    # dist1, dist2, _, _= CD(x, y)
    # print('orig', dist1)

[((12, 16, 12, 56), 'torch.sum', 'torch.sum', ({(12, 26, 12, 51): 'differences * differences', (12, 53, 12, 55): '-1'}, {}), '(differences * differences, -1)', False, 'import torch\n'), ((33, 8, 44, 6), 'torch.tensor', 'torch.tensor', ({(33, 21, 44, 5): '[[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], [[1.0, 1.0, 0.0], [\n    1.0, 2.0, 0.0], [0.0, 1.0, 0.0]]]'}, {}), '([[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], [[1.0, \n    1.0, 0.0], [1.0, 2.0, 0.0], [0.0, 1.0, 0.0]]])', False, 'import torch\n'), ((45, 8, 56, 6), 'torch.tensor', 'torch.tensor', ({(45, 21, 56, 5): '[[[0.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], [[1.0, 1.0, 0.0], [\n    1.0, 2.0, 0.0], [0.0, 1.0, 0.0]]]'}, {}), '([[[0.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], [[1.0, \n    1.0, 0.0], [1.0, 2.0, 0.0], [0.0, 1.0, 0.0]]])', False, 'import torch\n'), ((21, 20, 21, 67), 'torch.topk', 'torch.topk', (), '', False, 'import torch\n'), ((22, 14, 22, 45), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((22, 48, 22, 72), 'torch.tensor', 'torch.tensor', ({(22, 61, 22, 71): 'x.shape[1]'}, {}), '(x.shape[1])', False, 'import torch\n'), ((30, 15, 30, 33), 'torch.sum', 'torch.sum', ({(30, 25, 30, 32): 'chamfer'}, {}), '(chamfer)', False, 'import torch\n')]