import random
import torch
import torch.nn as nn

import numpy as np
from open3d import *

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


def show_point_cloud(points: np.ndarray, rgd=None):
    assert points.ndim == 2

    point_cloud = PointCloud()
    point_cloud.points = Vector3dVector(points)
    if rgd is not None:
        point_cloud.paint_uniform_color(rgd)
    draw_geometries([point_cloud])


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


def resample_pcd(pcd, n):
    """Drop or duplicate points so that pcd has exactly n points"""
    idx = np.random.permutation(pcd.shape[0])
    if idx.shape[0] < n:
        idx = np.concatenate([idx, np.random.randint(pcd.shape[0], size=n-pcd.shape[0])])
    return pcd[idx[:n]]


def get_batch_data(data, device, args):
    real_point, target = data
    # print(real_point.size())
    batch_size = real_point.size()[0]
    real_center = torch.FloatTensor(batch_size, 1, args.crop_point_num, 3)
    input_cropped1 = torch.FloatTensor(batch_size, args.pnum, 3)
    input_cropped1 = input_cropped1.data.copy_(real_point)
    real_point = torch.unsqueeze(real_point, 1)
    input_cropped1 = torch.unsqueeze(input_cropped1, 1)
    p_origin = [0, 0, 0]
    input_partial = torch.FloatTensor(batch_size,1,args.pnum-args.crop_point_num,3)

    if args.cropmethod == 'random_center':
        choice = [torch.Tensor([1, 0, 0]), torch.Tensor([0, 0, 1]), torch.Tensor([1, 0, 1]),
                    torch.Tensor([-1, 0, 0]), torch.Tensor([-1, 1, 0])]
        for m in range(batch_size):
            index = random.sample(choice, 1)
            distance_list = []
            p_center = index[0]
            for n in range(args.pnum):
                distance_list.append(distance_squre(real_point[m, 0, n], p_center))
            distance_order = sorted(enumerate(distance_list), key=lambda x: x[1])

            for sp in range(args.crop_point_num):
                input_cropped1.data[m, 0, distance_order[sp][0]] = torch.FloatTensor([0, 0, 0])
                real_center.data[m, 0, sp] = real_point[m, 0, distance_order[sp][0]]
            
            crop_num_list = []
            for num_ in range(args.pnum - args.crop_point_num):
                crop_num_list.append(distance_order[num_ + args.crop_point_num][0])
            indices = torch.LongTensor(crop_num_list)
            input_partial[m,0]= torch.index_select(real_point[m,0],0,indices)
    
    input_partial = input_partial.squeeze().to(device)
    real_point = real_point.to(device)
    # print('the real_point is ', real_point.size())
    real_center = real_center.to(device)
    # print('the real_center is ', real_center.size())
    input_cropped1 = input_cropped1.to(device)
    # print('the input_cropped1 is ', input_cropped1.size())
    ############################
    # (1) data prepare
    ###########################
    
    ## 128  个点
    input_cropped1 = torch.squeeze(input_cropped1, 1) #2048

    partial_input = input_partial
    coarse_gt = real_point.squeeze()
    dense_gt = real_point.squeeze()

    partial_input = partial_input.to(device)
    coarse_gt = coarse_gt.to(device)
    coarse_gt = coarse_gt.permute(0, 2, 1) # (B, 3, 2048)
    # print('the coarse gt is ', coarse_gt.size())
    dense_gt = dense_gt.to(device)
    dense_gt = dense_gt.permute(0, 2, 1)
    partial_input = partial_input.permute(0, 2, 1) # (B, 3, 2048)
    # print('the partial input is ', partial_input.size())

    return  partial_input, coarse_gt, dense_gt, real_center, target 


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
    
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)

class PointLoss(nn.Module):
    def __init__(self):
        super(PointLoss,self).__init__()
        
    def forward(self,array1,array2):
        return chamfer_distance_numpy(array1,array2)

class PointLoss_eval(nn.Module):
    def __init__(self):
        super(PointLoss_eval,self).__init__()
        
    def forward(self,array1,array2):
        return chamfer_distance_numpy_eval(array1,array2)

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

def chamfer_distance_numpy_eval(array1, array2):
    batch_size, num_point, num_features = array1.shape
    # print(array1.shape)
    dist_all = []
    dist1 = []
    dist2 = []

    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])

        av_dist1 = av_dist1.detach().cpu().numpy()
        av_dist2 = av_dist2.detach().cpu().numpy()

        dist_all.append((av_dist1+av_dist2))
        dist1.append(av_dist1)
        dist2.append(av_dist2)

    return dist_all, dist1, dist2
################### L1 distance #######################


################### L2 distance #######################


if __name__ == '__main__':
    pc = read_point_cloud('examples/chair.pcd')
    draw_geometries([pc])
    
    # resample pcd
    pc.points = Vector3dVector(resample_pcd(np.asarray(pc.points), 16384))
    print(len(pc.points))
    draw_geometries([pc])

    # for i in range(8):
    #     pc = read_point_cloud('examples/{}.pcd'.format(i))
    #     draw_geometries([pc])

import random
import torch
import torch.nn as nn

import numpy as np
from open3d import *

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


def show_point_cloud(points: np.ndarray, rgd=None):
    assert points.ndim == 2

    point_cloud = PointCloud()
    point_cloud.points = Vector3dVector(points)
    if rgd is not None:
        point_cloud.paint_uniform_color(rgd)
    draw_geometries([point_cloud])


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


