import numpy as np
import cv2
import matplotlib.pyplot as plt
import laspy
import sys
import os

# 将项目根目录加入 sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from team_code_fusiondrive.fusiondrive.configs.global_config import GlobalConfig
from team_code_fusiondrive.fusiondrive.utils import transfuser_utils as t_u
from team_code_fusiondrive.fusiondrive.utils.open3d_vis_util import draw_scenes


config = GlobalConfig()

# 读取图像（假设 img 是你的图像）
img = cv2.imread("/home/yaya/source/FusionDrive/tools/demo/0020.jpg")  # 读取图像
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # OpenCV 读取为 BGR，需要转换为 RGB

# 读取激光雷达点云数据 (XYZ 格式)
# 假设点云数据存储在 `lidar_points.laz`，格式为每行 "x y z"
las_object_new = laspy.read("/home/yaya/source/FusionDrive/tools/demo/0020.laz")
lidar_points = las_object_new.xyz

draw_scenes(lidar_points)

# # 取出点云数据
X, Y, Z = lidar_points[:, 0], lidar_points[:, 1], lidar_points[:, 2]

lidar_points = np.vstack((X, Y, Z, np.ones_like(Z)))

# 设定投影矩阵 (3×4)
projection_matrix = t_u.calculate_ego2image_proj(config).numpy()

# projection_matrix, projected_points = t_u.obtain_projection_points(
#     config, lidar_points
# ).numpy()
# u, v, Z = projected_points[:3]

# # 执行投影转换
projected_points = projection_matrix @ lidar_points  # 矩阵相乘

projected_points[:2] /= projected_points[2:3]  # 归一化

# # 获取投影后的像素坐标 (u, v)
u, v, Z = projected_points[0], projected_points[1], projected_points[2]

# 过滤掉超出图像范围的点
image_height, image_width, _ = img.shape
valid_mask = (u >= 0) & (u < image_width) & (v >= 0) & (v < image_height) & (Z > 0)
u, v, Z = u[valid_mask], v[valid_mask], Z[valid_mask]

# lidar_points = np.vstack((X, Y, Z))
# projected_points = t_u.obtain_projection_points(config, lidar_points).numpy()
# u, v, Z = projected_points[:3]

# 颜色编码 (根据深度 Z 值调整颜色)
depth_normalized = (Z - Z.min()) / (Z.max() - Z.min())  # 归一化深度
colors = plt.cm.jet(1 - depth_normalized)[:, :3]  # 伪彩色映射 (jet colormap)

# 画出投影点
fig, ax = plt.subplots(figsize=(10, 6))
ax.imshow(img)  # 显示原始图像
ax.scatter(u, v, c=colors, s=2, marker="o", alpha=0.8)  # 叠加点云
ax.set_title("Projected LiDAR Points onto Image")
ax.axis("off")  # 关闭坐标轴
plt.show()
