/*
 * GStreamer gstreamer-yolo26tensordecoder
 * Copyright (C) 2026 Collabora Ltd.
 *  @author: Olivier Crête <olivier.crete@collabora.com>
 *
 * gstyolo26tensordecoder.c
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA 02110-1301, USA.
 */

/**
 * SECTION:element-yolo26tensordec
 * @short_description: Decode tensors from YOLO26 object detection models
 *
 * This element can parse per-buffer inference tensor meta data generated by an upstream
 * inference element. It handles the YOLO26 object detection output format where
 * Non-Maximum Suppression is performed internally by the model, producing a fixed-size
 * list of finalised detections.
 *
 * The element expects a single output tensor with shape BATCH_SIZE x MAX_DET x 6
 * (float32, row-major) where each detection row contains:
 * X1 (left), Y1 (top), X2 (right), Y2 (bottom), confidence score, class index.
 * Bounding box coordinates are in absolute pixel space relative to the model input
 * resolution.
 *
 * ## Example launch commands:
 *
 * ### Object detection from a webcam:
 *
 * |[
 * gst-launch-1.0 v4l2src ! videoconvertscale !  \
 *   ! onnxinference model-file=yolo26n.onnx \
 *   ! yolo26tensordec score-threshold=0.3 label-file=COCO_classes.txt \
 *   ! objectdetectionoverlay ! videoconvertscale ! autovideosink
 * ]|
 *
 * ### Object detection from a JPEG image:
 *
 * |[
 * gst-launch-1.0 multifilesrc location=bus.jpg ! jpegdec ! videoconvertscale \
 *   ! onnxinference execution-provider=cpu model-file=yolo26n.onnx \
 *   ! yolo26tensordec score-threshold=0.3 label-file=COCO_classes.txt \
 *   ! objectdetectionoverlay ! videoconvertscale ! autovideosink
 * ]|
 *
 * The original repository of YOLO26 is located at
 * https://github.com/ultralytics/ultralytics.
 *
 * NOTE: This tensor decoder is for the "One-to-One head", which means
 * setting end2end=True in the training. When using the "One-to-Many
 * head", setting end2end=False, the YOLO26 model will produce the same output
 * as YOLOv8 and one should use the #yolov8tensordec tensor decoder.
 *
 * Since: 1.30
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "gstyolo26tensordecoder.h"

#include <math.h>

#include <gio/gio.h>

#include <gst/gst.h>
#include <gst/video/video.h>
#include <gst/analytics/analytics.h>

#define YOLO26_END2END_TENSOR_ID "yolo-26-end2end-out"

static const gfloat DEFAULT_SCORE_THRESHOLD = 0.3f;

GST_DEBUG_CATEGORY_STATIC (yolo26_tensor_decoder_debug);
#define GST_CAT_DEFAULT yolo26_tensor_decoder_debug

GST_ELEMENT_REGISTER_DEFINE (yolo26_tensor_decoder, "yolo26tensordec",
    GST_RANK_PRIMARY, GST_TYPE_YOLO26_TENSOR_DECODER);

enum
{
  PROP_0,
  PROP_LABEL_FILE,
  PROP_SCORE_THRESHOLD
};

static GQuark yolo26_tensor_id_quark;
static GQuark yolo26_none_quark;

static GstStaticPadTemplate gst_yolo26_tensor_decoder_src_template =
GST_STATIC_PAD_TEMPLATE ("src",
    GST_PAD_SRC,
    GST_PAD_ALWAYS,
    GST_STATIC_CAPS ("video/x-raw"));

/* *INDENT-OFF* */
static GstStaticPadTemplate gst_yolo26_tensor_decoder_sink_template =
GST_STATIC_PAD_TEMPLATE ("sink",
    GST_PAD_SINK,
    GST_PAD_ALWAYS,
    GST_STATIC_CAPS ("video/x-raw,"
      "tensors=(structure)["
        "tensorgroups,"
          YOLO26_END2END_TENSOR_ID"=(/uniquelist){"
            "(GstCaps)["
              "tensor/strided,"
                "tensor-id="YOLO26_END2END_TENSOR_ID","
                "dims=<(int)1,(int)[1,max],(int)6>,"
                "dims-order=(string)row-major,"
                "type=(string)float32;]"
          "}"
      "];"
    ));
