# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Loitering Detection: DLStreamer pipeline sample.
# See README.md for prerequisites and inputs. Run: python3 dlstreamer_sample.py

from collections import defaultdict
import json
import sys
import gi
gi.require_version("Gst", "1.0")
gi.require_version("GstAnalytics", "1.0")
gi.require_version("DLStreamerMeta", "1.0")
gi.require_version("DLStreamerWatermarkMeta", "1.0")
from gi.repository import Gst, GLib, GstAnalytics, DLStreamerMeta, DLStreamerWatermarkMeta

Gst.init([])

# Register DLStreamerMeta types so GstAnalytics iteration can handle them
_ov = sys.modules["gi.overrides.GstAnalytics"]
_ov.__mtd_types__[DLStreamerMeta.ZoneMtd.get_mtd_type()] = DLStreamerMeta.relation_meta_get_zone_mtd
_ov.__mtd_types__[DLStreamerMeta.TripwireMtd.get_mtd_type()] = DLStreamerMeta.relation_meta_get_tripwire_mtd

# Select the model: "yolo26n" (default) or "detr-resnet-50".
MODEL = "yolo26n"
MODEL_PATH = f"{MODEL}_openvino_model/{MODEL}.xml"
VIDEO = "VIRAT_S_000101.mp4"
ZONE_JSON = json.dumps([{
    "id": "loiter_zone",
    "type": "polygon",
    "points": [{"x": 0, "y": 200}, {"x": 300, "y": 200},
               {"x": 300, "y": 400}, {"x": 0, "y": 400}]
}])
LOITERING_SECONDS = 5.0

pipeline = Gst.parse_launch(
    f"filesrc location={VIDEO} ! decodebin3 ! videoconvert ! "
    f"gvadetect model={MODEL_PATH} device=GPU threshold=0.5 ! queue ! "
    f"gvatrack tracking-type=short-term-imageless ! queue ! "
    f"gvaanalytics name=analytics draw-zones=true ! "
    f"gvafpscounter ! identity name=probe ! gvawatermark name=watermark ! "
    f"videoconvert ! video/x-raw,format=I420 ! "
    f"openh264enc ! h264parse ! mp4mux ! filesink location=output_dlstreamer.mp4"
)

pipeline.get_by_name("analytics").set_property("zones", ZONE_JSON)
pipeline.get_by_name("watermark").set_property("displ-cfg", "hide-roi=person")

dwell = defaultdict(float)
last_seen = {}
flagged = set()

def on_buffer(pad, info):
    buf = info.get_buffer()
    now = buf.pts / Gst.SECOND if buf.pts != Gst.CLOCK_TIME_NONE else 0.0
    rmeta = GstAnalytics.buffer_get_analytics_relation_meta(buf)
    if not rmeta:
        return Gst.PadProbeReturn.OK

    # Iterate only over object-detection entries
    for od in rmeta.iter_on_type(GstAnalytics.ODMtd):
        label = GLib.quark_to_string(od.get_obj_type())
        if label != "person":
            continue

        # Find tracking ID via direct relation
        track_id = None
        for trk in od.iter_direct_related(GstAnalytics.RelTypes.RELATE_TO, GstAnalytics.TrackingMtd):
            success, tracking_id, *_ = trk.get_info()
            if success:
                track_id = tracking_id
            break
        if track_id is None:
            continue

        # Check if gvaanalytics placed this detection inside the zone
        in_zone = False
        for zone in od.iter_direct_related(GstAnalytics.RelTypes.RELATE_TO, DLStreamerMeta.ZoneMtd):
            in_zone = True
            break

        if not in_zone:
            continue

        # Accumulate dwell time for persons inside the zone
        dwell[track_id] += now - last_seen.get(track_id, now)
        last_seen[track_id] = now

        if dwell[track_id] >= LOITERING_SECONDS and track_id not in flagged:
            flagged.add(track_id)
            _, x, y, w, h, _ = od.get_location()
            print(f"LOITERING id={track_id} dwell={dwell[track_id]:.1f}s pos=({int(x + w/2)},{int(y + h)})")

    return Gst.PadProbeReturn.OK

pipeline.get_by_name("probe").get_static_pad("src").add_probe(Gst.PadProbeType.BUFFER, on_buffer)
pipeline.set_state(Gst.State.PLAYING)
pipeline.get_bus().timed_pop_filtered(Gst.CLOCK_TIME_NONE, Gst.MessageType.EOS | Gst.MessageType.ERROR)
pipeline.set_state(Gst.State.NULL)