# Install vmspy
# - Download the vmspy package that matches your system's Python version
!unzip vmspy.zip
!./vmspy/install-vmspy-collab.sh

# Install YOLO (Ultralytics)
%pip install ultralytics
import ultralytics
from ultralytics import YOLO

# Load YOLO model (lightweight version)
model = YOLO('yolo11n.pt')

# Import necessary modules
from tarfile import NUL  # Used as a null return value (note: None is generally recommended)
import cv2
import time
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon

# Input frame resolution (used for normalization)
input_width = 640
input_height = 640

# Configure the live_video object (VMS streaming interface)
live_video = vmspy.live_video()
live_video.set_keyframe_only(True)                 # Use keyframes only for performance optimization
live_video.set_image_size(input_width, input_height)
live_video.set_pixel_format("BGR")                 # OpenCV uses BGR format
live_video.init("url.to.vmsserver", 3300, "admin", "admin")

# Define the monitoring zone (normalized coordinates: 0.0 ~ 1.0)
# This zone will also be visualized on the VMS server
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)

# Draw the monitoring zone on the VMS display
live_video.draw_polygon(monitoring_zone_points, border_color="purple", border_thickness=5)

# Configure detection visualization in 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
)

# Configuration parameters (in seconds)
config = {
    "parking_violation_sec": 10,   # Time threshold to consider a parking violation
    "periodic_check_sec": 10,      # Interval for repeated reporting
    "remove_criteria_sec": 10      # Time to remove inactive tracked objects
}

# Simple tracking logic to manage object persistence and event triggering
class TrackingHistory():
    def __init__(self):
        self.history = {}

    def update(self, key, frame_info):
        """
        Update tracking information for a given object (track_id).
        Returns the initial frame_info when a reporting condition is met.
        """
        timestamp = frame_info["timestamp"]

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

            # Check if it's time to report again
            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:
            # First time seeing this object
            self.history[key] = {
                "frameinfo0": frame_info.copy(),  # Store initial frame info
                "timestamp_first": timestamp,
                "timestamp_last": timestamp,
                "timestamp_report": timestamp + config["parking_violation_sec"]
            }
            print(f"{key}: CREATED / n={len(self.history)}")

        return NUL  # No event to report

    def remove_old_items(self, frame_info):
        """
        Remove objects that have not been updated recently.
        """
        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()

# Start live video streaming from channel
ch_no = 1
live_video.start(ch_no)

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

    # Retrieve frame and metadata
    (frame_image, frame_info) = live_video.get_frame()

    # End condition: no more frames
    if frame_image.size == 0 and count > 2:
        print("\nEnd of stream")
        break

    # Run object detection / tracking
    timestamp_detect_start = time.time_ns()
    results = model.track(frame_image, persist=True, verbose=False)
    # Alternative: results = model.predict(frame_image, verbose=False)
    time_detect_elapsed = (time.time_ns() - timestamp_detect_start) / 1_000_000_000

    # Process detection results
    for box in results[0].boxes:

        # Extract tracking ID
        if box.id is not None:
            track_id = int(box.id[0].item())
        else:
            continue  # Skip objects without tracking ID

        # Convert bounding box to normalized coordinates
        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()

        # Check if object center is inside monitoring zone
        if monitoring_zone.contains(Point(x + w / 2, y + h / 2)):
            name = model.names[class_id]

            # Send detection to VMS
            live_video.input_detection(
                x, y, w, h,
                class_id=class_id,
                object_id=track_id,
                confidence=conf
            )

            # Update tracking state
            frame_info0 = tracking.update(track_id, frame_info)

            # If violation condition is met, report event
            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
                )

    # Cleanup stale tracking entries
    tracking.remove_old_items(frame_info)

    # Send all detections to VMS
    live_video.send_detections(frame_info)