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

import gi

gi.require_version("Gst", "1.0")
from gi.repository import Gst

Gst.init([])

import numpy as np
import cv2
from gstgva import VideoFrame

INPUT_VIDEO = "test_video.mp4"
OUTPUT_VIDEO = "output_dlstreamer.mp4"
DETECTION_MODEL = "intel/face-detection-adas-0001/FP16/face-detection-adas-0001.xml"
REID_MODEL = "intel/face-reidentification-retail-0095/FP16/face-reidentification-retail-0095.xml"
# For CPU: change "GPU" to "CPU". For NPU: change "GPU" to "NPU".
DEVICE = "GPU"
DET_THRESHOLD = 0.6
MATCH_THRESHOLD = 0.5

# Gallery of (numeric_id, embedding). recognize() returns an existing ID for a
# known face or enrolls a new one, keeping each person's ID stable over time.
gallery = []
next_id = 1


def recognize(embedding):
    global next_id
    best_index, best_sim = -1, 0.0
    for index, (_, gallery_emb) in enumerate(gallery):
        sim = float(np.dot(embedding, gallery_emb))
        if sim > best_sim:
            best_sim, best_index = sim, index
    if best_sim >= MATCH_THRESHOLD:
        person_id, gallery_emb = gallery[best_index]
        # Blend the embedding into the gallery entry to stay robust to pose.
        updated = 0.9 * gallery_emb + 0.1 * embedding
        gallery[best_index] = (person_id, updated / np.linalg.norm(updated))
        return person_id
    person_id = next_id
    next_id += 1
    gallery.append((person_id, embedding))
    print(f"Enrolled ID {person_id}", flush=True)
    return person_id


def face_embeddings(video_frame):
    """Yield ((x, y, w, h), normalized_embedding) for each classified face."""
    for region in video_frame.regions():
        rect = region.rect()
        emb = None
        for tensor in region.tensors():
            if tensor.is_detection():
                continue
            data = np.array(tensor.data(), dtype=np.float32)
            if data.size >= 256:
                emb = data[:256]
        if emb is None:
            continue
        emb = emb / (np.linalg.norm(emb) + 1e-9)
        yield (int(rect.x), int(rect.y), int(rect.w), int(rect.h)), emb


# Convert to BGR before inference so gvaclassify's embedding tensors survive to
# the appsink (a later format-changing videoconvert would strip them).
pipeline = Gst.parse_launch(
    f"filesrc location={INPUT_VIDEO} ! decodebin3 ! "
    "videoconvert ! video/x-raw,format=BGR ! "
    f"gvadetect model={DETECTION_MODEL} device={DEVICE} "
    f"threshold={DET_THRESHOLD} ! queue ! "
    f"gvaclassify model={REID_MODEL} device={DEVICE} ! queue ! "
    "appsink name=sink emit-signals=true sync=false max-buffers=4 drop=false"
)
sink = pipeline.get_by_name("sink")

writer = {"w": None}


def on_video(sink):
    sample = sink.emit("pull-sample")
    if sample is None:
        return Gst.FlowReturn.OK
    vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps())
    labeled = []
    for (x, y, w, h), emb in face_embeddings(vf):
        labeled.append((x, y, w, h, recognize(emb)))

    with vf.data() as mat:
        frame = mat.copy()

    if writer["w"] is None:
        frame_h, frame_w = frame.shape[:2]
        structure = sample.get_caps().get_structure(0)
        ok_fr, fps_n, fps_d = structure.get_fraction("framerate")
        fps = fps_n / fps_d if ok_fr and fps_d else 12
        writer["w"] = cv2.VideoWriter(
            OUTPUT_VIDEO, cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame_w, frame_h))

    for x, y, w, h, person_id in labeled:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(frame, f"ID {person_id}", (x, max(15, y - 8)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    writer["w"].write(frame)
    return Gst.FlowReturn.OK


sink.connect("new-sample", on_video)
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)
if writer["w"] is not None:
    writer["w"].release()
print(f"Total identities recognized: {len(gallery)}", flush=True)
print(f"Saved: {OUTPUT_VIDEO}", flush=True)