import open3d as o3d
import numpy as np
import math
import math
from functools import reduce
import sys
import copy

# pcd = o3d.io.read_point_cloud("data/stockpile.ply")
# pcd = o3d.io.read_point_cloud("data/2box-SG-t-MnQ0zoEQIBX.ply")
# pcd = o3d.io.read_point_cloud("data/4box-1-SG-t-AsSRO3iA0XT.ply")
# pcd = o3d.io.read_point_cloud("data/4box-2-SG-t-cp3kSewxMKi.ply")
pcd = o3d.io.read_point_cloud("data/4box-3-SG-t-kuuv0moEP0C.ply")
print(pcd)
assert (pcd.has_normals())
# o3d.visualization.draw_geometries([pcd])

axes = o3d.geometry.TriangleMesh.create_coordinate_frame()

# 找出地面所在平面
plane_model, inliers = pcd.segment_plane(distance_threshold=0.01,
                                         ransac_n=5,
                                         num_iterations=10000)

# 把地面上的点和其他点分开
[a, b, c, d] = plane_model
plane_pcd = pcd.select_by_index(inliers)
plane_pcd.paint_uniform_color([1.0, 0, 0])
stockpile_pcd = pcd.select_by_index(inliers, invert=True)
# stockpile_pcd = pcd
stockpile_pcd.paint_uniform_color([0, 0, 1.0])
# o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd, axes], point_show_normal=True)
o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd])
# o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd, axes])

# 移动到原点
plane_pcd = plane_pcd.translate((0,0,d/c))
stockpile_pcd = stockpile_pcd.translate((0,0,d/c))

# 旋转至与坐标轴对齐
cos_theta = c / math.sqrt(a**2 + b**2 + c**2)
sin_theta = math.sqrt((a**2+b**2)/(a**2 + b**2 + c**2))
u_1 = b / math.sqrt(a**2 + b**2 )
u_2 = -a / math.sqrt(a**2 + b**2)
rotation_matrix = np.array([[cos_theta + u_1**2 * (1-cos_theta), u_1*u_2*(1-cos_theta), u_2*sin_theta],
                            [u_1*u_2*(1-cos_theta), cos_theta + u_2**2*(1- cos_theta), -u_1*sin_theta],
                            [-u_2*sin_theta, u_1*sin_theta, cos_theta]])
plane_pcd.rotate(rotation_matrix)
stockpile_pcd.rotate(rotation_matrix)
o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd, axes])

# 用statistical outlier算法去除噪点
cl, ind = stockpile_pcd.remove_statistical_outlier(nb_neighbors=100,
                                                    std_ratio=1)
stockpile_pcd = stockpile_pcd.select_by_index(ind)
o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd, axes])

# 依据法向量去除噪点
kd_tree_3d = o3d.geometry.KDTreeFlann(stockpile_pcd)
noise_list_by_normal = []
for idx in range (0, len(stockpile_pcd.points)):
  [k, neighbor_idx_list, _] = kd_tree_3d.search_radius_vector_3d(stockpile_pcd.points[idx], 0.1)
  if len(neighbor_idx_list) == 0:
    continue
  normal_sum = np.zeros((3))
  count = 0
  for neib_idx in neighbor_idx_list:
    normal_sum += np.asarray(stockpile_pcd.normals[neib_idx])
    count += 1
  normal_mean = normal_sum / count
  normal_angle = normal_mean[0] * stockpile_pcd.normals[idx][0] + normal_mean[1] * stockpile_pcd.normals[idx][1] + normal_mean[2] * stockpile_pcd.normals[idx][2]
  if normal_angle < 0:
    print(normal_angle)
    noise_list_by_normal.append(idx)
noise_by_normal_pcd = stockpile_pcd.select_by_index(noise_list_by_normal)
noise_by_normal_pcd.paint_uniform_color([0, 1.0, 0])
stockpile_pcd = stockpile_pcd.select_by_index(noise_list_by_normal, invert=True)
o3d.visualization.draw_geometries([plane_pcd, noise_by_normal_pcd, axes])
o3d.visualization.draw_geometries([plane_pcd, stockpile_pcd, axes])


# 降采样
# downpcd = stockpile_pcd.voxel_down_sample(voxel_size=0.1)
downpcd = stockpile_pcd
print('有效点云：', np.asarray(downpcd.points))

# 对stockpile建立二维网格
bbox = downpcd.get_axis_aligned_bounding_box()
print('bbox: ', bbox)
bin_num_1d = 250
bin_size_x = (bbox.max_bound[0] - bbox.min_bound[0]) / bin_num_1d
bin_size_y = (bbox.max_bound[1] - bbox.min_bound[1]) / bin_num_1d
grid2d = np.zeros((bin_num_1d, bin_num_1d, 2))
print('bin_size: ', bin_size_x, bin_size_y)

