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=3,
                                         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])

# 移动到原点
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)

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

# 建立二维网格
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])

# 建立二维kdtree以便查询邻居，对二维网格中没有点的网格做插值
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 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)