# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Face Matching: 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"
REFERENCE_IMAGE = "face_a.jpg"
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.4


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


def run_pipeline(source_desc, on_frame):
    # 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"{source_desc} ! 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")
    sink.connect("new-sample", on_frame)
    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)


# 1. Compute the reference embedding from the reference image.
ref = {"emb": None, "area": 0}


def on_reference(sink):
    sample = sink.emit("pull-sample")
    if sample is None:
        return Gst.FlowReturn.OK
    vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps())
    for (x, y, w, h), emb in face_embeddings(vf):
        if w * h > ref["area"]:
            ref["area"] = w * h
            ref["emb"] = emb
    return Gst.FlowReturn.OK


run_pipeline(f"filesrc location={REFERENCE_IMAGE} ! jpegdec", on_reference)
if ref["emb"] is None:
    raise SystemExit("No face detected in the reference image")
ref_emb = ref["emb"]

# 2. Process the video, boxing only faces that match the reference identity.
writer = {"w": None}
match_frames = 0


def on_video(sink):
    global match_frames
    sample = sink.emit("pull-sample")
    if sample is None:
        return Gst.FlowReturn.OK
    vf = VideoFrame(sample.get_buffer(), caps=sample.get_caps())
    matches = []
    for (x, y, w, h), emb in face_embeddings(vf):
        similarity = float(np.dot(emb, ref_emb))
        if similarity >= MATCH_THRESHOLD:
            matches.append((x, y, w, h, similarity))

    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, similarity in matches:
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(frame, f"MATCH {similarity:.2f}", (x, max(15, y - 8)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
    if matches:
        match_frames += 1
    writer["w"].write(frame)
    return Gst.FlowReturn.OK


run_pipeline(f"filesrc location={INPUT_VIDEO} ! decodebin3", on_video)
if writer["w"] is not None:
    writer["w"].release()
print(f"Frames with a matched face: {match_frames}", flush=True)
print(f"Saved: {OUTPUT_VIDEO}", flush=True)