# 安装 vmspy
# - 下载与当前 Python 版本匹配的 vmspy 软件包
!unzip vmspy.zip
!./vmspy/install-vmspy-collab.sh

# 安装 YOLO（Ultralytics）
%pip install ultralytics
import ultralytics
from ultralytics import YOLO

# 加载 YOLO 模型（轻量版本）
model = YOLO('yolo11n.pt')

# 导入必要的模块
from tarfile import NUL  # 用作空返回值（通常建议使用 None）
import cv2
import time
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon

# 输入帧分辨率（用于归一化）
input_width = 640
input_height = 640

# 配置 live_video 对象（VMS 流接口）
live_video = vmspy.live_video()
live_video.set_keyframe_only(True)                 # 为了性能优化，仅使用关键帧
live_video.set_image_size(input_width, input_height)
live_video.set_pixel_format("BGR")                 # OpenCV 使用 BGR 格式
live_video.init("url.to.vmsserver", 3300, "admin", "admin")

# 定义监控区域（归一化坐标：0.0 ~ 1.0）
# 该区域也会在 VMS 服务器上进行可视化显示
monitoring_zone_points = [
    (0.3380152329749104, 0.3314180107526881),
    (0.42672491039426524, 0.3404905913978495),
    (0.41261200716845886, 0.9241599462365592),
    (0.1323700716845878, 0.9060147849462368)
]

monitoring_zone = Polygon(monitoring_zone_points)

# 在 VMS 画面上绘制监控区域
live_video.draw_polygon(monitoring_zone_points, border_color="purple", border_thickness=5)

# 配置 VMS 中的检测结果显示
live_video.set_detection_class(model.names)
live_video.set_detection_box(border_color="blue", border_thickness=5)
live_video.set_detection_caption(
    show_object_id=True,
    show_class_name=True,
    show_confidence=False
)

# 配置参数（单位：秒）
config = {
    "parking_violation_sec": 10,   # 判定为停车违规的时间阈值
    "periodic_check_sec": 10,      # 周期性重复上报的时间间隔
    "remove_criteria_sec": 10      # 移除不活跃目标的时间阈值
}

# 简单的跟踪逻辑，用于管理目标持续状态和事件触发
class TrackingHistory():
    def __init__(self):
        self.history = {}

    def update(self, key, frame_info):
        """
        更新指定对象（track_id）的跟踪信息。
        当满足上报条件时，返回初始的 frame_info。
        """
        timestamp = frame_info["timestamp"]

        if key in self.history:
            track_item = self.history[key]
            track_item["timestamp_last"] = timestamp

            # 检查是否到达再次上报的时间
            if timestamp > track_item["timestamp_report"]:
                track_item["timestamp_report"] = timestamp + config["periodic_check_sec"]
                print(f"{key}: REPORTED / n={len(self.history)}")
                return track_item["frameinfo0"]
        else:
            # 首次检测到该对象
            self.history[key] = {
                "frameinfo0": frame_info.copy(),  # 保存初始帧信息
                "timestamp_first": timestamp,
                "timestamp_last": timestamp,
                "timestamp_report": timestamp + config["parking_violation_sec"]
            }
            print(f"{key}: CREATED / n={len(self.history)}")

        return NUL  # 当前无需上报事件

    def remove_old_items(self, frame_info):
        """
        移除长时间未更新的对象。
        """
        current_timestamp = frame_info["timestamp"]
        min_sec = config["remove_criteria_sec"]

        keys_to_remove = [
            key for key, value in self.history.items()
            if current_timestamp - value["timestamp_last"] > min_sec
        ]

        for key in keys_to_remove:
            sec = current_timestamp - self.history[key]["timestamp_first"]
            del self.history[key]
            print(f"{key}: REMOVED ({sec}sec) / n={len(self.history)}")

        return len(keys_to_remove)

tracking = TrackingHistory()

# 从指定通道启动实时视频流
ch_no = 1
live_video.start(ch_no)

count = 0
while count < 10000:
    count += 1

    # 获取视频帧和元数据
    (frame_image, frame_info) = live_video.get_frame()

    # 结束条件：没有更多帧数据
    if frame_image.size == 0 and count > 2:
        print("\nEnd of stream")
        break

    # 执行目标检测 / 跟踪
    timestamp_detect_start = time.time_ns()
    results = model.track(frame_image, persist=True, verbose=False)
    # 可选：results = model.predict(frame_image, verbose=False)
    time_detect_elapsed = (time.time_ns() - timestamp_detect_start) / 1_000_000_000

    # 处理检测结果
    for box in results[0].boxes:

        # 获取跟踪 ID
        if box.id is not None:
            track_id = int(box.id[0].item())
        else:
            continue  # 跳过没有跟踪 ID 的目标

        # 将边界框转换为归一化坐标
        xyxy = box.xyxy[0].tolist()
        x = xyxy[0] / input_width
        y = xyxy[1] / input_height
        w = (xyxy[2] - xyxy[0]) / input_width
        h = (xyxy[3] - xyxy[1]) / input_height

        class_id = int(box.cls[0].item())
        conf = box.conf[0].item()

        # 检查目标中心点是否位于监控区域内
        if monitoring_zone.contains(Point(x + w / 2, y + h / 2)):
            name = model.names[class_id]

            # 将检测结果发送到 VMS
            live_video.input_detection(
                x, y, w, h,
                class_id=class_id,
                object_id=track_id,
                confidence=conf
            )

            # 更新跟踪状态
            frame_info0 = tracking.update(track_id, frame_info)

            # 如果满足违规条件，则上报事件
            if frame_info0 != NUL:
                duration = float(frame_info["timestamp"] - frame_info0["timestamp"])
                print(f"duration: {duration}")

                live_video.report_event(
                    frame_info0,
                    x, y, w, h,
                    duration=duration,
                    event_type_id=0,
                    class_id=class_id,
                    object_id=track_id
                )

    # 清理过期的跟踪对象
    tracking.remove_old_items(frame_info)

    # 将所有检测结果发送到 VMS
    live_video.send_detections(frame_info)