
import torch
import torch.nn as nn
from knn_cuda import KNN
from pointnet2_utils import grouping_operation
from pointnet2 import pointnet2_utils as pn2_utils
from emd import earth_mover_distance
import numpy as np

def array2samples_distance(array1, array2):
    """
    arguments:
        array1: the array, size: (num_point, num_feature)
        array2: the samples, size: (num_point, num_feature)
    returns:
        distances: each entry is the distance from a sample to array1
    """
    num_point1, num_features1 = array1.shape
    num_point2, num_features2 = array2.shape
    expanded_array1 = array1.repeat(num_point2, 1)
    expanded_array2 = torch.reshape(
        torch.unsqueeze(array2, 1).repeat(1, num_point1, 1),
        (-1, num_features2))
    distances = (expanded_array1 - expanded_array2) * (expanded_array1 - expanded_array2)
    #    distances = torch.sqrt(distances)
    distances = torch.sum(distances, dim=1)
    distances = torch.reshape(distances, (num_point2, num_point1))
    distances = torch.min(distances, dim=1)[0]
    distances = torch.mean(distances)
    return distances


def chamfer_distance_numpy(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist = 0
    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])
        dist = dist + (0.5 * av_dist1 + 0.5 * av_dist2) / batch_size
    return dist * 100


def chamfer_distance_numpy_test(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist_all = 0
    dist1 = 0
    dist2 = 0
    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])
        dist_all = dist_all + (av_dist1 + av_dist2) / batch_size
        dist1 = dist1 + av_dist1 / batch_size
        dist2 = dist2 + av_dist2 / batch_size
    return dist_all, dist1, dist2


class PointLoss(nn.Module):
    def __init__(self):
        super(PointLoss, self).__init__()

    def forward(self, array1, array2):
        return chamfer_distance_numpy(array1, array2)

class EMDLoss(nn.Module):
    # must use cuda
    def __init__(self):
        super(EMDLoss, self).__init__()


    def forward(self, array1, array2):
        return torch.mean(earth_mover_distance(array1, array2,False))*10/256

# class PointEMDLoss(nn.Module):
#     def __init__(self):
#         super(PointLoss, self).__init__()
#         self.EMD=emd.emdModule()
#
#     def forward(self, array1, array2,eps,iters):
#         return self.EMD(array1, array2,eps,iters)

#  add  repu loss
def knn_point(group_size, point_cloud, query_cloud, transpose_mode=False):
    knn_obj = KNN(k=group_size, transpose_mode=transpose_mode)
    dist, idx = knn_obj(point_cloud, query_cloud)
    return dist, idx

class RepulsionLoss(nn.Module):
    def __init__(self, alpha=1.0, nn_size=5, radius=0.07, h=0.03, eps=1e-12):
        super(RepulsionLoss, self).__init__()
        self.alpha = alpha
        self.nn_size = nn_size
        self.radius = radius
        self.h = h
        self.eps = eps

    def get_repulsion_loss(self, pred):
        _, idx = knn_point(self.nn_size, pred, pred, transpose_mode=True)
        idx = idx[:, :, 1:].to(torch.int32) # remove first one
        idx = idx.contiguous() # B, N, nn

        pred = pred.transpose(1, 2).contiguous() # B, 3, N
        grouped_points = grouping_operation(pred, idx) # (B, 3, N), (B, N, nn) => (B, 3, N, nn)

        grouped_points = grouped_points - pred.unsqueeze(-1)
        dist2 = torch.sum(grouped_points ** 2, dim=1)
        dist2 = torch.max(dist2, torch.tensor(self.eps).cuda())
        dist = torch.sqrt(dist2)
        weight = torch.exp(- dist2 / self.h ** 2)

        uniform_loss = torch.mean((self.radius - dist) * weight)
        # uniform_loss = torch.mean(self.radius - dist * weight) # punet
        return uniform_loss

    def forward(self,array1 ):
        return self.alpha * self.get_repulsion_loss(array1)

class PointLoss_test(nn.Module):
    def __init__(self):
        super(PointLoss_test, self).__init__()

    def forward(self, array1, array2):
        return chamfer_distance_numpy_test(array1, array2)


def distance_squre(p1, p2):
    tensor = torch.FloatTensor(p1) - torch.FloatTensor(p2)
    val = tensor.mul(tensor)
    val = val.sum()
    return val


def index_points(points, idx):
    """
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points


def farthest_point_sample(xyz, npoint, RAN=True):
    """
    Input:
        xyz: pointcloud data, [B, N, C]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    if RAN:
        farthest = torch.randint(0, 1, (B,), dtype=torch.long).to(device)
    else:
        farthest = torch.randint(1, 2, (B,), dtype=torch.long).to(device)

    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids


if __name__ == '__main__':
    x1 = torch.rand(20, 1024, 3).cuda()
    x2 = torch.rand(20, 1024, 3).cuda()
    x1 = torch.tensor([[[0.9,0.6,0.7],[0.0,0,0]]]).cuda()
    x2 = torch.tensor([[[1.0,1.0,1.0],[0.5,0.5,0.5]]]).cuda()
    p1 = torch.from_numpy(np.array([[[1.7, -0.1, 0.1], [0.1, 1.2, 0.3]]], dtype=np.float32)).cuda()
    print(p1.shape)
    p1 = p1.repeat(3, 1, 1)
    p2 = torch.from_numpy(np.array([[[0.3, 1.8, 0.2], [1.2, -0.2, 0.3]]], dtype=np.float32)).cuda()
    p2 = p2.repeat(3, 1, 1)
    print(p1.shape)
    p = EMDLoss()
    print(p(p1, p2))

    print(PointLoss())