/* *INDENT-ON* */

static GArray *read_labels (const char *labels_file);

static gboolean gst_yolo26_tensor_decoder_set_caps (GstBaseTransform * trans,
    GstCaps * incaps, GstCaps * outcaps);
static void gst_yolo26_tensor_decoder_set_property (GObject * object,
    guint prop_id, const GValue * value, GParamSpec * pspec);
static void gst_yolo26_tensor_decoder_get_property (GObject * object,
    guint prop_id, GValue * value, GParamSpec * pspec);
static void gst_yolo26_tensor_decoder_finalize (GObject * object);
static GstFlowReturn gst_yolo26_tensor_decoder_transform_ip (GstBaseTransform *
    trans, GstBuffer * buf);

G_DEFINE_TYPE (GstYolo26TensorDecoder, gst_yolo26_tensor_decoder,
    GST_TYPE_BASE_TRANSFORM);

static GArray *
read_labels (const char *labels_file)
{
  GArray *array;
  GFile *file = g_file_new_for_path (labels_file);
  GFileInputStream *file_stream;
  GDataInputStream *data_stream;
  GError *error = NULL;
  gchar *line;

  file_stream = g_file_read (file, NULL, &error);
  g_object_unref (file);
  if (!file_stream) {
    GST_WARNING ("Could not open file %s: %s\n", labels_file, error->message);
    g_clear_error (&error);
    return NULL;
  }

  data_stream = g_data_input_stream_new (G_INPUT_STREAM (file_stream));
  g_object_unref (file_stream);

  array = g_array_new (FALSE, FALSE, sizeof (GQuark));

  while ((line = g_data_input_stream_read_line (data_stream, NULL, NULL,
              &error))) {
    GQuark label = g_quark_from_string (line);
    g_array_append_val (array, label);
    g_free (line);
  }

  g_object_unref (data_stream);

  if (error) {
    GST_WARNING ("Could not open file %s: %s", labels_file, error->message);
    g_array_free (array, TRUE);
    g_clear_error (&error);
    return NULL;
  }

  if (array->len == 0) {
    g_array_free (array, TRUE);
    return NULL;
  }

  return array;
}

static void
gst_yolo26_tensor_decoder_class_init (GstYolo26TensorDecoderClass * klass)
{
  GObjectClass *gobject_class = (GObjectClass *) klass;
  GstElementClass *element_class = (GstElementClass *) klass;
  GstBaseTransformClass *basetransform_class = (GstBaseTransformClass *) klass;

  GST_DEBUG_CATEGORY_INIT (yolo26_tensor_decoder_debug, "yolo26tensordec",
      0, "YOLO26 tensor decoder");

  gobject_class->set_property = gst_yolo26_tensor_decoder_set_property;
  gobject_class->get_property = gst_yolo26_tensor_decoder_get_property;
  gobject_class->finalize = gst_yolo26_tensor_decoder_finalize;

  g_object_class_install_property (gobject_class, PROP_LABEL_FILE,
      g_param_spec_string ("label-file",
          "Label file", "Label file", NULL, (GParamFlags)
          (G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS)));

  g_object_class_install_property (gobject_class, PROP_SCORE_THRESHOLD,
      g_param_spec_float ("score-threshold",
          "Score threshold",
          "Threshold for deciding when to remove boxes based on score",
          0.0, 1.0, DEFAULT_SCORE_THRESHOLD,
          (GParamFlags) (G_PARAM_READWRITE | G_PARAM_STATIC_STRINGS)));

  gst_element_class_set_static_metadata (element_class,
      "YOLO26 object detection tensor decoder", "Tensordecoder/Video",
      "Decode tensors output from inference of YOLO26 object detection"
      " models with end2end enabled.",
      "Olivier Crête <olivier.crete@collabora.com>");

  gst_element_class_add_pad_template (element_class,
      gst_static_pad_template_get (&gst_yolo26_tensor_decoder_src_template));
  gst_element_class_add_pad_template (element_class,
      gst_static_pad_template_get (&gst_yolo26_tensor_decoder_sink_template));

  basetransform_class->transform_ip =
      GST_DEBUG_FUNCPTR (gst_yolo26_tensor_decoder_transform_ip);
  basetransform_class->set_caps =
      GST_DEBUG_FUNCPTR (gst_yolo26_tensor_decoder_set_caps);

  yolo26_tensor_id_quark =
      g_quark_from_static_string (YOLO26_END2END_TENSOR_ID);
  yolo26_none_quark = g_quark_from_static_string ("yolo26-none");
}

