import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt

# 断面を取得し、重ね合わせた図を作成する関数
def process_sections(axis, levels, label1, label2, idx1, idx2, idx3):
    overlay_x = []
    overlay_y = []
    overlay_alpha = []

    # 各断面の点数をカウント
    section_counts = []
    for level in levels:
        section_points = points[(points[:, idx3] >= level - slice_thickness) & (points[:, idx3] <= level + slice_thickness)]
        section_counts.append(section_points.shape[0])

    # 最大・最小のデータ点数を取得（色の調整に使用）
    max_count = max(section_counts) if max(section_counts) > 0 else 1  # 0除算を防ぐ
    min_count = min(section_counts) if min(section_counts) > 0 else 1

    # 色の濃さを点数に応じて設定（点が多いほど薄く、少ないほど濃く）
    normalized_counts = [((max_count*1.1)-count) / (max_count-min_count + 1) for count in section_counts]
    for i, level in enumerate(levels):
        section_points = points[(points[:, idx3] >= level - slice_thickness) & (points[:, idx3] <= level + slice_thickness)]

        if section_points.shape[0] == 0:
            continue

        x, y = section_points[:, idx1], section_points[:, idx2]

        # 重ね合わせ用のデータを保存
        overlay_x.extend(x)
        overlay_y.extend(y)
        overlay_alpha.extend([normalized_counts[i]] * len(x))  # 各点に対応する透明度を設定

    # 重ね合わせた図を作成
    if len(overlay_x) > 0:
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.scatter(overlay_x, overlay_y, s=0.5, c=overlay_alpha, cmap="gray", alpha=0.8)
        ax.set_aspect("equal")
        ax.set_xlabel(label1)
        ax.set_ylabel(label2)
        ax.set_title(f"Overlaid {axis} Sections")
        ax.grid(True)
        plt.savefig(f"output_2d_overlay_{axis}.png", dpi=300, bbox_inches="tight")
        plt.close(fig)  # 画面表示を抑制

# PLYファイルのパスを指定
ply_file_path = "plydata/20250223_room.ply"

# PLYファイルを読み込む
point_cloud = o3d.io.read_point_cloud(ply_file_path)

# 点群の範囲を指定してトリミング
min_bound = np.array([-10, -10, -10])  # 最小座標 (X_min, Y_min, Z_min)
max_bound = np.array([10, 10, 0.5])     # 最大座標 (X_max, Y_max, Z_max)

# 指定範囲内の点のみ抽出
cropped_pc = point_cloud.crop(o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound))

# 可視化
o3d.visualization.draw_geometries([cropped_pc],
                                  window_name="PLY Viewer",
                                  width=800,
                                  height=600,
                                  left=50,
                                  top=50,
                                  point_show_normal=False)


# PLYファイルを読み込む
point_cloud = o3d.io.read_point_cloud(ply_file_path)

# 点群の範囲を指定してトリミング
min_bound = np.array([-10, -10, -10])  # 最小座標 (X_min, Y_min, Z_min)
max_bound = np.array([10, 10, 0.5])     # 最大座標 (X_max, Y_max, Z_max)

# 指定範囲内の点のみ抽出
cropped_pc = point_cloud.crop(o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound))

# 可視化
o3d.visualization.draw_geometries([cropped_pc],
                                  window_name="PLY Viewer",
                                  width=800,
                                  height=600,
                                  left=50,
                                  top=50,
                                  point_show_normal=False)

# 点群データを numpy 配列に変換
points = np.asarray(point_cloud.points)

# 断面の設定
num_slices = 20  # 断面の数
slice_thickness = 0.1  # 断面の厚み

# 各軸の範囲を取得
x_min, x_max = np.min(points[:, 0]), np.max(points[:, 0])
y_min, y_max = np.min(points[:, 1]), np.max(points[:, 1])
z_min, z_max = np.min(points[:, 2]), np.max(points[:, 2])

x_levels = np.linspace(x_min, x_max, num_slices + 2)[1:-1]  # 最上と最下を除外
y_levels = np.linspace(y_min, y_max, num_slices + 2)[1:-1]
z_levels = np.linspace(z_min, z_max, num_slices + 2)[1:-1]

# XY 平面の断面（Z方向）
process_sections("XY", z_levels, "X", "Y", 0, 1, 2)
# XZ 平面の断面（Y方向）
#process_sections("XZ", y_levels, "X", "Z", 0, 2, 1)
# YZ 平面の断面（X方向）
#process_sections("YZ", x_levels, "Y", "Z", 1, 2, 0)