import pygame
import sys
import pandas as pd
import os
import numpy as np

# 初始化Pygame
pygame.init()

# 设置常量
WIDTH, HEIGHT = 1200, 1000
BACKGROUND_COLOR = (255, 255, 255)
LEFT_LANE_COLOR = (0, 128, 255)  # 左车道线颜色
RIGHT_LANE_COLOR = (255, 0, 0)    # 右车道线颜色
CENTER_LINE_COLOR = (0, 255, 0)    # 中心线颜色
VEHICLE_COLOR_DEFAULT = (0, 0, 255)  # 默认车辆颜色
X_OFFSET = 258109.4239876
Y_OFFSET = 4149969.964821


# 创建颜色映射
vehicle_colors = {
    1: (255, 0, 0),  # 红色
}

# 提供一些额外的颜色
additional_colors = [
    (0, 0, 255),     # 蓝色
    (0, 255, 0),     # 绿色
    (255, 165, 0),   # 橙色
    (128, 0, 128),   # 紫色
    (255, 255, 0),   # 黄色
]

# 创建窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Lane Visualization")
font = pygame.font.Font(None, 24)  # 默认字体

# 读取数据
csv_data_path = "./"
csv_path = os.path.join(csv_data_path, 'LaneMap.csv')
df = pd.read_csv(csv_path)

car_csv_path = os.path.join(csv_data_path, 'ObjState.csv')
car_df = pd.read_csv(car_csv_path)

# 提取车道线数据
left_lane_x = df['lane_left_x'].values
left_lane_y = df['lane_left_y'].values
right_lane_x = df['lane_right_x'].values
right_lane_y = df['lane_right_y'].values
center_line_x = df['centerLine_x'].values
center_line_y = df['centerLine_y'].values

# 提取车辆数据，将车辆信息按 playerId 分组
car_ids = car_df['playerId'].unique()
car_data = {player_id: car_df[car_df['playerId'] == player_id].reset_index(drop=True) for player_id in car_ids}

# 为未分配颜色的车辆分配颜色
for idx, player_id in enumerate(car_ids):
    print(player_id)
    if player_id not in vehicle_colors:
        vehicle_colors[player_id] = additional_colors[idx % len(additional_colors)]

# 计算坐标范围
min_x = min(left_lane_x.min(), right_lane_x.min(), center_line_x.min())
max_x = max(left_lane_x.max(), right_lane_x.max(), center_line_x.max())
min_y = min(left_lane_y.min(), right_lane_y.min(), center_line_y.min())
max_y = max(left_lane_y.max(), right_lane_y.max(), center_line_y.max())

# 计算缩放比例
scale_x = WIDTH / (max_x - min_x)
scale_y = HEIGHT / (max_y - min_y)
scale = min(scale_x, scale_y)

# 计算偏移量
offset_x = -min_x * scale
offset_y = -min_y * scale

# 主循环
car_indices = {player_id: 0 for player_id in car_ids}  # 保存每辆车的当前索引

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    # 清屏
    screen.fill(BACKGROUND_COLOR)

    # 绘制车道线
    for i in range(len(left_lane_x) - 1):
        start_pos = (left_lane_x[i] * scale + offset_x, HEIGHT - (left_lane_y[i] * scale + offset_y))
        end_pos = (left_lane_x[i + 1] * scale + offset_x, HEIGHT - (left_lane_y[i + 1] * scale + offset_y))
        pygame.draw.line(screen, LEFT_LANE_COLOR, start_pos, end_pos, 2)

    for i in range(len(right_lane_x) - 1):
        start_pos = (right_lane_x[i] * scale + offset_x, HEIGHT - (right_lane_y[i] * scale + offset_y))
        end_pos = (right_lane_x[i + 1] * scale + offset_x, HEIGHT - (right_lane_y[i + 1] * scale + offset_y))
        pygame.draw.line(screen, RIGHT_LANE_COLOR, start_pos, end_pos, 2)

    for i in range(len(center_line_x) - 1):
        start_pos = (center_line_x[i] * scale + offset_x, HEIGHT - (center_line_y[i] * scale + offset_y))
        end_pos = (center_line_x[i + 1] * scale + offset_x, HEIGHT - (center_line_y[i + 1] * scale + offset_y))
        pygame.draw.line(screen, CENTER_LINE_COLOR, start_pos, end_pos, 2)

    # 绘制车辆
    for player_id, car_df in car_data.items():
        car_index = car_indices[player_id]
        if car_index < len(car_df):
            posX = (car_df['posX'][car_index] + X_OFFSET)* scale + offset_x
            posY = HEIGHT - ((car_df['posY'][car_index] + Y_OFFSET) * scale + offset_y)  # 反转Y坐标
            posH = car_df['posH'][car_index]  # 车辆方向（弧度）

            # 计算矩形的四个角点
            # vehicle_length = car_df['dimX'][car_index] * scale  # 车辆长度，以像素表示
            # vehicle_width = car_df['dimY'][car_index] * scale    # 车辆宽度，以像素表示
            vehicle_length = 5.0 * scale  # 车辆长度，以像素表示
            vehicle_width = 2.0 * scale    # 车辆宽度，以像素表示

            corners = np.array([
                [-vehicle_length / 2, -vehicle_width / 2],
                [ vehicle_length / 2, -vehicle_width / 2],
                [ vehicle_length / 2,  vehicle_width / 2],
                [-vehicle_length / 2,  vehicle_width / 2]
            ])

            # 旋转矩形
            rotation_matrix = np.array([
                [np.cos(posH), -np.sin(posH)],
                [np.sin(posH),  np.cos(posH)]
            ])
            rotated_corners = corners.dot(rotation_matrix) + np.array([posX, posY])

            # 获取车辆颜色
            vehicle_color = vehicle_colors[player_id]

            # 创建矩形对象并绘制
            pygame.draw.polygon(screen, vehicle_color, rotated_corners)

            # 更新车辆的索引
            car_indices[player_id] += 1
            # 如果车的索引超过数据长度，重置为0
            if car_indices[player_id] >= len(car_df):
                car_indices[player_id] = 0

    # 更新显示
    pygame.display.flip()

    # 控制帧率
    pygame.time.Clock().tick(60)  # 每秒60帧