static void
gst_yolo26_tensor_decoder_init (GstYolo26TensorDecoder * self)
{
  self->score_threshold = DEFAULT_SCORE_THRESHOLD;
  gst_base_transform_set_passthrough (GST_BASE_TRANSFORM (self), FALSE);
  GST_PAD_UNSET_ACCEPT_INTERSECT (self->basetransform.sinkpad);
}

static void
gst_yolo26_tensor_decoder_finalize (GObject * object)
{
  GstYolo26TensorDecoder *self = GST_YOLO26_TENSOR_DECODER (object);

  g_free (self->label_file);
  g_clear_pointer (&self->labels, g_array_unref);

  G_OBJECT_CLASS (gst_yolo26_tensor_decoder_parent_class)->finalize (object);
}

static void
gst_yolo26_tensor_decoder_set_property (GObject * object, guint prop_id,
    const GValue * value, GParamSpec * pspec)
{
  GstYolo26TensorDecoder *self = GST_YOLO26_TENSOR_DECODER (object);
  const gchar *filename;

  switch (prop_id) {
    case PROP_LABEL_FILE:
    {
      GArray *labels;

      filename = g_value_get_string (value);
      labels = read_labels (filename);

      if (labels) {
        g_free (self->label_file);
        self->label_file = g_strdup (filename);
        g_clear_pointer (&self->labels, g_array_unref);
        self->labels = labels;
      } else {
        GST_WARNING_OBJECT (self, "Label file '%s' not found!", filename);
      }
      break;
    }
    case PROP_SCORE_THRESHOLD:
      self->score_threshold = g_value_get_float (value);
      break;
    default:
      G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
      break;
  }
}

static void
gst_yolo26_tensor_decoder_get_property (GObject * object, guint prop_id,
    GValue * value, GParamSpec * pspec)
{
  GstYolo26TensorDecoder *self = GST_YOLO26_TENSOR_DECODER (object);

  switch (prop_id) {
    case PROP_LABEL_FILE:
      g_value_set_string (value, self->label_file);
      break;
    case PROP_SCORE_THRESHOLD:
      g_value_set_float (value, self->score_threshold);
      break;
    default:
      G_OBJECT_WARN_INVALID_PROPERTY_ID (object, prop_id, pspec);
      break;
  }
}

static gboolean
gst_yolo26_tensor_decoder_set_caps (GstBaseTransform * trans, GstCaps * incaps,
    GstCaps * outcaps)
{
  GstYolo26TensorDecoder *self = GST_YOLO26_TENSOR_DECODER (trans);

  if (!gst_video_info_from_caps (&self->video_info, incaps)) {
    GST_ERROR_OBJECT (self, "Failed to parse caps");
    return FALSE;
  }

  return TRUE;
}

static const GstTensor *
gst_yolo26_tensor_decoder_get_tensor (GstYolo26TensorDecoder * self,
    GstBuffer * buf)
{
  GstMeta *meta = NULL;
  gpointer iter_state = NULL;
  static const gsize dims[] = { 1, G_MAXSIZE, 6 };

  if (!gst_buffer_get_meta (buf, GST_TENSOR_META_API_TYPE)) {
    GST_DEBUG_OBJECT (self,
        "missing tensor meta from buffer %" GST_PTR_FORMAT, buf);
    return NULL;
  }

  while ((meta = gst_buffer_iterate_meta_filtered (buf, &iter_state,
              GST_TENSOR_META_API_TYPE))) {
    GstTensorMeta *tensor_meta = (GstTensorMeta *) meta;
    const GstTensor *tensor;

    tensor = gst_tensor_meta_get_typed_tensor (tensor_meta,
        yolo26_tensor_id_quark, GST_TENSOR_DATA_TYPE_FLOAT32,
        GST_TENSOR_DIM_ORDER_ROW_MAJOR, G_N_ELEMENTS (dims), dims);

    if (!tensor)
      continue;

    return tensor;
  }

  return NULL;
}

