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
import geopandas as gpd

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

# 读取数据
csv_data_path = "/home/kevin/kevin/zhaoyuan/sqlite3_demo/docker_build/preprocess_run/data/V2I_CSAE53-2020_RedLightViolationW_LST_01-01/"
# 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"))

# 创建颜色映射
# unique_vehicle_ids = car_df['playerId'].unique()
unique_vehicle_ids = (
    car_df["playerId"].unique()
    if "playerId" in car_df.columns
    else car_df.assign(playerId=1)["playerId"].unique()
)

colors = [
    "red",
    "blue",
    "green",
    "orange",
    "purple",
    "cyan",
    "magenta",
    "yellow",
    "black",
    "brown",
]
vehicle_colors = {
    vehicle_id: colors[i] for i, vehicle_id in enumerate(unique_vehicle_ids)
}

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

# 加载地图数据
gdfs = []
folder_path = (
    "/home/kevin/kevin/zhaoyuan/zhaoyuan_0227_v0.3/zhaoyuan/modules/map/data/shp"
)
file_names = [
    "hd_arrow.shp",
    "hd_lane_change.shp",
    "hd_link_change.shp",
    "hd_parkingspace.shp",
    "hd_boundary_change.shp",
    "hd_crosswalk.shp",
    "hd_intersection.shp",
    "hd_link_boundary.shp",
    "hd_noparkingarea.shp",
    "hd_trafficlight.shp",
    "hd_stopline.shp",
    "hd_lane.shp",
]

for file_name in file_names:
    gdf = gpd.read_file(f"{folder_path}/{file_name}")
    gdfs.append(gdf)

gdfs = [gdf for gdf in gdfs if not gdf.empty]
all_gdf = gpd.GeoDataFrame(pd.concat(gdfs, ignore_index=True))

lane_gdf = gdfs[file_names.index("hd_lane.shp")]
lane_gdf["geometry"] = lane_gdf.apply(
    lambda row: row["geometry"].buffer(row["lane_width"] / 2), axis=1
)

# 设置画布
fig, ax = plt.subplots(figsize=(12, 8))
lane_gdf.plot(ax=ax, color="gray", edgecolor="black", alpha=0.7)
all_gdf.plot(ax=ax, edgecolor="black", alpha=0.5, linestyle="--")



# 查找车辆位置的极值
min_x = car_df["posX"].min() + X_OFFSET
max_x = car_df["posX"].max() + X_OFFSET
min_y = car_df["posY"].min() + Y_OFFSET
max_y = car_df["posY"].max() + Y_OFFSET

# 添加缓冲区
buffer_x = 50  # 可以根据需要修改
buffer_y = 50  # 可以根据需要修改

# 设置适当的 x 和 y 轴范围
x_min, x_max = min_x - buffer_x, max_x + buffer_x
y_min, y_max = min_y - buffer_y, max_y + buffer_y

ax.set_xlim(x_min, x_max)
ax.set_ylim(y_min, y_max)


# 动态绘制车辆
def update(frame):
    # 清除之前的矩形
    for patch in ax.patches:
        patch.remove()
    for text in ax.texts:
        text.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
        print(f"row[posH]: {row['posH']}")
        print(f"row['speedH']: {row['speedH']}")
        # posH =np.pi/ 2 -  math.radians(row["posH"]) 
        posH =np.pi/ 2 -np.deg2rad(row["posH"]) 
        # posH = np.pi/ 2
        print(f"posH: {posH}")
        # dimX = row['dimX']np.deg2rad(degrees)
        # dimY = row['dimY']
        dimX = 5.99
        dimY = 2.065

        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])

        

        # # 更安全的坐标转换
        # rotated_corners = np.dot(corners, rotation_matrix) + np.array([[posX], [posY]])

        rotation_matrix = np.array(
            [[np.cos(posH), np.sin(posH)], [-np.sin(posH), np.cos(posH)]]
        )
        

        # 关键修复：使用广播机制自动对齐形状
        rotated_corners = np.dot(corners, rotation_matrix) + [posX, posY]


        vehicle_color = vehicle_colors[row["playerId"]]

        vehicle = patches.Polygon(
            rotated_corners,
            closed=True,
            fill=True,
            color=vehicle_color,
            alpha=0.5,
            edgecolor="black",
        )
        ax.add_patch(vehicle)
        # 添加文本标注，显示坐标值
        ax.text(
            posX,
            posY,
            f"({posX:.1f}, {posY:.1f})",
            fontsize=8,
            ha="center",
            va="center",
            color="black",
        )
        # 添加黑色圆点，显示车辆坐标位置
        # ax.dot(posX, posY, 'o', color='black', markersize=1)  # 你可以调整markersize来改变圆点大小
    plt.title(f"Time: {current_time:.5f}s")


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

# 添加标题并显示地图
plt.title("Map Visualization with Vehicle Animation", fontsize=15)
plt.show()