points = np.asarray(downpcd.points)
bin_idx_x = -1
bin_idx_y = -1
for point in points:
  bin_idx_x = math.floor((point[0] - bbox.min_bound[0]) / (bbox.max_bound[0] - bbox.min_bound[0]) * bin_num_1d)
  if bin_idx_x <= -1: # 浮点数溢出
    bin_idx_x = 0
  if bin_idx_x >= bin_num_1d:
    bin_idx_x = bin_num_1d - 1

  bin_idx_y = math.floor((point[1] - bbox.min_bound[1]) / (bbox.max_bound[1] - bbox.min_bound[1]) * bin_num_1d)
  if bin_idx_y <= -1: # 浮点数溢出
    bin_idx_y = 0
  if bin_idx_y >= bin_num_1d:
    bin_idx_y = bin_num_1d - 1

  # print('bin before ', bin_idx_x, bin_idx_y, grid2d[bin_idx_x, bin_idx_y, 0], grid2d[bin_idx_x, bin_idx_y, 1])
  # print('point', point)
  grid2d[bin_idx_x, bin_idx_y, 0] += (point[2]) # 注意！认为地面高度为0！
  grid2d[bin_idx_x, bin_idx_y, 1] += 1
  # print('bin after  ', bin_idx_x, bin_idx_y, grid2d[bin_idx_x, bin_idx_y, 0], grid2d[bin_idx_x, bin_idx_y, 1])

# 对地面建立二维网格,空的格子就是被物体盖住的。
grid_ground = np.zeros((bin_num_1d, bin_num_1d))
for ground_point in np.asarray(plane_pcd.points):
  if ground_point[0] >= bbox.min_bound[0] and ground_point[0] <= bbox.max_bound[0] and ground_point[1] >= bbox.min_bound[1] and ground_point[1] <= bbox.max_bound[1]:
    bin_idx_x = math.floor((ground_point[0] - bbox.min_bound[0]) / (bbox.max_bound[0] - bbox.min_bound[0]) * bin_num_1d)
    if bin_idx_x <= -1: # 浮点数溢出
      bin_idx_x = 0
    if bin_idx_x >= bin_num_1d:
      bin_idx_x = bin_num_1d - 1

    bin_idx_y = math.floor((ground_point[1] - bbox.min_bound[1]) / (bbox.max_bound[1] - bbox.min_bound[1]) * bin_num_1d)
    if bin_idx_y <= -1: # 浮点数溢出
      bin_idx_y = 0
    if bin_idx_y >= bin_num_1d:
      bin_idx_y = bin_num_1d - 1

    grid_ground[bin_idx_x, bin_idx_y] += 1

for i in range(0, len(stockpile_pcd.points)):
  point = stockpile_pcd.points[i]
  bin_idx_x = math.floor((point[0] - bbox.min_bound[0]) / (bbox.max_bound[0] - bbox.min_bound[0]) * bin_num_1d)
  if bin_idx_x <= -1: # 浮点数溢出
    bin_idx_x = 0
  if bin_idx_x >= bin_num_1d:
    bin_idx_x = bin_num_1d - 1

  bin_idx_y = math.floor((point[1] - bbox.min_bound[1]) / (bbox.max_bound[1] - bbox.min_bound[1]) * bin_num_1d)
  if bin_idx_y <= -1: # 浮点数溢出
    bin_idx_y = 0
  if bin_idx_y >= bin_num_1d:
    bin_idx_y = bin_num_1d - 1

# 对stockpile建立二维kdtree以便查询邻居，对stockpile二维网格中没有点的网格做插值
downpcd_2d = copy.deepcopy(downpcd)
for point in downpcd_2d.points:
  point[2] = 0
# print('build kdtree...')
kd_tree_2d = o3d.geometry.KDTreeFlann(downpcd_2d)
# print('build over')

# 二维网格体积累积
sum = 0
interCount = 0

for i in range(0, bin_num_1d):
  for j in range(0, bin_num_1d):
    # print('~~~~~~~~~~~~~')
    if grid_ground[i, j] == 1:
      continue

    if grid2d[i, j, 1] == 0:
      # print('做插值……')
      interCount += 1
      [k, neighbor_idx_list, _] = kd_tree_2d.search_knn_vector_3d([bin_size_x * (i + 0.5) + bbox.min_bound[0], bin_size_y * (j + 0.5) + bbox.min_bound[1], 0], 50)
      
      # 针对物体有垂直的侧面的情况，拿k近邻中z值最小的
      neighbor_pcd = downpcd.select_by_index(neighbor_idx_list)
      neighbor_pcd_pos = np.asarray(neighbor_pcd.points)
      z_min = 100
      for pos in neighbor_pcd_pos:
        if pos[2] < z_min:
          z_min = pos[2]

      grid2d[i, j, 0] = z_min
      grid2d[i, j, 1] = 1

    if(grid2d[i, j, 1]):
      sum += (grid2d[i, j, 0] / grid2d[i, j, 1] * bin_size_x * bin_size_y)
    # print('本网格idx', i, j)
    # print('本网格sum',grid2d[i, j, 0])
    # print('本网格点数',grid2d[i, j, 1])
    
    # print('sum: ', sum)
    # print('~~~~~~~~~~~~~')

print('体积：', sum)
# print('插值次数：', interCount)

exit(0)