static gboolean
gst_yolo26_tensor_decoder_decode_f32 (GstYolo26TensorDecoder * self,
    GstAnalyticsRelationMeta * rmeta, const GstTensor * detections_tensor)
{
  GstMapInfo map = GST_MAP_INFO_INIT;
  gsize max_det;
  const gfloat *detections;

  if (!gst_buffer_map (detections_tensor->data, &map, GST_MAP_READ)) {
    GST_ELEMENT_ERROR (self, STREAM, FAILED, (NULL),
        ("Cannot map detection tensor data"));
    return FALSE;
  }

  max_det = detections_tensor->dims[1];
  detections = (const gfloat *) map.data;

  GST_LOG_OBJECT (self, "Decoding up to %zu objects", max_det);

  for (gsize i = 0; i < max_det; i++) {
    const gfloat *det = &detections[i * 6];
    gfloat x1 = det[0];
    gfloat y1 = det[1];
    gfloat x2 = det[2];
    gfloat y2 = det[3];
    gfloat score = det[4];
    gfloat class_idx_f = det[5];
    gint x;
    gint y;
    gint width;
    gint height;
    gint class_idx;
    GQuark class_quark;

    if (score < self->score_threshold)
      continue;

    width = (gint) roundf (x2 - x1);
    height = (gint) roundf (y2 - y1);
    x = (gint) roundf (x1);
    y = (gint) roundf (y1);

    if (width <= 0 || height <= 0) {
      GST_WARNING_OBJECT (self,
          "Ignoring detection with non-positive bounding box dimensions [%"
          G_GSIZE_FORMAT "]: " "(%d, %d, %d, %d)", i, x, y, width, height);
      continue;
    }

    class_idx = (gint) class_idx_f;
    if (self->labels && class_idx >= 0 && class_idx < (gint) self->labels->len)
      class_quark = g_array_index (self->labels, GQuark, class_idx);
    else
      class_quark = yolo26_none_quark;

    if (!gst_analytics_relation_meta_add_od_mtd (rmeta, class_quark,
            x, y, width, height, score, NULL)) {
      GST_WARNING_OBJECT (self, "Could not add detection to meta");
      continue;
    }

    GST_LOG_OBJECT (self,
        "Detected class=%d label=%s score=%f bbox=(%d,%d) %dx%d",
        class_idx, g_quark_to_string (class_quark), score, x, y, width, height);
  }

  gst_buffer_unmap (detections_tensor->data, &map);
  return TRUE;
}

static GstFlowReturn
gst_yolo26_tensor_decoder_transform_ip (GstBaseTransform * trans,
    GstBuffer * buf)
{
  GstYolo26TensorDecoder *self = GST_YOLO26_TENSOR_DECODER (trans);
  GstAnalyticsRelationMeta *rmeta;
  const GstTensor *detections_tensor;

  detections_tensor = gst_yolo26_tensor_decoder_get_tensor (self, buf);
  if (detections_tensor == NULL) {
    GST_WARNING_OBJECT (self,
        "Could not find YOLO26 detection tensor, skipping");
    return GST_FLOW_OK;
  }

  rmeta = gst_buffer_add_analytics_relation_meta (buf);
  if (rmeta == NULL) {
    GST_ELEMENT_ERROR (trans, STREAM, FAILED, (NULL),
        ("Analytics Relation meta allocation failed"));
    return GST_FLOW_ERROR;
  }

  if (!gst_yolo26_tensor_decoder_decode_f32 (self, rmeta, detections_tensor))
    return GST_FLOW_ERROR;

  return GST_FLOW_OK;
}
