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

ratio_factor = 0.9
neighbor_point_min_num = 8

# 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()
my_bbox = {
  'min_bound': (bbox.min_bound[0], bbox.min_bound[1], bbox.min_bound[2]),
  'max_bound': ((bbox.max_bound[0], bbox.max_bound[1], bbox.max_bound[2] + 0.5)),
}
print('my_bbox: ', my_bbox)
bin_num_1d = 30
bin_size_x = (my_bbox['max_bound'][0] - my_bbox['min_bound'][0]) / bin_num_1d
bin_size_y = (my_bbox['max_bound'][1] - my_bbox['min_bound'][1]) / bin_num_1d
bin_size_z = (my_bbox['max_bound'][2] - my_bbox['min_bound'][2]) / bin_num_1d
grid = np.zeros((bin_num_1d, bin_num_1d, bin_num_1d))
print('bin_size: ', bin_size_x, bin_size_y, bin_size_z)

points = np.asarray(downpcd.points)
for point in points: 
  bin_idx_x = math.floor((point[0] - my_bbox['min_bound'][0]) / (my_bbox['max_bound'][0] - my_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] - my_bbox['min_bound'][1]) / (my_bbox['max_bound'][1] - my_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

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

  # print('point', point)
  # print(grid[bin_idx_x, bin_idx_y, bin_idx_z])
  grid[bin_idx_x, bin_idx_y, bin_idx_z] += 1

min_point_num_threshold = 1
assert(grid[0, 0, bin_num_1d - 1] < min_point_num_threshold)
grid_voxel = np.zeros((bin_num_1d, bin_num_1d, bin_num_1d)) # 0: 初始状态；-1：已加入待处理列表； 1：已处理，无点；2：已处理，有点；
cellToProcessList = [(0, 0, bin_num_1d - 1)]
grid_voxel[0, 0, bin_num_1d - 1] = -1
count1 = 0
count2 = 0

def processCell():
  global min_point_num_threshold
  global grid
  global grid_voxel
  global cellToProcessList
  global count1
  global count2
  i = cellToProcessList[0][0]
  j = cellToProcessList[0][1]
  k = cellToProcessList[0][2]
  assert(grid_voxel[i, j, k] == -1)
  if grid[i, j, k] < min_point_num_threshold:
    grid_voxel[i, j, k] = 1
    count1 += 1
    if i - 1 >= 0 and grid_voxel[i - 1, j, k] == 0:
      cellToProcessList.append((i - 1, j, k))
      grid_voxel[i - 1, j, k] = -1
    if i + 1 <= bin_num_1d - 1 and grid_voxel[i + 1, j, k] == 0:
      cellToProcessList.append((i + 1, j, k))
      grid_voxel[i + 1, j, k] = -1
    if j - 1 >= 0 and grid_voxel[i, j - 1, k] == 0:
      cellToProcessList.append((i, j - 1, k))
      grid_voxel[i, j - 1, k] = -1
    if j + 1 <= bin_num_1d - 1 and grid_voxel[i, j + 1, k] == 0:
      cellToProcessList.append((i, j + 1, k))
      grid_voxel[i, j + 1, k] = -1
    if k - 1 >= 0 and grid_voxel[i, j, k - 1] == 0:
      cellToProcessList.append((i, j, k - 1))
      grid_voxel[i, j, k - 1] = -1
    if k + 1 <= bin_num_1d - 1 and grid_voxel[i, j, k + 1] == 0:
      cellToProcessList.append((i, j, k + 1))
      grid_voxel[i, j, k + 1] = -1
  else:
    grid_voxel[i, j, k] = 2
    count2 += 1
  cellToProcessList.pop(0)


while len(cellToProcessList) > 0:
  processCell()

print('空虚格子数：', count1)
print('边界格子数：', count2)
print('内部格子数：', bin_num_1d * bin_num_1d * bin_num_1d - count1 - 0.5 * count2)
print('物体占比：', (bin_num_1d * bin_num_1d * bin_num_1d - count1 - 0.5 * count2) / (bin_num_1d * bin_num_1d * bin_num_1d))
print('体积：', (bin_num_1d * bin_num_1d * bin_num_1d - count1 - 0.5 * count2) * bin_size_x * bin_size_y * bin_size_z)
print('外部体积：', (count1 + 0.5 * count2) * bin_size_x * bin_size_y * bin_size_z)
print('0, 0, 最高：', grid_voxel[0, 0, bin_num_1d - 1])
print('0, 0, 最低：', grid_voxel[0, 0, 0])