def resample_pcd(pcd, n):
    """Drop or duplicate points so that pcd has exactly n points"""
    idx = np.random.permutation(pcd.shape[0])
    if idx.shape[0] < n:
        idx = np.concatenate([idx, np.random.randint(pcd.shape[0], size=n-pcd.shape[0])])
    return pcd[idx[:n]]


def get_batch_data(data, device, args):
    real_point, target = data
    # print(real_point.size())
    batch_size = real_point.size()[0]
    real_center = torch.FloatTensor(batch_size, 1, args.crop_point_num, 3)
    input_cropped1 = torch.FloatTensor(batch_size, args.pnum, 3)
    input_cropped1 = input_cropped1.data.copy_(real_point)
    real_point = torch.unsqueeze(real_point, 1)
    input_cropped1 = torch.unsqueeze(input_cropped1, 1)
    p_origin = [0, 0, 0]
    input_partial = torch.FloatTensor(batch_size,1,args.pnum-args.crop_point_num,3)

    if args.cropmethod == 'random_center':
        choice = [torch.Tensor([1, 0, 0]), torch.Tensor([0, 0, 1]), torch.Tensor([1, 0, 1]),
                    torch.Tensor([-1, 0, 0]), torch.Tensor([-1, 1, 0])]
        for m in range(batch_size):
            index = random.sample(choice, 1)
            distance_list = []
            p_center = index[0]
            for n in range(args.pnum):
                distance_list.append(distance_squre(real_point[m, 0, n], p_center))
            distance_order = sorted(enumerate(distance_list), key=lambda x: x[1])

            for sp in range(args.crop_point_num):
                input_cropped1.data[m, 0, distance_order[sp][0]] = torch.FloatTensor([0, 0, 0])
                real_center.data[m, 0, sp] = real_point[m, 0, distance_order[sp][0]]
            
            crop_num_list = []
            for num_ in range(args.pnum - args.crop_point_num):
                crop_num_list.append(distance_order[num_ + args.crop_point_num][0])
            indices = torch.LongTensor(crop_num_list)
            input_partial[m,0]= torch.index_select(real_point[m,0],0,indices)
    
    input_partial = input_partial.squeeze().to(device)
    real_point = real_point.to(device)
    # print('the real_point is ', real_point.size())
    real_center = real_center.to(device)
    # print('the real_center is ', real_center.size())
    input_cropped1 = input_cropped1.to(device)
    # print('the input_cropped1 is ', input_cropped1.size())
    ############################
    # (1) data prepare
    ###########################
    
    ## 128  个点
    input_cropped1 = torch.squeeze(input_cropped1, 1) #2048

    partial_input = input_partial
    coarse_gt = real_point.squeeze()
    dense_gt = real_point.squeeze()

    partial_input = partial_input.to(device)
    coarse_gt = coarse_gt.to(device)
    coarse_gt = coarse_gt.permute(0, 2, 1) # (B, 3, 2048)
    # print('the coarse gt is ', coarse_gt.size())
    dense_gt = dense_gt.to(device)
    dense_gt = dense_gt.permute(0, 2, 1)
    partial_input = partial_input.permute(0, 2, 1) # (B, 3, 2048)
    # print('the partial input is ', partial_input.size())

    return  partial_input, coarse_gt, dense_gt, real_center, target 


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
    
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)

class PointLoss(nn.Module):
    def __init__(self):
        super(PointLoss,self).__init__()
        
    def forward(self,array1,array2):
        return chamfer_distance_numpy(array1,array2)

class PointLoss_eval(nn.Module):
    def __init__(self):
        super(PointLoss_eval,self).__init__()
        
    def forward(self,array1,array2):
        return chamfer_distance_numpy_eval(array1,array2)

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

def chamfer_distance_numpy_eval(array1, array2):
    batch_size, num_point, num_features = array1.shape
    # print(array1.shape)
    dist_all = []
    dist1 = []
    dist2 = []

    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])

        av_dist1 = av_dist1.detach().cpu().numpy()
        av_dist2 = av_dist2.detach().cpu().numpy()

        dist_all.append((av_dist1+av_dist2))
        dist1.append(av_dist1)
        dist2.append(av_dist2)

    return dist_all, dist1, dist2
################### L1 distance #######################


################### L2 distance #######################


if __name__ == '__main__':
    pc = read_point_cloud('examples/chair.pcd')
    draw_geometries([pc])
    
    # resample pcd
    pc.points = Vector3dVector(resample_pcd(np.asarray(pc.points), 16384))
    print(len(pc.points))
    draw_geometries([pc])

    # for i in range(8):
    #     pc = read_point_cloud('examples/{}.pcd'.format(i))
    #     draw_geometries([pc])

Drop or duplicate points so that pcd has exactly n points # print(real_point.size()) # print('the real_point is ', real_point.size()) # print('the real_center is ', real_center.size()) # print('the input_cropped1 is ', input_cropped1.size()) ############################ # (1) data prepare ########################### ## 128  个点 #2048 # (B, 3, 2048) # print('the coarse gt is ', coarse_gt.size()) # (B, 3, 2048) # print('the partial input is ', partial_input.size()) 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 #    distances = torch.sqrt(distances) # print(array1.shape) ################### L1 distance ####################### ################### L2 distance ####################### # resample pcd # for i in range(8): #     pc = read_point_cloud('examples/{}.pcd'.format(i)) #     draw_geometries([pc])