import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.animation import FuncAnimation
import os

# 位置偏移常量
X_OFFSET = 258109.4239876
Y_OFFSET = 4149969.964821

# 读取数据
csv_data_path = "/home/kevin/kevin/zhaoyuan/zhaoyuan/data/processed/data_zhaoyuan1"
lane_map_df = pd.read_csv(os.path.join(csv_data_path, 'LaneMap.csv'))
car_df = pd.read_csv(os.path.join(csv_data_path, 'ObjState.csv'))

# 创建颜色映射
# 你可以通过 playerId 生成颜色，也可以使用 matplotlib 的 colormap
unique_vehicle_ids = car_df['playerId'].unique()
colors = ['red', 'blue', 'green', 'orange', 'purple', 'cyan', 'magenta', 'yellow', 'black', 'brown']
# colors = plt.cm.viridis(np.linspace(0, 1, len(unique_vehicle_ids)))  # 使用 Viridis colormap

# 创建颜色字典
vehicle_colors = {vehicle_id: colors[i] for i, vehicle_id in enumerate(unique_vehicle_ids)}

# 依据时间分组车辆数据，按 simTime 进行分组
grouped_car_data = car_df.groupby("simTime")

# 设置画布
fig, ax = plt.subplots()
ax.set_aspect("equal")
ax.set_xlim([lane_map_df["centerLine_x"].min() - 20, lane_map_df["centerLine_x"].max() + 20])
ax.set_ylim([lane_map_df["centerLine_y"].min() - 20, lane_map_df["centerLine_y"].max() + 20])

# 绘制车道线
ax.plot(lane_map_df["centerLine_x"], lane_map_df["centerLine_y"], color="red", linewidth=2, linestyle="--", label="Vehicle Center Line")
ax.plot(lane_map_df["lane_left_x"], lane_map_df["lane_left_y"], color="blue", linewidth=2, label="Left Lane Line")
ax.plot(lane_map_df["lane_right_x"], lane_map_df["lane_right_y"], color="blue", linewidth=2, label="Right Lane Line")

# 动态绘制车辆
def update(frame):
    # 清除之前的矩形
    for patch in ax.patches:
        patch.remove()
    
    # 获取当前时间的车辆数据
    current_time = list(grouped_car_data.groups.keys())[frame]  # 获取当前帧的时间
    current_data = grouped_car_data.get_group(current_time)

    # 遍历当前时间点上的所有车辆
    for index, row in current_data.iterrows():
        posX = row['posX'] + X_OFFSET
        posY = row['posY'] + Y_OFFSET
        posH = row['posH']  # 假设已在弧度或者直接为弧度
        dimX = row['dimX']
        dimY = row['dimY']

        # 计算矩形的四个角点
        corners = np.array([
            [-dimX / 2, -dimY / 2],
            [dimX / 2, -dimY / 2],
            [dimX / 2, dimY / 2],
            [-dimX / 2, dimY / 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[row['playerId']]
        
        # 创建矩形对象
        vehicle = patches.Polygon(
            rotated_corners, closed=True, fill=True, color=vehicle_color, alpha=0.5
        )
        ax.add_patch(vehicle)

    plt.title(f'Time: {current_time:.2f}s')

# 设置动画，frames为时间段的数量
ani = FuncAnimation(fig, update, frames=len(grouped_car_data.groups), repeat=True, interval=50)

plt.show()