import cv2
import numpy as np
from ultralytics import YOLO
from collections import defaultdict
import time
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 加载 YOLOv11 模型
model = YOLO("yolo11n.pt")  # 确保模型文件名称正确

# 打开视频
cap = cv2.VideoCapture("test5.avi")
assert cap.isOpened(), "Error reading video file"

# 获取视频参数
w, h, fps = (
    int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
    int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
    int(cap.get(cv2.CAP_PROP_FPS)),
)

# 视频输出
video_writer = cv2.VideoWriter(
    "output_with_trajectory.mp4",
    cv2.VideoWriter_fourcc(*"mp4v"),
    fps,
    (w, h)
)

# 定义检测线
LINE_Y = h // 2
LINE_START = (0, LINE_Y)
LINE_END = (w, LINE_Y)
LINE_COLOR = (0, 255, 0)  # 绿色
LINE_THICKNESS = 2

# 存储轨迹和计数
track_history = defaultdict(list)
crossed_objects = set()
count = 0

# 类别名称映射（YOLO默认COCO类别）
CLASS_NAMES = {
    # 0: 'person',
    # 1: 'bicycle',
    2: "car",
    3: "motorcycle",
    5: "bus",
    7: "truck",
    # 可根据需要添加更多类别
}

# 速度估计参数
PIXELS_PER_METER = 15  # 需要根据实际场景校准
SECONDS_FOR_SPEED = 2  # 计算速度的时间窗口（秒）

# 轨迹可视化参数
TRAJECTORY_LENGTH = 30  # 保留的轨迹点数
TRAJECTORY_COLOR = (0, 255, 255)  # 黄色轨迹线

# 创建日志文件
with open("speed_log.txt", "w") as f:
    f.write("ID,Type,Direction,Speed(km/h),Timestamp\n")

# 处理视频
while cap.isOpened():
    success, frame = cap.read()
    if not success:
        break

    # YOLOv11 目标跟踪
    results = model.track(frame, persist=True, classes=[2, 3, 5, 7])  # 只检测车辆类

    # 绘制检测线
    cv2.line(frame, LINE_START, LINE_END, LINE_COLOR, LINE_THICKNESS)

    if results[0].boxes.id is not None:  # 如果有检测到目标
        boxes = results[0].boxes.xywh.cpu()
        track_ids = results[0].boxes.id.int().cpu().tolist()
        classes = results[0].boxes.cls.int().cpu().tolist()
        confidences = results[0].boxes.conf.cpu().tolist()  # 获取置信度

        for box, track_id, cls, conf in zip(boxes, track_ids, classes, confidences):
            x, y, width, height = box
            center_x, center_y = int(x), int(y)

            # 记录轨迹（坐标、时间戳、类别、置信度）
            track = track_history[track_id]
            track.append((center_x, center_y, time.time(), cls, conf))
            if len(track) > TRAJECTORY_LENGTH:
                track.pop(0)

            # 绘制轨迹线
            if len(track) >= 2:
                points = np.array([(x, y) for x, y, _, _, _ in track], dtype=np.int32)
                cv2.polylines(frame, [points], False, TRAJECTORY_COLOR, 2)

            # 获取类别名称
            vehicle_type = CLASS_NAMES.get(cls, "unknown")

            # 绘制带类别和ID的边界框
            label = f"{vehicle_type} {track_id} ({conf:.2f})"
            cv2.rectangle(
                frame,
                (int(x - width / 2), int(y - height / 2)),
                (int(x + width / 2), int(y + height / 2)),
                (255, 0, 0), 2
            )
            cv2.putText(
                frame, label,
                (int(x - width / 2), int(y - height / 2) - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.6,
                (255, 255, 255), 2
            )

            # 检测是否穿过线
            if len(track) >= 2:
                prev_y = track[-2][1]
                current_y = track[-1][1]

                # 方向判断和计数
                if (prev_y < LINE_Y and current_y >= LINE_Y) or (prev_y > LINE_Y and current_y <= LINE_Y):
                    if track_id not in crossed_objects:
                        count += 1
                        crossed_objects.add(track_id)

                        # 计算速度
                        if len(track) >= 2:
                            x1, y1, t1, _, _ = track[0]
                            x2, y2, t2, _, _ = track[-1]

                            distance_pixels = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
                            distance_meters = distance_pixels / PIXELS_PER_METER
                            time_diff = t2 - t1

                            if time_diff > 0:
                                speed_kmh = (distance_meters / time_diff) * 3.6
                                direction = "down" if current_y > prev_y else "up"

                                # 记录日志
                                with open("speed_log.txt", "a") as f:
                                    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
                                    f.write(
                                        f"{count},{track_id},{vehicle_type},{direction},{speed_kmh:.2f},{timestamp}\n")

                                # 在画面上显示速度
                                speed_text = f"{speed_kmh:.1f} km/h"
                                cv2.putText(
                                    frame, speed_text,
                                    (int(x + width / 2 + 10), int(y)),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                                    (0, 255, 255), 2
                                )

    # 显示统计信息
    info_text = f"Total: {count} | FPS: {fps} | Time: {time.strftime('%H:%M:%S')}"
    cv2.putText(
        frame, info_text,
        (20, 40), cv2.FONT_HERSHEY_SIMPLEX,
        0.8, (0, 0, 255), 2
    )

    # 写入视频
    video_writer.write(frame)

    # 显示实时画面
    cv2.imshow("Vehicle Tracking", frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

cap.release()
video_writer.release()
cv2.destroyAllWindows()