# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import os
import logging
import warnings
import threading
from math import log
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple

import backoff
import numpy as np
import json
import pandas as pd

from nv_ingest_api.internal.primitives.nim import ModelInterface
import tritonclient.grpc as grpcclient
from nv_ingest_api.internal.primitives.nim.model_interface.decorators import global_cache
from nv_ingest_api.internal.primitives.nim.model_interface.decorators import lock
from nv_ingest_api.internal.primitives.nim.model_interface.decorators import multiprocessing_cache
from nv_ingest_api.internal.primitives.nim.model_interface.helpers import get_model_name
from nv_ingest_api.util.image_processing import scale_image_to_encoding_size
from nv_ingest_api.util.image_processing.transforms import numpy_to_base64

logger = logging.getLogger(__name__)

YOLOX_PAGE_DEFAULT_VERSION = "nemotron-page-elements-v3"

# Note: local-vs-remote selection is driven by `infer_protocol` ("local" vs "grpc"/"http")
# at the `create_inference_client(...)` layer, not by environment variables.

# yolox-page-elements-v2 and v3 common contants
YOLOX_PAGE_CONF_THRESHOLD = 0.01
YOLOX_PAGE_IOU_THRESHOLD = 0.5
YOLOX_PAGE_MIN_SCORE = 0.1
YOLOX_PAGE_NIM_MAX_IMAGE_SIZE = 512_000
YOLOX_PAGE_IMAGE_PREPROC_HEIGHT = 1024
YOLOX_PAGE_IMAGE_PREPROC_WIDTH = 1024
YOLOX_PAGE_IMAGE_FORMAT = os.getenv("YOLOX_PAGE_IMAGE_FORMAT", "PNG")

# yolox-page-elements-v3 contants
YOLOX_PAGE_FINAL_SCORE = YOLOX_PAGE_V3_FINAL_SCORE = {
    "table": 0.1,
    "chart": 0.01,
    "title": 0.1,
    "infographic": 0.01,
    "paragraph": 0.1,
    "header_footer": 0.1,
}
YOLOX_PAGE_CLASS_LABELS = YOLOX_PAGE_V3_CLASS_LABELS = [
    "table",
    "chart",
    "title",
    "infographic",
    "paragraph",
    "header_footer",
]

# yolox-page-elements-v2 contants
YOLOX_PAGE_V2_FINAL_SCORE = {"table": 0.1, "chart": 0.01, "infographic": 0.01}
YOLOX_PAGE_V2_CLASS_LABELS = [
    "table",
    "chart",
    "title",
    "infographic",
]


# yolox-graphic-elements-v1 contants
YOLOX_GRAPHIC_CONF_THRESHOLD = 0.01
YOLOX_GRAPHIC_IOU_THRESHOLD = 0.25
YOLOX_GRAPHIC_MIN_SCORE = 0.1
YOLOX_GRAPHIC_NIM_MAX_IMAGE_SIZE = 512_000


YOLOX_GRAPHIC_CLASS_LABELS = [
    "chart_title",
    "x_title",
    "y_title",
    "xlabel",
    "ylabel",
    "other",
    "legend_label",
    "legend_title",
    "mark_label",
    "value_label",
]


# yolox-table-structure-v1 contants
YOLOX_TABLE_CONF_THRESHOLD = 0.01
YOLOX_TABLE_IOU_THRESHOLD = 0.25
YOLOX_TABLE_MIN_SCORE = 0.1
YOLOX_TABLE_NIM_MAX_IMAGE_SIZE = 512_000

YOLOX_TABLE_IMAGE_PREPROC_HEIGHT = 1024
YOLOX_TABLE_IMAGE_PREPROC_WIDTH = 1024

YOLOX_TABLE_CLASS_LABELS = [
    "border",
    "cell",
    "row",
    "column",
    "header",
]


# YoloxModelInterfaceBase implements methods that are common to yolox-page-elements and yolox-graphic-elements
class YoloxModelInterfaceBase(ModelInterface):
    """
    An interface for handling inference with a Yolox object detection model, supporting both gRPC and HTTP protocols.
    """

    def __init__(
        self,
        nim_max_image_size: Optional[int] = None,
        conf_threshold: Optional[float] = None,
        iou_threshold: Optional[float] = None,
        min_score: Optional[float] = None,
        class_labels: Optional[List[str]] = None,
        endpoints: Optional[Tuple[str, str]] = None,
    ):
        """
        Initialize the YOLOX model interface.
        Parameters
        ----------
        """
        self.nim_max_image_size = nim_max_image_size
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.min_score = min_score
        self.class_labels = class_labels

        if endpoints:
            self._yolox_grpc_endpoint = endpoints[0]
            self._model_name = None
            self._grpc_uses_bls_value = None  # Resolved on first use
        else:
            self._yolox_grpc_endpoint = None
            self._model_name = None
            self._grpc_uses_bls_value = False

    def _resolve_yolox_model_name_if_needed(self) -> None:
        """Resolve model name and BLS flag from the gRPC endpoint on first use. Cached on the instance."""
        if self._yolox_grpc_endpoint is None:
            return
        if self._model_name is not None:
            return
        self._model_name = get_yolox_model_name(self._yolox_grpc_endpoint, default_model_name="yolox_ensemble")
        self._grpc_uses_bls_value = self._model_name == "pipeline"

    @property
    def model_name(self) -> Optional[str]:
        self._resolve_yolox_model_name_if_needed()
        return self._model_name

    @model_name.setter
    def model_name(self, value: Optional[str]) -> None:
        self._model_name = value

    @property
    def _grpc_uses_bls(self) -> bool:
        self._resolve_yolox_model_name_if_needed()
        return bool(self._grpc_uses_bls_value)

    def prepare_data_for_inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Prepare input data for inference by resizing images and storing their original shapes.

        Parameters
        ----------
        data : dict
            The input data containing a list of images.

        Returns
        -------
        dict
            The updated data dictionary with resized images and original image shapes.
        """
        if (not isinstance(data, dict)) or ("images" not in data):
            raise KeyError("Input data must be a dictionary containing an 'images' key with a list of images.")

        if not all(isinstance(x, np.ndarray) for x in data["images"]):
            raise ValueError("All elements in the 'images' list must be numpy.ndarray objects.")

        original_images = data["images"]
        data["original_image_shapes"] = [image.shape for image in original_images]

        return data

    def format_input(
        self, data: Dict[str, Any], protocol: str, max_batch_size: int, **kwargs
    ) -> Tuple[List[Any], List[Dict[str, Any]]]:
        """
        Format input data for the specified protocol, returning a tuple of:
          (formatted_batches, formatted_batch_data)
        where:
          - For gRPC: formatted_batches is a list of NumPy arrays, each of shape (B, H, W, C)
            with B <= max_batch_size.
          - For HTTP: formatted_batches is a list of JSON-serializable dict payloads.
          - In both cases, formatted_batch_data is a list of dicts that coalesce the original
            images and their original shapes in the same order as provided.

        Parameters
        ----------
        data : dict
            The input data to format. Must include:
              - "images": a list of numpy.ndarray images.
              - "original_image_shapes": a list of tuples with each image's (height, width),
                as set by prepare_data_for_inference.
        protocol : str
            The protocol to use ("grpc" or "http").
        max_batch_size : int
            The maximum number of images per batch.

        Returns
        -------
        tuple
            A tuple (formatted_batches, formatted_batch_data).

        Raises
        ------
        ValueError
            If the protocol is invalid.
        """

        # Helper functions to chunk a list into sublists of length up to chunk_size.
        def chunk_list(lst: list, chunk_size: int) -> List[list]:
            chunk_size = max(1, chunk_size)
            return [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)]

        def chunk_list_geometrically(lst: list, max_size: int) -> List[list]:
            # TRT engine in Yolox NIM (gRPC) only allows a batch size in powers of 2.
            chunks = []
            i = 0
            while i < len(lst):
                chunk_size = max(1, min(2 ** int(log(len(lst) - i, 2)), max_size))
                chunks.append(lst[i : i + chunk_size])
                i += chunk_size
            return chunks

        if protocol == "grpc":
            logger.debug("Formatting input for gRPC Yolox Ensemble model")
            b64_images = [numpy_to_base64(image, format=YOLOX_PAGE_IMAGE_FORMAT) for image in data["images"]]
            b64_chunks = chunk_list_geometrically(b64_images, max_batch_size)
            original_chunks = chunk_list_geometrically(data["images"], max_batch_size)
            shape_chunks = chunk_list_geometrically(data["original_image_shapes"], max_batch_size)

            batched_inputs = []
            formatted_batch_data = []
            for b64_chunk, orig_chunk, shapes in zip(b64_chunks, original_chunks, shape_chunks):
                input_array = np.array(b64_chunk, dtype=np.object_)

                if self._grpc_uses_bls:
                    # For BLS with dynamic batching (max_batch_size > 0), we need to add explicit batch dimension
                    # Shape [N] becomes [1, N] to indicate: batch of 1, containing N images
                    input_array = input_array.reshape(1, -1)
                    thresholds = np.array([[self.conf_threshold, self.iou_threshold]], dtype=np.float32)
                else:
                    current_batch_size = input_array.shape[0]
                    single_threshold_pair = [self.conf_threshold, self.iou_threshold]
                    thresholds = np.tile(single_threshold_pair, (current_batch_size, 1)).astype(np.float32)

                batched_inputs.append([input_array, thresholds])
                formatted_batch_data.append({"images": orig_chunk, "original_image_shapes": shapes})

            return batched_inputs, formatted_batch_data

        elif protocol == "http":
            logger.debug("Formatting input for HTTP Yolox model")
            content_list: List[Dict[str, Any]] = []
            for image in data["images"]:
                # Convert to uint8 if needed.
                if image.dtype != np.uint8:
                    image = (image * 255).astype(np.uint8)

                # Get original size directly from numpy array (width, height)
                original_size = (image.shape[1], image.shape[0])
                # Convert numpy array directly to base64 using OpenCV
                image_b64 = numpy_to_base64(image, format=YOLOX_PAGE_IMAGE_FORMAT)
                # Scale the image if necessary.
                scaled_image_b64, new_size = scale_image_to_encoding_size(
                    image_b64, max_base64_size=self.nim_max_image_size
                )
                if new_size != original_size:
                    logger.debug(f"Image was scaled from {original_size} to {new_size}.")

                content_list.append({"type": "image_url", "url": f"data:image/png;base64,{scaled_image_b64}"})

            # Chunk the payload content, the original images, and their shapes.
            content_chunks = chunk_list(content_list, max_batch_size)
            original_chunks = chunk_list(data["images"], max_batch_size)
            shape_chunks = chunk_list(data["original_image_shapes"], max_batch_size)

            payload_batches = []
            formatted_batch_data = []
            for chunk, orig_chunk, shapes in zip(content_chunks, original_chunks, shape_chunks):
                payload = {"input": chunk}
                payload_batches.append(payload)
                formatted_batch_data.append({"images": orig_chunk, "original_image_shapes": shapes})
            return payload_batches, formatted_batch_data

        else:
            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")

    def parse_output(self, response: Any, protocol: str, data: Optional[Dict[str, Any]] = None, **kwargs) -> Any:
        """
        Parse the output from the model's inference response.

        Parameters
        ----------
        response : Any
            The response from the model inference.
        protocol : str
            The protocol used ("grpc" or "http").
        data : dict, optional
            Additional input data passed to the function.

        Returns
        -------
        Any
            The parsed output data.

        Raises
        ------
        ValueError
            If an invalid protocol is specified or the response format is unexpected.
        """

        if protocol == "grpc":
            logger.debug("Parsing output from gRPC Yolox model")
            return response  # For gRPC, response is already a numpy array
        elif protocol == "http":
            logger.debug("Parsing output from HTTP Yolox model")

            processed_outputs = []

            batch_results = response.get("data", [])
            for detections in batch_results:
                new_bounding_boxes = {label: [] for label in self.class_labels}

                bounding_boxes = detections.get("bounding_boxes", [])
                for obj_type, bboxes in bounding_boxes.items():
                    for bbox in bboxes:
                        xmin = bbox["x_min"]
                        ymin = bbox["y_min"]
                        xmax = bbox["x_max"]
                        ymax = bbox["y_max"]
                        confidence = bbox["confidence"]

                        new_bounding_boxes[obj_type].append([xmin, ymin, xmax, ymax, confidence])

                processed_outputs.append(new_bounding_boxes)

            return processed_outputs
        else:
            raise ValueError("Invalid protocol specified. Must be 'grpc' or 'http'.")

    def process_inference_results(self, output: Any, protocol: str, **kwargs) -> List[Dict[str, Any]]:
        """
        Process the results of the Yolox model inference and return the final annotations.

        Parameters
        ----------
        output_array : np.ndarray
            The raw output from the Yolox model.
        kwargs : dict
            Additional parameters for processing, including thresholds and number of classes.

        Returns
        -------
        list[dict]
            A list of annotation dictionaries for each image in the batch.
        """
        if protocol == "http":
            # For http, the output already has postprocessing applied. Skip to table/chart expansion.
            results = output

        elif protocol == "grpc":
            results = []
            # For grpc, apply the same NIM postprocessing.
            for out in output:
                if isinstance(out, bytes):
                    out = out.decode("utf-8")
                if isinstance(out, dict):
                    continue
                results.append(json.loads(out))
        inference_results = self.postprocess_annotations(results, **kwargs)
        return inference_results

    def postprocess_annotations(self, annotation_dicts, **kwargs):
        raise NotImplementedError()

    def transform_normalized_coordinates_to_original(self, results, original_image_shapes):
        """ """
        transformed_results = []

        for annotation_dict, shape in zip(results, original_image_shapes):
            new_dict = {}
            for label, bboxes_and_scores in annotation_dict.items():
                new_dict[label] = []
                for bbox_and_score in bboxes_and_scores:
                    bbox = bbox_and_score[:4]
                    transformed_bbox = [
                        bbox[0] * shape[1],
                        bbox[1] * shape[0],
                        bbox[2] * shape[1],
                        bbox[3] * shape[0],
                    ]
                    transformed_bbox += bbox_and_score[4:]
                    new_dict[label].append(transformed_bbox)
            transformed_results.append(new_dict)

        return transformed_results


class YoloxPageElementsModelInterface(YoloxModelInterfaceBase):
    """
    An interface for handling inference with yolox-page-elements model, supporting both gRPC and HTTP protocols.
    """

    # ---- Local Nemotron singleton (lazy) ----
    _local_nemotron_lock = threading.Lock()
    _local_nemotron_instance = None
    _local_nemotron_device = None

    def __init__(
        self,
        version: str = YOLOX_PAGE_DEFAULT_VERSION,
        endpoints: Optional[Tuple[str, str]] = None,
        *,
        backend: Optional[str] = None,
    ):
        """
        Initialize the yolox-page-elements model interface.
        """
        self.version = version
        self.backend = (backend or "nim").strip().lower()

        super().__init__(
            nim_max_image_size=YOLOX_PAGE_NIM_MAX_IMAGE_SIZE,
            conf_threshold=YOLOX_PAGE_CONF_THRESHOLD,
            iou_threshold=YOLOX_PAGE_IOU_THRESHOLD,
            min_score=YOLOX_PAGE_MIN_SCORE,
            class_labels=YOLOX_PAGE_V3_CLASS_LABELS if self.version.endswith("-v3") else YOLOX_PAGE_V2_CLASS_LABELS,
            endpoints=endpoints,
        )

    @classmethod
    def _get_local_nemotron_page_elements_v3(cls, *, device: Optional[str] = None):
        """
        Lazily create and return a process-wide singleton NemotronPageElementsV3 model.

        Notes
        -----
        - Imports are deferred so the API can run without nemo-retriever/torch installed unless local backend is used.
        - Device selection is best-effort; tensors are moved to the chosen device before invocation.
        """
        if cls._local_nemotron_instance is not None:
            return cls._local_nemotron_instance, cls._local_nemotron_device

        with cls._local_nemotron_lock:
            if cls._local_nemotron_instance is not None:
                return cls._local_nemotron_instance, cls._local_nemotron_device

            try:
                import torch  # local-only dependency
            except Exception as e:  # pragma: no cover
                raise RuntimeError(
                    "Local YOLOX backend requested but 'torch' is not available. "
                    "Install the nemo-retriever/local model dependencies or run with infer_protocol='grpc'/'http'."
                ) from e

            try:
                # Import path used by nemo_retriever stages.
                from nemo_retriever.model.local.nemotron_page_elements_v3 import NemotronPageElementsV3  # type: ignore
            except Exception as e:  # pragma: no cover
                raise RuntimeError(
                    "Local YOLOX backend requested but 'nemo-retriever' package is not importable. "
                    "Ensure the 'nemo-retriever' project is installed on the PYTHONPATH, or "
                    "run with infer_protocol='grpc'/'http' and backend='local'."  # noqa: E501
                ) from e

            # Choose device (default: cuda if available else cpu).
            dev_str = (device or "").strip()
            if dev_str:
                dev = torch.device(dev_str)
            else:
                dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            model = NemotronPageElementsV3()
            # Best-effort: move underlying module to device if present.
            try:
                if hasattr(model, "_model") and model._model is not None:
                    model._model = model._model.to(dev)  # type: ignore[attr-defined]
            except Exception:
                # Some wrappers manage device internally; we'll always move tensors below.
                pass

            cls._local_nemotron_instance = model
            cls._local_nemotron_device = dev
            logger.info(f"Initialized local NemotronPageElementsV3 singleton on device={dev}.")

            return cls._local_nemotron_instance, cls._local_nemotron_device

    def infer(self, data: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        Unified inference entrypoint.

        - backend="nim": this interface is meant to be invoked via NimClient (existing behavior).
        - backend="local": runs a local singleton NemotronPageElementsV3 and returns yolox-style annotations.
        """
        if self.backend == "local":
            return self._infer_local_nemotron_page_elements_v3(data, **kwargs)

        raise RuntimeError(
            "YoloxPageElementsModelInterface.infer() is only supported for backend='local'. "
            "For backend='nim', construct a NimClient and call NimClient.infer()."
        )

    def _infer_local_nemotron_page_elements_v3(self, data: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        Local inference using Nemotron Page Elements v3.

        Returns
        -------
        List[Dict[str, Any]]
            Same structure as NimClient inference results for yolox-page-elements:
            per image: {label: [[x1,y1,x2,y2,score], ...], ...} with coordinates in original pixel space.
        """
        # Validate + record original shapes.
        data = self.prepare_data_for_inference(data)
        images = data.get("images", [])
        if not images:
            return []

        # Nemotron label id mapping used by nemo_retriever local stages.
        # Nemotron emits "text" where nv-ingest yolox expects "paragraph".
        id_to_label = {
            0: "table",
            1: "chart",
            2: "title",
            3: "infographic",
            4: "paragraph",  # nemotron: "text"
            5: "header_footer",
        }
        alias_label = {"text": "paragraph"}

        model, dev = self._get_local_nemotron_page_elements_v3(device=kwargs.get("device"))

        try:
            import torch  # local-only dependency
        except Exception as e:  # pragma: no cover
            raise RuntimeError("Local YOLOX backend requested but 'torch' is not available.") from e

        # Convert numpy HWC -> torch CHW uint8 on device; preprocess (resize/pad).
        tensors: List["torch.Tensor"] = []
        orig_shapes_hw: List[Tuple[int, int]] = []
        for img in images:
            if not isinstance(img, np.ndarray):
                raise ValueError("All elements in data['images'] must be numpy.ndarray objects for local inference.")
            if img.ndim != 3 or img.shape[2] not in (3, 4):
                raise ValueError(f"Expected HWC RGB/RGBA image; got shape={getattr(img, 'shape', None)}")

            # Ensure RGB uint8.
            if img.shape[2] == 4:
                img = img[:, :, :3]
            if img.dtype != np.uint8:
                img = (img * 255).astype(np.uint8) if np.issubdtype(img.dtype, np.floating) else img.astype(np.uint8)

            h, w = int(img.shape[0]), int(img.shape[1])
            orig_shapes_hw.append((h, w))

            t = torch.from_numpy(img).permute(2, 0, 1).contiguous()
            t = t.to(device=dev, dtype=torch.uint8, non_blocking=(getattr(dev, "type", "") == "cuda"))
            tensors.append(model.preprocess(t))

        # Best-effort batching (mirrors nemo_retriever local stage behavior).
        per_image_preds: Optional[List[Any]] = None
        with torch.inference_mode():
            with torch.autocast(device_type="cuda"):
                try:
                    batch_tensor = torch.stack(tensors, dim=0)
                    m = model.model  # may raise if remote-only
                    try:
                        raw = m(batch_tensor, list(orig_shapes_hw))
                    except Exception:
                        raw = m(batch_tensor, orig_shapes_hw[0])
                    raw0 = raw[0] if isinstance(raw, (tuple, list)) and len(raw) > 0 else raw
                    if isinstance(raw0, list) and len(raw0) == int(batch_tensor.shape[0]):
                        per_image_preds = raw0
                except Exception:
                    per_image_preds = None

                annotation_dicts: List[Dict[str, Any]] = []
                for i in range(len(tensors)):
                    preds = (
                        per_image_preds[i]
                        if per_image_preds is not None
                        else model.invoke(tensors[i], orig_shapes_hw[i])
                    )
                    boxes, labels, scores = model.postprocess(preds)

                ann: Dict[str, List[List[float]]] = {lab: [] for lab in YOLOX_PAGE_V3_CLASS_LABELS}
                for box, lab, score in zip(boxes, labels, scores):
                    # coerce label -> string
                    if isinstance(lab, str):
                        label_name = alias_label.get(lab, lab)
                    else:
                        try:
                            lab_i = int(lab.item()) if hasattr(lab, "item") else int(lab)
                        except Exception:
                            continue
                        label_name = id_to_label.get(lab_i, f"unknown_{lab_i}")

                    if label_name not in ann:
                        # Ignore unknown labels rather than polluting downstream keys.
                        continue

                    # box is xyxy normalized
                    if hasattr(box, "detach"):
                        box_list = box.detach().cpu().tolist()
                    else:
                        box_list = box.tolist() if hasattr(box, "tolist") else list(box)

                    try:
                        s = float(score.item()) if hasattr(score, "item") else float(score)
                    except Exception:
                        s = 0.0

                    ann[label_name].append(
                        [float(box_list[0]), float(box_list[1]), float(box_list[2]), float(box_list[3]), s]
                    )

                annotation_dicts.append(ann)

        # Reuse the existing nv-ingest postprocessing + coord transform.
        return self.postprocess_annotations(
            annotation_dicts,
            original_image_shapes=data.get("original_image_shapes", []),
        )

    def name(
        self,
    ) -> str:
        """
        Returns the name of the Yolox model interface.

        Returns
        -------
        str
            The name of the model interface.
        """

        return "yolox-page-elements"

    def postprocess_annotations(self, annotation_dicts, final_score=None, **kwargs):
        original_image_shapes = kwargs.get("original_image_shapes", [])

        # Decide v2 vs v3 post-processing.
        #
        # IMPORTANT: Do not infer v3 purely from output keys.
        # It's common for model outputs (especially gRPC JSON) to omit classes that had no detections
        # on a given page, which would incorrectly skip v3 post-processing for v3 models.
        running_v3 = str(getattr(self, "version", "")).endswith("-v3")
        if (not running_v3) and annotation_dicts:
            # Back-compat / best-effort auto-detect if version string isn't v3.
            # Treat presence of any v3-only classes as v3.
            running_v3 = any(k in annotation_dicts[0] for k in ("paragraph", "header_footer"))

        if not final_score:
            if running_v3:
                final_score = YOLOX_PAGE_V3_FINAL_SCORE
            else:
                final_score = YOLOX_PAGE_V2_FINAL_SCORE

        if running_v3:
            expected_final_score_keys = YOLOX_PAGE_V3_FINAL_SCORE
        else:
            expected_final_score_keys = [x for x in YOLOX_PAGE_V2_FINAL_SCORE if x != "title"]

        if (not isinstance(final_score, dict)) or (sorted(final_score.keys()) != sorted(expected_final_score_keys)):
            raise ValueError(
                "yolox-page-elements requires a dictionary of thresholds per each class: "
                f"{expected_final_score_keys}"
            )

        if annotation_dicts and running_v3:
            annotation_dicts = [
                postprocess_page_elements_v3(annotation_dict, labels=YOLOX_PAGE_V3_CLASS_LABELS)
                for annotation_dict in annotation_dicts
            ]
        else:
            # Table/chart expansion is "business logic" specific to nv-ingest
            annotation_dicts = [expand_table_bboxes(annotation_dict) for annotation_dict in annotation_dicts]
            annotation_dicts = [expand_chart_bboxes(annotation_dict) for annotation_dict in annotation_dicts]

        inference_results = []

        # Filter out bounding boxes below the final threshold
        # This final thresholding is "business logic" specific to nv-ingest
        for annotation_dict in annotation_dicts:
            new_dict = {}
            if running_v3:
                for label in YOLOX_PAGE_V3_CLASS_LABELS:
                    if label in annotation_dict and label in final_score:
                        threshold = final_score[label]
                        new_dict[label] = [bb for bb in annotation_dict[label] if bb[4] >= threshold]
            else:
                for label in YOLOX_PAGE_V2_CLASS_LABELS:
                    if label in annotation_dict:
                        if label == "title":
                            new_dict[label] = annotation_dict[label]
                        elif label in final_score:
                            threshold = final_score[label]
                            new_dict[label] = [bb for bb in annotation_dict[label] if bb[4] >= threshold]

            inference_results.append(new_dict)

        inference_results = self.transform_normalized_coordinates_to_original(inference_results, original_image_shapes)

        return inference_results


class YoloxGraphicElementsModelInterface(YoloxModelInterfaceBase):
    """
    An interface for handling inference with yolox-graphic-elemenents model, supporting both gRPC and HTTP protocols.
    """

    # ---- Local Nemotron singleton (lazy) ----
    _local_nemotron_lock = threading.Lock()
    _local_nemotron_instance = None
    _local_nemotron_device = None

    def __init__(
        self,
        endpoints: Optional[Tuple[str, str]] = None,
        *,
        backend: Optional[str] = None,
    ):
        """
        Initialize the yolox-graphic-elements model interface.
        """
        self.backend = (backend or "nim").strip().lower()
        super().__init__(
            nim_max_image_size=YOLOX_GRAPHIC_NIM_MAX_IMAGE_SIZE,
            conf_threshold=YOLOX_GRAPHIC_CONF_THRESHOLD,
            iou_threshold=YOLOX_GRAPHIC_IOU_THRESHOLD,
            min_score=YOLOX_GRAPHIC_MIN_SCORE,
            class_labels=YOLOX_GRAPHIC_CLASS_LABELS,
            endpoints=endpoints,
        )
        # Ensure a stable model_name even when endpoints are absent (common for local backend).
        if not hasattr(self, "model_name") or not getattr(self, "model_name", None):
            self.model_name = "yolox-graphic-elements"

    @classmethod
    def _get_local_nemotron_graphic_elements_v1(cls, *, device: Optional[str] = None):
        """
        Lazily create and return a process-wide singleton Nemotron Graphic Elements v1 model.

        Note
        ----
        This intentionally avoids importing the monorepo `nemo-retriever` package so local
        chart extraction can work in environments where only `nv_ingest_api` + the
        Nemotron HF packages are installed.
        """
        if cls._local_nemotron_instance is not None:
            return cls._local_nemotron_instance, cls._local_nemotron_device

        with cls._local_nemotron_lock:
            if cls._local_nemotron_instance is not None:
                return cls._local_nemotron_instance, cls._local_nemotron_device

            try:
                import torch  # local-only dependency
            except Exception as e:  # pragma: no cover
                raise RuntimeError(
                    "Local YOLOX (graphic-elements) backend requested but 'torch' is not available. "
                    "Install the nemo-retriever/local model dependencies or run with infer_protocol='grpc'/'http'."
                ) from e

            try:
                from nemotron_graphic_elements_v1.model import (  # type: ignore
                    define_model as define_model_graphic_elements,
                )
                from nemotron_graphic_elements_v1.model import (  # type: ignore
                    resize_pad as resize_pad_graphic_elements,
                )
            except Exception as e:  # pragma: no cover
                raise RuntimeError(
                    "Local YOLOX (graphic-elements) backend requested but the "
                    "`nemotron-graphic-elements-v1` package is not importable. "
                    "Install it (and its deps) or run with infer_protocol='grpc'/'http'."
                ) from e

            dev_str = (device or "").strip()
            if dev_str:
                dev = torch.device(dev_str)
            else:
                dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            # Create minimal wrapper matching the methods we use in local inference.
            # Use the same model id string used by the nemo_retriever shim.
            model_id = "Nemotron Graphic Elements v1"
            hf_model = define_model_graphic_elements(model_id)
            input_shape = (1024, 1024)

            # Best-effort: move underlying module to device if present.
            try:
                if hf_model is not None and hasattr(hf_model, "to"):
                    hf_model = hf_model.to(dev)
            except Exception:
                pass

            class _LocalNemotronGraphicElements:
                def __init__(self, m, shape):
                    self._model = m
                    self._shape = shape

                @property
                def input_shape(self):
                    return self._shape

                def preprocess(self, tensor):
                    return resize_pad_graphic_elements(tensor, self.input_shape)

                def invoke(self, input_data, orig_shape):
                    if self._model is None:
                        raise RuntimeError("Local Nemotron Graphic Elements v1 model was not initialized.")
                    return self._model(input_data, orig_shape)[0]

            model = _LocalNemotronGraphicElements(hf_model, input_shape)

            cls._local_nemotron_instance = model
            cls._local_nemotron_device = dev
            logger.info(f"Initialized local NemotronGraphicElementsV1 singleton on device={dev}.")

            return cls._local_nemotron_instance, cls._local_nemotron_device

    def infer(self, data: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        Unified inference entrypoint.

        - backend="nim": this interface is meant to be invoked via NimClient (existing behavior).
        - backend="local": runs a local NemotronGraphicElementsV1 and returns yolox-style annotations.
        """
        if self.backend == "local":
            return self._infer_local_nemotron_graphic_elements_v1(data, **kwargs)

        raise RuntimeError(
            "YoloxGraphicElementsModelInterface.infer() is only supported for backend='local'. "
            "For backend='nim', construct a NimClient and call NimClient.infer()."
        )

    def _infer_local_nemotron_graphic_elements_v1(self, data: Dict[str, Any], **kwargs) -> List[Dict[str, Any]]:
        """
        Local inference using Nemotron Graphic Elements v1.

        Returns:
          per image: {label: [[x1,y1,x2,y2], ...], ...} with coordinates in original pixel space.
        """
        data = self.prepare_data_for_inference(data)
        images = data.get("images", [])
        if not images:
            return []

        model, dev = self._get_local_nemotron_graphic_elements_v1(device=kwargs.get("device"))

        try:
            import torch  # local-only dependency
        except Exception as e:  # pragma: no cover
            raise RuntimeError("Local YOLOX backend requested but 'torch' is not available.") from e

        def _to_list(v: Any) -> Any:
            try:
                if isinstance(v, torch.Tensor):
                    return v.detach().cpu().tolist()
                if isinstance(v, np.ndarray):
                    return v.tolist()
            except Exception:
                pass
            return v

        def _extract_boxes_labels_scores(preds: Any) -> Tuple[Any, Any, Any]:
            # Common shapes emitted by HF detectors: dict with keys.
            if isinstance(preds, dict):
                # IMPORTANT: do not use `or` directly on torch.Tensors (truthiness is ambiguous).
                def _first_present(d: Dict[str, Any], keys: List[str]) -> Any:
                    for k in keys:
                        if k in d and d.get(k) is not None:
                            return d.get(k)
                    return None

                boxes = _first_present(preds, ["boxes", "bboxes", "pred_boxes"])
                labels = _first_present(preds, ["labels", "classes", "pred_classes"])
                scores = _first_present(preds, ["scores", "confidences", "pred_scores"])
                return boxes, labels, scores
            # Sometimes wrapped as list/tuple with dict payload.
            if isinstance(preds, (list, tuple)) and preds:
                p0 = preds[0]
                if isinstance(p0, dict):
                    return _extract_boxes_labels_scores(p0)
            return None, None, None

        # Convert images to torch tensors and run per-image (best-effort batching handled by underlying model).
        annotation_dicts: List[Dict[str, Any]] = []
        orig_shapes = data.get("original_image_shapes", [])
        for img, shape in zip(images, orig_shapes):
            if not isinstance(img, np.ndarray):
                raise ValueError("All elements in data['images'] must be numpy.ndarray objects for local inference.")
            if img.ndim != 3 or img.shape[2] not in (3, 4):
                raise ValueError(f"Expected HWC RGB/RGBA image; got shape={getattr(img, 'shape', None)}")

            # Ensure RGB uint8.
            if img.shape[2] == 4:
                img = img[:, :, :3]
            if img.dtype != np.uint8:
                img = (img * 255).astype(np.uint8) if np.issubdtype(img.dtype, np.floating) else img.astype(np.uint8)

            h, w = int(img.shape[0]), int(img.shape[1])
            t = torch.from_numpy(img).permute(2, 0, 1).contiguous()
            t = t.to(device=dev, dtype=torch.uint8, non_blocking=(getattr(dev, "type", "") == "cuda"))
            x = model.preprocess(t)

            with torch.inference_mode():
                with torch.autocast(device_type="cuda"):
                    preds = model.invoke(x, (h, w))

            boxes, labels, scores = _extract_boxes_labels_scores(preds)
            boxes_l = _to_list(boxes) if boxes is not None else []
            labels_l = _to_list(labels) if labels is not None else []
            scores_l = _to_list(scores) if scores is not None else []

            # Normalize shapes.
            if not isinstance(boxes_l, list):
                boxes_l = []
            if not isinstance(labels_l, list):
                labels_l = []
            if not isinstance(scores_l, list):
                scores_l = []

            n = min(len(boxes_l), len(labels_l), len(scores_l)) if scores_l else min(len(boxes_l), len(labels_l))
            if not scores_l:
                scores_l = [1.0 for _ in range(n)]

            # If labels are tensors nested, try flatten.
            if labels_l and isinstance(labels_l[0], list):
                labels_l = [x[0] if isinstance(x, list) and x else x for x in labels_l]
            if scores_l and isinstance(scores_l[0], list):
                scores_l = [x[0] if isinstance(x, list) and x else x for x in scores_l]

            # Determine if boxes are pixel coords (values > 1) vs normalized.
            mx = 0.0
            try:
                mx = float(max(max(b[:4]) for b in boxes_l[:n])) if n else 0.0
            except Exception:
                mx = 0.0
            is_normalized = mx <= 1.5

            ann: Dict[str, List[List[float]]] = {lab: [] for lab in YOLOX_GRAPHIC_CLASS_LABELS}

            for b, lab, sc in zip(boxes_l[:n], labels_l[:n], scores_l[:n]):
                if not (isinstance(b, list) and len(b) >= 4):
                    continue
                x1, y1, x2, y2 = [float(v) for v in b[:4]]
                if not is_normalized:
                    # Convert pixels -> normalized.
                    if w > 0 and h > 0:
                        x1, x2 = x1 / float(w), x2 / float(w)
                        y1, y2 = y1 / float(h), y2 / float(h)
                # Clamp.
                x1 = max(0.0, min(1.0, x1))
                y1 = max(0.0, min(1.0, y1))
                x2 = max(0.0, min(1.0, x2))
                y2 = max(0.0, min(1.0, y2))

                label_name: Optional[str] = None
                if isinstance(lab, str):
                    label_name = lab
                else:
                    try:
                        li = int(lab.item()) if hasattr(lab, "item") else int(lab)
                        if 0 <= li < len(YOLOX_GRAPHIC_CLASS_LABELS):
                            label_name = YOLOX_GRAPHIC_CLASS_LABELS[li]
                    except Exception:
                        label_name = None

                if label_name not in ann:
                    continue
                try:
                    score = float(sc.item()) if hasattr(sc, "item") else float(sc)
                except Exception:
                    score = 0.0

                ann[label_name].append([x1, y1, x2, y2, score])

            annotation_dicts.append(ann)

        # Reuse existing nv-ingest postprocessing + coord transform to pixel-space box dict.
        return self.postprocess_annotations(annotation_dicts, original_image_shapes=orig_shapes)

    def name(
        self,
    ) -> str:
        """
        Returns the name of the Yolox model interface.

        Returns
        -------
        str
            The name of the model interface.
        """

        return "yolox-graphic-elements"

    def postprocess_annotations(self, annotation_dicts, **kwargs):
        original_image_shapes = kwargs.get("original_image_shapes", [])

        annotation_dicts = self.transform_normalized_coordinates_to_original(annotation_dicts, original_image_shapes)

        inference_results = []

        # bbox extraction: additional postprocessing speicifc to nv-ingest
        for pred, shape in zip(annotation_dicts, original_image_shapes):
            bbox_dict = get_bbox_dict_yolox_graphic(
                pred,
                shape,
                self.class_labels,
                self.min_score,
            )
            # convert numpy arrays to list
            bbox_dict = {
                label: array.tolist() if isinstance(array, np.ndarray) else array for label, array in bbox_dict.items()
            }
            inference_results.append(bbox_dict)

        return inference_results


class YoloxTableStructureModelInterface(YoloxModelInterfaceBase):
    """
    An interface for handling inference with yolox-graphic-elemenents model, supporting both gRPC and HTTP protocols.
    """

    def __init__(self, endpoints: Optional[Tuple[str, str]] = None):
        """
        Initialize the yolox-graphic-elements model interface.
        """
        super().__init__(
            nim_max_image_size=YOLOX_TABLE_NIM_MAX_IMAGE_SIZE,
            conf_threshold=YOLOX_TABLE_CONF_THRESHOLD,
            iou_threshold=YOLOX_TABLE_IOU_THRESHOLD,
            min_score=YOLOX_TABLE_MIN_SCORE,
            class_labels=YOLOX_TABLE_CLASS_LABELS,
            endpoints=endpoints,
        )

    def name(
        self,
    ) -> str:
        """
        Returns the name of the Yolox model interface.

        Returns
        -------
        str
            The name of the model interface.
        """

        return "yolox-table-structure"

    def postprocess_annotations(self, annotation_dicts, **kwargs):
        original_image_shapes = kwargs.get("original_image_shapes", [])

        annotation_dicts = self.transform_normalized_coordinates_to_original(annotation_dicts, original_image_shapes)

        inference_results = []

        # bbox extraction: additional postprocessing speicifc to nv-ingest
        for pred, shape in zip(annotation_dicts, original_image_shapes):
            bbox_dict = get_bbox_dict_yolox_table(
                pred,
                shape,
                self.class_labels,
                self.min_score,
            )
            # convert numpy arrays to list
            bbox_dict = {
                label: array.tolist() if isinstance(array, np.ndarray) else array for label, array in bbox_dict.items()
            }
            inference_results.append(bbox_dict)

        return inference_results


def expand_table_bboxes(annotation_dict, labels=None):
    """
    Additional preprocessing for tables: extend the upper bounds to capture titles if any.
    Args:
        annotation_dict: output of postprocess_results, a dictionary with keys "table", "figure", "title"

    Returns:
        annotation_dict: same as input, with expanded bboxes for charts

    """
    if not labels:
        labels = list(annotation_dict.keys())

    if not annotation_dict or len(annotation_dict["table"]) == 0:
        return annotation_dict

    new_annotation_dict = {label: [] for label in labels}

    for label, bboxes in annotation_dict.items():
        for bbox_and_score in bboxes:
            bbox, score = bbox_and_score[:4], bbox_and_score[4]

            if label == "table":
                height = bbox[3] - bbox[1]
                bbox[1] = max(0.0, min(1.0, bbox[1] - height * 0.2))

            new_annotation_dict[label].append([round(float(x), 4) for x in bbox + [score]])

    return new_annotation_dict


def expand_chart_bboxes(annotation_dict, labels=None):
    """
    Expand bounding boxes of charts and titles based on the bounding boxes of the other class.
    Args:
        annotation_dict: output of postprocess_results, a dictionary with keys "table", "figure", "title"

    Returns:
        annotation_dict: same as input, with expanded bboxes for charts

    """
    if not labels:
        labels = list(annotation_dict.keys())

    if not annotation_dict or len(annotation_dict["chart"]) == 0:
        return annotation_dict

    bboxes = []
    confidences = []
    label_idxs = []
    for i, label in enumerate(labels):
        label_annotations = np.array(annotation_dict[label])

        if len(label_annotations) > 0:
            bboxes.append(label_annotations[:, :4])
            confidences.append(label_annotations[:, 4])
            label_idxs.append(np.full(len(label_annotations), i))
    bboxes = np.concatenate(bboxes)
    confidences = np.concatenate(confidences)
    label_idxs = np.concatenate(label_idxs)

    pred_wbf, confidences_wbf, labels_wbf = weighted_boxes_fusion(
        bboxes[:, None],
        confidences[:, None],
        label_idxs[:, None],
        merge_type="biggest",
        conf_type="max",
        iou_thr=0.01,
        class_agnostic=False,
    )

    chart_bboxes = pred_wbf[labels_wbf == 1]
    chart_confidences = confidences_wbf[labels_wbf == 1]
    title_bboxes = pred_wbf[labels_wbf == 2]

    found_title_idxs, no_found_title_idxs = [], []
    for i in range(len(chart_bboxes)):
        match = match_with_title_v1(chart_bboxes[i], title_bboxes, iou_th=0.01)
        if match is not None:
            chart_bboxes[i] = match[0]
            title_bboxes = match[1]
            found_title_idxs.append(i)
        else:
            no_found_title_idxs.append(i)

    chart_bboxes[found_title_idxs] = expand_boxes_v1(chart_bboxes[found_title_idxs], r_x=1.05, r_y=1.1)
    chart_bboxes[no_found_title_idxs] = expand_boxes_v1(chart_bboxes[no_found_title_idxs], r_x=1.1, r_y=1.25)

    annotation_dict = {
        "table": annotation_dict["table"],
        "chart": np.concatenate([chart_bboxes, chart_confidences[:, None]], axis=1).tolist(),
        "title": annotation_dict["title"],
    }

    return annotation_dict


def postprocess_page_elements_v3(annotation_dict, labels=None):
    """
    Expand bounding boxes of tables/charts/infographics and titles based on the bounding boxes of the other class.
    Args:
        annotation_dict: output of postprocess_results, a dictionary with keys:
        "table", "chart", "infographics", "title", "paragraph", "header_footer".

    Returns:
        annotation_dict: same as input, with expanded bboxes for page elements.

    """
    if not labels:
        labels = list(annotation_dict.keys())

    if not annotation_dict:
        return annotation_dict

    bboxes = []
    confidences = []
    label_idxs = []

    for i, label in enumerate(labels):
        if label not in annotation_dict:
            continue

        label_annotations = np.array(annotation_dict[label])

        if len(label_annotations) > 0:
            bboxes.append(label_annotations[:, :4])
            confidences.append(label_annotations[:, 4])
            label_idxs.append(np.full(len(label_annotations), i))

    if not bboxes:
        return annotation_dict

    bboxes = np.concatenate(bboxes)
    confidences = np.concatenate(confidences)
    label_idxs = np.concatenate(label_idxs)

    bboxes, confidences, label_idxs = remove_overlapping_boxes_using_wbf(bboxes, confidences, label_idxs)
    bboxes, confidences, label_idxs, found_title = match_structured_boxes_with_title(
        bboxes, confidences, label_idxs, labels
    )
    bboxes, confidences, label_idxs = expand_tables_and_charts(bboxes, confidences, label_idxs, labels, found_title)
    bboxes, confidences, label_idxs = postprocess_included_texts(bboxes, confidences, label_idxs, labels)

    order = np.argsort(bboxes[:, 1] * 10 + bboxes[:, 0])
    bboxes, confidences, label_idxs = bboxes[order], confidences[order], label_idxs[order]

    new_annotation_dict = {}
    for i, label in enumerate(labels):
        selected_bboxes = bboxes[label_idxs == i]
        selected_confidences = confidences[label_idxs == i]
        new_annotation_dict[label] = np.concatenate([selected_bboxes, selected_confidences[:, None]], axis=1).tolist()

    return new_annotation_dict


def weighted_boxes_fusion(
    boxes_list,
    scores_list,
    labels_list,
    iou_thr=0.5,
    skip_box_thr=0.0,
    conf_type="avg",
    merge_type="weighted",
    class_agnostic=False,
):
    """
    Custom wbf implementation that supports a class_agnostic mode and a biggest box fusion.
    Boxes are expected to be in normalized (x0, y0, x1, y1) format.

    Args:
        boxes_list (list[np array[n x 4]]): List of boxes. One list per model.
        scores_list (list[np array[n]]): List of confidences.
        labels_list (list[np array[n]]): List of labels
        iou_thr (float, optional): IoU threshold for matching. Defaults to 0.55.
        skip_box_thr (float, optional): Exclude boxes with score < skip_box_thr. Defaults to 0.0.
        conf_type (str, optional): Confidence merging type. Defaults to "avg".
        merge_type (str, optional): Merge type "weighted" or "biggest". Defaults to "weighted".
        class_agnostic (bool, optional): If True, merge boxes from different classes. Defaults to False.

    Returns:
        np array[N x 4]: Merged boxes,
        np array[N]: Merged confidences,
        np array[N]: Merged labels.
    """
    weights = np.ones(len(boxes_list))

    assert conf_type in ["avg", "max"], 'Conf type must be "avg" or "max"'
    assert merge_type in [
        "weighted",
        "biggest",
    ], 'Conf type must be "weighted" or "biggest"'

    filtered_boxes = prefilter_boxes(
        boxes_list,
        scores_list,
        labels_list,
        weights,
        skip_box_thr,
        class_agnostic=class_agnostic,
    )
    if len(filtered_boxes) == 0:
        return np.zeros((0, 4)), np.zeros((0,)), np.zeros((0,))

    overall_boxes = []
    for label in filtered_boxes:
        boxes = filtered_boxes[label]
        np.empty((0, 8))

        clusters = []

        # Clusterize boxes
        for j in range(len(boxes)):
            ids = [i for i in range(len(boxes)) if i != j]
            index, best_iou = find_matching_box_fast(boxes[ids], boxes[j], iou_thr)

            if index != -1:
                index = ids[index]
                cluster_idx = [clust_idx for clust_idx, clust in enumerate(clusters) if (j in clust or index in clust)]
                if len(cluster_idx):
                    cluster_idx = cluster_idx[0]
                    clusters[cluster_idx] = list(set(clusters[cluster_idx] + [index, j]))
                else:
                    clusters.append([index, j])
            else:
                clusters.append([j])

        for j, c in enumerate(clusters):
            if merge_type == "weighted":
                weighted_box = get_weighted_box(boxes[c], conf_type)
            elif merge_type == "biggest":
                weighted_box = get_biggest_box(boxes[c], conf_type)

            if conf_type == "max":
                weighted_box[1] = weighted_box[1] / weights.max()
            else:  # avg
                weighted_box[1] = weighted_box[1] * len(c) / weights.sum()
            overall_boxes.append(weighted_box)

    overall_boxes = np.array(overall_boxes)
    overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]
    boxes = overall_boxes[:, 4:]
    scores = overall_boxes[:, 1]
    labels = overall_boxes[:, 0]
    return boxes, scores, labels


def prefilter_boxes(boxes, scores, labels, weights, thr, class_agnostic=False):
    """
    Reformats and filters boxes.
    Output is a dict of boxes to merge separately.

    Args:
        boxes (list[np array[n x 4]]): List of boxes. One list per model.
        scores (list[np array[n]]): List of confidences.
        labels (list[np array[n]]): List of labels.
        weights (list): Model weights.
        thr (float): Confidence threshold
        class_agnostic (bool, optional): If True, merge boxes from different classes. Defaults to False.

    Returns:
        dict[np array [? x 8]]: Filtered boxes.
    """
    # Create dict with boxes stored by its label
    new_boxes = dict()

    for t in range(len(boxes)):
        if len(boxes[t]) != len(scores[t]):
            print(
                "Error. Length of boxes arrays not equal to length of scores array: {} != {}".format(
                    len(boxes[t]), len(scores[t])
                )
            )
            exit()

        if len(boxes[t]) != len(labels[t]):
            print(
                "Error. Length of boxes arrays not equal to length of labels array: {} != {}".format(
                    len(boxes[t]), len(labels[t])
                )
            )
            exit()

        for j in range(len(boxes[t])):
            score = scores[t][j]
            if score < thr:
                continue
            label = int(labels[t][j])
            box_part = boxes[t][j]
            x1 = float(box_part[0])
            y1 = float(box_part[1])
            x2 = float(box_part[2])
            y2 = float(box_part[3])

            # Box data checks
            if x2 < x1:
                warnings.warn("X2 < X1 value in box. Swap them.")
                x1, x2 = x2, x1
            if y2 < y1:
                warnings.warn("Y2 < Y1 value in box. Swap them.")
                y1, y2 = y2, y1
            if x1 < 0:
                warnings.warn("X1 < 0 in box. Set it to 0.")
                x1 = 0
            if x1 > 1:
                warnings.warn("X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.")
                x1 = 1
            if x2 < 0:
                warnings.warn("X2 < 0 in box. Set it to 0.")
                x2 = 0
            if x2 > 1:
                warnings.warn("X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.")
                x2 = 1
            if y1 < 0:
                warnings.warn("Y1 < 0 in box. Set it to 0.")
                y1 = 0
            if y1 > 1:
                warnings.warn("Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.")
                y1 = 1
            if y2 < 0:
                warnings.warn("Y2 < 0 in box. Set it to 0.")
                y2 = 0
            if y2 > 1:
                warnings.warn("Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.")
                y2 = 1
            if (x2 - x1) * (y2 - y1) == 0.0:
                warnings.warn("Zero area box skipped: {}.".format(box_part))
                continue

            # [label, score, weight, model index, x1, y1, x2, y2]
            b = [int(label), float(score) * weights[t], weights[t], t, x1, y1, x2, y2]

            label_k = "*" if class_agnostic else label
            if label_k not in new_boxes:
                new_boxes[label_k] = []
            new_boxes[label_k].append(b)

    # Sort each list in dict by score and transform it to numpy array
    for k in new_boxes:
        current_boxes = np.array(new_boxes[k])
        new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]]

    return new_boxes


def find_matching_box_fast(boxes_list, new_box, match_iou):
    """
    Reimplementation of find_matching_box with numpy instead of loops. Gives significant speed up for larger arrays
    (~100x). This was previously the bottleneck since the function is called for every entry in the array.
    """

    def bb_iou_array(boxes, new_box):
        # bb interesection over union
        xA = np.maximum(boxes[:, 0], new_box[0])
        yA = np.maximum(boxes[:, 1], new_box[1])
        xB = np.minimum(boxes[:, 2], new_box[2])
        yB = np.minimum(boxes[:, 3], new_box[3])

        interArea = np.maximum(xB - xA, 0) * np.maximum(yB - yA, 0)

        # compute the area of both the prediction and ground-truth rectangles
        boxAArea = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
        boxBArea = (new_box[2] - new_box[0]) * (new_box[3] - new_box[1])

        iou = interArea / (boxAArea + boxBArea - interArea)

        return iou

    if boxes_list.shape[0] == 0:
        return -1, match_iou

    ious = bb_iou_array(boxes_list[:, 4:], new_box[4:])
    # ious[boxes[:, 0] != new_box[0]] = -1

    best_idx = np.argmax(ious)
    best_iou = ious[best_idx]

    if best_iou <= match_iou:
        best_iou = match_iou
        best_idx = -1

    return best_idx, best_iou


def get_biggest_box(boxes, conf_type="avg"):
    """
    Merges boxes by using the biggest box.

    Args:
        boxes (np array [n x 8]): Boxes to merge.
        conf_type (str, optional): Confidence merging type. Defaults to "avg".

    Returns:
        np array [8]: Merged box.
    """
    box = np.zeros(8, dtype=np.float32)
    box[4:] = boxes[0][4:]
    conf_list = []
    w = 0
    for b in boxes:
        box[4] = min(box[4], b[4])
        box[5] = min(box[5], b[5])
        box[6] = max(box[6], b[6])
        box[7] = max(box[7], b[7])
        conf_list.append(b[1])
        w += b[2]

    box[0] = merge_labels(np.array([b[0] for b in boxes]), np.array([b[1] for b in boxes]))
    #     print(box[0], np.array([b[0] for b in boxes]))

    box[1] = np.max(conf_list) if conf_type == "max" else np.mean(conf_list)
    box[2] = w
    box[3] = -1  # model index field is retained for consistency but is not used.
    return box


def merge_labels(labels, confs):
    """
    Custom function for merging labels.
    If all labels are the same, return the unique value.
    Else, return the label of the most confident non-title (class 2) box.

    Args:
        labels (np array [n]): Labels.
        confs (np array [n]): Confidence.

    Returns:
        int: Label.
    """
    if len(np.unique(labels)) == 1:
        return labels[0]
    else:  # Most confident and not a title
        confs = confs[confs != 2]
        labels = labels[labels != 2]
        return labels[np.argmax(confs)]


def match_with_title_v1(chart_bbox, title_bboxes, iou_th=0.01):
    if not len(title_bboxes):
        return None

    dist_above = np.abs(title_bboxes[:, 3] - chart_bbox[1])
    dist_below = np.abs(chart_bbox[3] - title_bboxes[:, 1])

    dist_left = np.abs(title_bboxes[:, 0] - chart_bbox[0])

    ious = bb_iou_array(title_bboxes, chart_bbox)

    matches = None
    if np.max(ious) > iou_th:
        matches = np.where(ious > iou_th)[0]
    else:
        dists = np.min([dist_above, dist_below], 0)
        dists += dist_left
        #         print(dists)
        if np.min(dists) < 0.1:
            matches = [np.argmin(dists)]

    if matches is not None:
        new_bbox = chart_bbox
        for match in matches:
            new_bbox = merge_boxes(new_bbox, title_bboxes[match])
        title_bboxes = title_bboxes[[i for i in range(len(title_bboxes)) if i not in matches]]
        return new_bbox, title_bboxes

    else:
        return None


def bb_iou_array(boxes, new_box):
    # bb interesection over union
    xA = np.maximum(boxes[:, 0], new_box[0])
    yA = np.maximum(boxes[:, 1], new_box[1])
    xB = np.minimum(boxes[:, 2], new_box[2])
    yB = np.minimum(boxes[:, 3], new_box[3])

    interArea = np.maximum(xB - xA, 0) * np.maximum(yB - yA, 0)

    # compute the area of both the prediction and ground-truth rectangles
    boxAArea = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    boxBArea = (new_box[2] - new_box[0]) * (new_box[3] - new_box[1])

    iou = interArea / (boxAArea + boxBArea - interArea)

    return iou


def merge_boxes(b1, b2):
    b = b1.copy()
    b[0] = min(b1[0], b2[0])
    b[1] = min(b1[1], b2[1])
    b[2] = max(b1[2], b2[2])
    b[3] = max(b1[3], b2[3])
    return b


def expand_boxes_v1(boxes, r_x=1, r_y=1):
    dw = (boxes[:, 2] - boxes[:, 0]) / 2 * (r_x - 1)
    boxes[:, 0] -= dw
    boxes[:, 2] += dw

    dh = (boxes[:, 3] - boxes[:, 1]) / 2 * (r_y - 1)
    boxes[:, 1] -= dh
    boxes[:, 3] += dh

    boxes = np.clip(boxes, 0, 1)
    return boxes


def get_weighted_box(boxes, conf_type="avg"):
    """
    Merges boxes by using the weighted fusion.

    Args:
        boxes (np array [n x 8]): Boxes to merge.
        conf_type (str, optional): Confidence merging type. Defaults to "avg".

    Returns:
        np array [8]: Merged box.
    """
    box = np.zeros(8, dtype=np.float32)
    conf = 0
    conf_list = []
    w = 0
    for b in boxes:
        box[4:] += b[1] * b[4:]
        conf += b[1]
        conf_list.append(b[1])
        w += b[2]

    box[0] = merge_labels(np.array([b[0] for b in boxes]), np.array([b[1] for b in boxes]))

    box[1] = np.max(conf_list) if conf_type == "max" else np.mean(conf_list)
    box[2] = w
    box[3] = -1  # model index field is retained for consistency but is not used.
    box[4:] /= conf
    return box


def batched_overlaps(A, B):
    """
    Calculate the Intersection over Union (IoU) between
    two sets of bounding boxes in a batched manner.
    Normalization is modified to only use the area of A boxes, hence computing the overlaps.
    Args:
        A (ndarray): Array of bounding boxes of shape (N, 4) in format [x1, y1, x2, y2].
        B (ndarray): Array of bounding boxes of shape (M, 4) in format [x1, y1, x2, y2].
    Returns:
        ndarray: Array of IoU values of shape (N, M) representing the overlaps
         between each pair of bounding boxes.
    """
    A = A.copy()
    B = B.copy()

    A = A[None].repeat(B.shape[0], 0)
    B = B[:, None].repeat(A.shape[1], 1)

    low = np.s_[..., :2]
    high = np.s_[..., 2:]

    A, B = A.copy(), B.copy()
    A[high] += 1
    B[high] += 1

    intrs = (np.maximum(0, np.minimum(A[high], B[high]) - np.maximum(A[low], B[low]))).prod(-1)
    ious = intrs / (A[high] - A[low]).prod(-1)

    return ious


def find_boxes_inside(boxes, boxes_to_check, threshold=0.9):
    """
    Find all boxes that are inside another box based on
    the intersection area divided by the area of the smaller box,
    and removes them.
    """
    overlaps = batched_overlaps(boxes_to_check, boxes)
    to_keep = (overlaps >= threshold).sum(0) <= 1
    return boxes_to_check[to_keep]


def get_bbox_dict_yolox_graphic(preds, shape, class_labels, threshold_=0.1) -> Dict[str, np.ndarray]:
    """
    Extracts bounding boxes from YOLOX model predictions:
    - Applies thresholding
    - Reformats boxes
    - Cleans the `other` detections: removes the ones that are included  in other detections.
    - If no title is found, the biggest `other` box is used if it is larger than 0.3*img_w.
    Args:
        preds (np.ndarray): YOLOX model predictions including bounding boxes, scores, and labels.
        shape (tuple): Original image shape.
        threshold_ (float): Score threshold to filter bounding boxes.
    Returns:
        Dict[str, np.ndarray]: Dictionary of bounding boxes, organized by class.
    """
    bbox_dict = {label: np.array([]) for label in class_labels}

    for i, label in enumerate(class_labels):
        if label not in preds:
            continue

        bboxes_class = np.array(preds[label])

        if bboxes_class.size == 0:
            continue

        # Try to find a chart_title box
        threshold = threshold_ if label != "chart_title" else min(threshold_, bboxes_class[:, -1].max())
        bboxes_class = bboxes_class[bboxes_class[:, -1] >= threshold][:, :4].astype(int)

        sort = ["x0", "y0"] if label != "ylabel" else ["y0", "x0"]
        idxs = (
            pd.DataFrame(
                {
                    "y0": bboxes_class[:, 1],
                    "x0": bboxes_class[:, 0],
                }
            )
            .sort_values(sort, ascending=label != "ylabel")
            .index
        )
        bboxes_class = bboxes_class[idxs]
        bbox_dict[label] = bboxes_class

    # Remove other included
    if len(bbox_dict.get("other", [])):
        other = find_boxes_inside(
            np.concatenate(list([v for v in bbox_dict.values() if len(v)])), bbox_dict["other"], threshold=0.7
        )
        del bbox_dict["other"]
        if len(other):
            bbox_dict["other"] = other

    # Biggest other is title if no title
    if not len(bbox_dict.get("chart_title", [])) and len(bbox_dict.get("other", [])):
        boxes = bbox_dict["other"]
        ws = boxes[:, 2] - boxes[:, 0]
        if np.max(ws) > shape[1] * 0.3:
            bbox_dict["chart_title"] = boxes[np.argmax(ws)][None].copy()
            bbox_dict["other"] = np.delete(boxes, (np.argmax(ws)), axis=0)

    # Make sure other key not lost
    bbox_dict["other"] = bbox_dict.get("other", [])

    return bbox_dict


def get_bbox_dict_yolox_table(preds, shape, class_labels, threshold=0.1, delta=0.0):
    """
    Extracts bounding boxes from YOLOX model predictions:
    - Applies thresholding
    - Reformats boxes
    - Reorders predictions

    Args:
        preds (np.ndarray): YOLOX model predictions including bounding boxes, scores, and labels.
        shape (tuple): Original image shape.
        config: Model configuration, including size for bounding box adjustment.
        threshold (float): Score threshold to filter bounding boxes.
        delta (float): How much the table was cropped upwards.

    Returns:
        dict[str, np.ndarray]: Dictionary of bounding boxes, organized by class.
    """
    bbox_dict = {label: np.array([]) for label in class_labels}

    for i, label in enumerate(class_labels):
        if label not in ["cell", "row", "column"]:
            continue  # Ignore useless classes

        bboxes_class = np.array(preds[label])

        if bboxes_class.size == 0:
            continue

        # Threshold and clip
        bboxes_class = bboxes_class[bboxes_class[:, -1] >= threshold][:, :4].astype(int)
        bboxes_class[:, [0, 2]] = np.clip(bboxes_class[:, [0, 2]], 0, shape[1])
        bboxes_class[:, [1, 3]] = np.clip(bboxes_class[:, [1, 3]], 0, shape[0])

        # Reorder
        sort = ["x0", "y0"] if label != "row" else ["y0", "x0"]
        df = pd.DataFrame(
            {
                "y0": (bboxes_class[:, 1] + bboxes_class[:, 3]) / 2,
                "x0": (bboxes_class[:, 0] + bboxes_class[:, 2]) / 2,
            }
        )
        idxs = df.sort_values(sort).index
        bboxes_class = bboxes_class[idxs]

        bbox_dict[label] = bboxes_class

    # Enforce spanning the entire table
    if len(bbox_dict["row"]):
        bbox_dict["row"][:, 0] = 0
        bbox_dict["row"][:, 2] = shape[1]
    if len(bbox_dict["column"]):
        bbox_dict["column"][:, 1] = 0
        bbox_dict["column"][:, 3] = shape[0]

    # Shift back if cropped
    for k in bbox_dict:
        if len(bbox_dict[k]):
            bbox_dict[k][:, [1, 3]] = np.add(bbox_dict[k][:, [1, 3]], delta, casting="unsafe")

    return bbox_dict


def match_with_title_v3(bbox, title_bboxes, match_dist=0.1, delta=1.5, already_matched=[]):
    """
    Matches a bounding box with a title bounding box based on IoU or proximity.

    Args:
        bbox (numpy.ndarray): Bounding box to match with title [x_min, y_min, x_max, y_max].
        title_bboxes (numpy.ndarray): Array of title bounding boxes with shape (N, 4).
        match_dist (float, optional): Maximum distance for matching. Defaults to 0.1.
        delta (float, optional): Multiplier for matching several titles. Defaults to 1.5.
        already_matched (list, optional): List of already matched title indices. Defaults to [].

    Returns:
        tuple or None: If matched, returns a tuple of (merged_bbox, updated_title_bboxes).
                       If no match is found, returns None, None.
    """
    if not len(title_bboxes):
        return None, None

    dist_above = np.abs(title_bboxes[:, 3] - bbox[1])
    dist_below = np.abs(bbox[3] - title_bboxes[:, 1])

    dist_left = np.abs(title_bboxes[:, 0] - bbox[0])
    dist_center = np.abs(title_bboxes[:, 0] + title_bboxes[:, 2] - bbox[0] - bbox[2]) / 2

    dists = np.min([dist_above, dist_below], 0)
    dists += np.min([dist_left, dist_center], 0) / 2

    ious = bb_iou_array(title_bboxes, bbox)
    dists = np.where(ious > 0, min(match_dist, np.min(dists)), dists)

    if len(already_matched):
        dists[already_matched] = match_dist * 10  # Remove already matched titles

    # print(dists)
    matches = None  # noqa
    if np.min(dists) <= match_dist:
        matches = np.where(dists <= min(match_dist, np.min(dists) * delta))[0]

    if matches is not None:
        new_bbox = bbox
        for match in matches:
            new_bbox = merge_boxes(new_bbox, title_bboxes[match])
        return new_bbox, list(matches)
    else:
        return None, None


def match_boxes_with_title(boxes, confs, labels, classes, to_match_labels=["chart"], remove_matched_titles=False):
    """
    Matches charts with title.

    Args:
        boxes (numpy.ndarray): Array of bounding boxes with shape (N, 4).
        confs (numpy.ndarray): Array of confidence scores with shape (N,).
        labels (numpy.ndarray): Array of labels with shape (N,).
        classes (list): List of class names.
        to_match_labels (list): List of class names to match with titles.
        remove_matched_titles (bool): Whether to remove matched titles from the boxes.

    Returns:
        boxes (numpy.ndarray): Array of bounding boxes with shape (M, 4).
        confs (numpy.ndarray): Array of confidence scores with shape (M,).
        labels (numpy.ndarray): Array of labels with shape (M,).
        found_title (list): List of indices of matched titles.
        no_found_title (list): List of indices of unmatched titles.
    """
    # Put titles at the end
    title_ids = np.where(labels == classes.index("title"))[0]
    order = np.concatenate([np.delete(np.arange(len(boxes)), title_ids), title_ids])
    boxes = boxes[order]
    confs = confs[order]
    labels = labels[order]

    # Ids
    title_ids = np.where(labels == classes.index("title"))[0]
    to_match = np.where(np.isin(labels, [classes.index(c) for c in to_match_labels]))[0]

    # Matching
    found_title, already_matched = [], []
    for i in range(len(boxes)):
        if i not in to_match:
            continue
        merged_box, matched_title_ids = match_with_title_v3(
            boxes[i],
            boxes[title_ids],
            already_matched=already_matched,
        )
        if matched_title_ids is not None:
            # print(f'Merged {classes[int(labels[i])]} at idx #{i} with title {matched_title_ids[-1]}')  # noqa
            boxes[i] = merged_box
            already_matched += matched_title_ids
            found_title.append(i)

    if remove_matched_titles and len(already_matched):
        boxes = np.delete(boxes, title_ids[already_matched], axis=0)
        confs = np.delete(confs, title_ids[already_matched], axis=0)
        labels = np.delete(labels, title_ids[already_matched], axis=0)

    return boxes, confs, labels, found_title


def expand_boxes_v3(boxes, r_x=(1, 1), r_y=(1, 1), size_agnostic=True):
    """
    Expands bounding boxes by a specified ratio.
    Expected box format is normalized [x_min, y_min, x_max, y_max].

    Args:
        boxes (numpy.ndarray): Array of bounding boxes with shape (N, 4).
        r_x (tuple, optional): Left, right expansion ratios. Defaults to (1, 1) (no expansion).
        r_y (tuple, optional): Up, down expansion ratios. Defaults to (1, 1) (no expansion).
        size_agnostic (bool, optional): Expand independently of the bbox shape. Defaults to True.

    Returns:
        numpy.ndarray: Adjusted bounding boxes clipped to the [0, 1] range.
    """
    old_boxes = boxes.copy()

    if not size_agnostic:
        h = boxes[:, 3] - boxes[:, 1]
        w = boxes[:, 2] - boxes[:, 0]
    else:
        h, w = 1, 1

    boxes[:, 0] -= w * (r_x[0] - 1)  # left
    boxes[:, 2] += w * (r_x[1] - 1)  # right
    boxes[:, 1] -= h * (r_y[0] - 1)  # up
    boxes[:, 3] += h * (r_y[1] - 1)  # down

    boxes = np.clip(boxes, 0, 1)

    # Enforce non-overlapping boxes
    for i in range(len(boxes)):
        for j in range(i + 1, len(boxes)):
            iou = bb_iou_array(boxes[i][None], boxes[j])[0]
            old_iou = bb_iou_array(old_boxes[i][None], old_boxes[j])[0]
            # print(iou, old_iou)
            if iou > 0.05 and old_iou < 0.1:
                if boxes[i, 1] < boxes[j, 1]:  # i above j
                    boxes[j, 1] = min(old_boxes[j, 1], boxes[i, 3])
                    if old_iou > 0:
                        boxes[i, 3] = max(old_boxes[i, 3], boxes[j, 1])
                else:
                    boxes[i, 1] = min(old_boxes[i, 1], boxes[j, 3])
                    if old_iou > 0:
                        boxes[j, 3] = max(old_boxes[j, 3], boxes[i, 1])

    return boxes


def get_overlaps(boxes, other_boxes, normalize="box_only"):
    """
    Checks if a box overlaps with any other box.
    Boxes are expeceted in format (x0, y0, x1, y1)

    Args:
        boxes (np array [4] or [n x 4]): Boxes.
        other_boxes (np array [m x 4]): Other boxes.

    Returns:
        np array [n x m]: Overlaps.
    """
    if boxes.ndim == 1:
        boxes = boxes[None, :]

    x0, y0, x1, y1 = (boxes[:, 0][:, None], boxes[:, 1][:, None], boxes[:, 2][:, None], boxes[:, 3][:, None])
    areas = (y1 - y0) * (x1 - x0)

    x0_other, y0_other, x1_other, y1_other = (
        other_boxes[:, 0][None, :],
        other_boxes[:, 1][None, :],
        other_boxes[:, 2][None, :],
        other_boxes[:, 3][None, :],
    )
    areas_other = (y1_other - y0_other) * (x1_other - x0_other)

    # Intersection
    inter_y0 = np.maximum(y0, y0_other)
    inter_y1 = np.minimum(y1, y1_other)
    inter_x0 = np.maximum(x0, x0_other)
    inter_x1 = np.minimum(x1, x1_other)
    inter_area = np.maximum(0, inter_y1 - inter_y0) * np.maximum(0, inter_x1 - inter_x0)

    # Overlap
    if normalize == "box_only":  # Only consider box included in other box
        overlaps = inter_area / areas
    elif normalize == "all":  # Consider box included in other box and other box included in box
        overlaps = inter_area / np.minimum(areas, areas_other[:, None])
    else:
        raise ValueError(f"Invalid normalization: {normalize}")
    return overlaps


def postprocess_included(boxes, labels, confs, class_="title", classes=["table", "chart", "title", "infographic"]):
    """
    Post process title predictions.
    - Remove titles that are included in other boxes

    Args:
        boxes (numpy.ndarray [N, 4]): Array of bounding boxes.
        labels (numpy.ndarray [N]): Array of labels.
        confs (numpy.ndarray [N]): Array of confidences.
        class_ (str, optional): Class to postprocess. Defaults to "title".
        classes (list, optional): Classes. Defaults to ["table", "chart", "title", "infographic"].

    Returns:
        boxes (numpy.ndarray): Array of bounding boxes.
        labels (numpy.ndarray): Array of labels.
        confs (numpy.ndarray): Array of confidences.
    """
    boxes_to_pp = boxes[labels == classes.index(class_)]
    confs_to_pp = confs[labels == classes.index(class_)]

    order = np.argsort(confs_to_pp)  # least to most confident for NMS
    boxes_to_pp, confs_to_pp = boxes_to_pp[order], confs_to_pp[order]

    if len(boxes_to_pp) == 0:
        return boxes, labels, confs

    # other_boxes = boxes[labels != classes.index("title")]

    inclusion_classes = ["table", "infographic", "chart"]
    if class_ in ["header_footer", "title"]:
        inclusion_classes.append("paragraph")

    other_boxes = boxes[np.isin(labels, [classes.index(c) for c in inclusion_classes])]

    # Remove boxes included in other_boxes
    kept_boxes, kept_confs = [], []
    for i, b in enumerate(boxes_to_pp):
        if len(other_boxes) > 0:
            overlaps = get_overlaps(b, other_boxes, normalize="box_only")
            if overlaps.max() > 0.9:
                continue

        kept_boxes.append(b)
        kept_confs.append(confs_to_pp[i])

    # Aggregate
    kept_boxes = np.stack(kept_boxes) if len(kept_boxes) else np.empty((0, 4))
    kept_confs = np.stack(kept_confs) if len(kept_confs) else np.empty(0)

    boxes_pp = np.concatenate([boxes[labels != classes.index(class_)], kept_boxes])
    confs_pp = np.concatenate([confs[labels != classes.index(class_)], kept_confs])
    labels_pp = np.concatenate(
        [labels[labels != classes.index(class_)], np.ones(len(kept_boxes)) * classes.index(class_)]
    )

    return boxes_pp, labels_pp, confs_pp


def remove_overlapping_boxes_using_wbf(boxes, confs, labels):
    """
    Remove overlapping boxes using WBF
    """
    # Applied twice because once is not enough in some rare cases
    for _ in range(2):
        boxes, confs, labels = weighted_boxes_fusion(
            boxes[:, None],
            confs[:, None],
            labels[:, None],
            merge_type="biggest",
            conf_type="max",
            iou_thr=0.01,
            class_agnostic=False,
        )

    return boxes, confs, labels


def match_structured_boxes_with_title(boxes, confs, labels, classes):
    # Reorder by y, x
    order = np.argsort(boxes[:, 1] * 10 + boxes[:, 0])
    boxes, confs, labels = boxes[order], confs[order], labels[order]

    # Match with title
    # Although the model should detect titles, additional post-processing helps retrieve FNs
    found_title = []
    boxes, confs, labels, found_title = match_boxes_with_title(
        boxes,
        confs,
        labels,
        classes,
        to_match_labels=["chart", "table", "infographic"],
        remove_matched_titles=True,
    )

    return boxes, confs, labels, found_title


def expand_tables_and_charts(boxes, confs, labels, classes, found_title):
    # This is mostly to retrieve titles, but this also helps when YOLOX boxes are too tight.
    # Boxes with titles matched are expanded less.
    # Expansion is different for tables and charts
    no_found_title = [i for i in range(len(boxes)) if i not in found_title]
    ids = np.arange(len(boxes))

    if len(found_title):  # Boxes with title matched are expanded less
        ids_ = ids[found_title][labels[found_title] == classes.index("chart")]
        boxes[ids_] = expand_boxes_v3(
            boxes[ids_],
            r_x=(1.025, 1.025),
            r_y=(1.05, 1.05),
            size_agnostic=False,
        )
        ids_ = ids[found_title][labels[found_title] == classes.index("table")]
        boxes[ids_] = expand_boxes_v3(
            boxes[ids_],
            r_x=(1.01, 1.01),
            r_y=(1.05, 1.01),
        )

    ids_ = ids[no_found_title][labels[no_found_title] == classes.index("chart")]
    boxes[ids_] = expand_boxes_v3(
        boxes[ids_],
        r_x=(1.05, 1.05),
        r_y=(1.125, 1.125),
        size_agnostic=False,
    )

    ids_ = ids[no_found_title][labels[no_found_title] == classes.index("table")]
    boxes[ids_] = expand_boxes_v3(
        boxes[ids_],
        r_x=(1.02, 1.02),
        r_y=(1.05, 1.05),
    )

    order = np.argsort(boxes[:, 1] * 10 + boxes[:, 0])
    boxes, labels, confs = boxes[order], labels[order], confs[order]

    return boxes, labels, confs


def postprocess_included_texts(boxes, confs, labels, classes):
    for c in ["title", "paragraph", "header_footer"]:
        boxes, labels, confs = postprocess_included(boxes, labels, confs, c, classes)
    return boxes, labels, confs


@backoff.on_predicate(backoff.expo, max_time=30)
def get_yolox_model_name(yolox_grpc_endpoint, default_model_name="yolox"):
    # If a gRPC endpoint isn't provided (common when using HTTP-only NIM endpoints),
    # skip Triton repository introspection entirely.
    if not yolox_grpc_endpoint:
        return default_model_name

    # Guard against accidentally passing an HTTP URL into the gRPC Triton client.
    if isinstance(yolox_grpc_endpoint, str) and (
        yolox_grpc_endpoint.startswith("http://") or yolox_grpc_endpoint.startswith("https://")
    ):
        return default_model_name

    key = (
        "get_yolox_model_name",
        (yolox_grpc_endpoint,),
        frozenset({"default_model_name": default_model_name}.items()),
    )
    with lock:
        if key in global_cache:
            return global_cache[key]

    try:
        client = grpcclient.InferenceServerClient(yolox_grpc_endpoint)
        model_index = client.get_model_repository_index(as_json=True)
        model_names = [x.get("name") for x in model_index.get("models", []) if isinstance(x, dict)]

        # Prefer explicit known orchestrations first.
        for preferred in (
            "pipeline",  # BLS pipeline
            "yolox_ensemble",
            "yolox",
            "yolox-page-elements",
            "page-elements",
            "nemotron-page-elements-v3",
            "nemoretriever-page-elements-v3",
            "nemoretriever-page-elements-v2",
        ):
            if preferred in model_names:
                result = preferred
                with lock:
                    global_cache[key] = result
                return result

        # Otherwise pick a best-effort match for newer model names.
        candidates = [m for m in model_names if isinstance(m, str) and ("yolox" in m or "page-elements" in m)]
        if candidates:
            result = sorted(candidates)[0]
            with lock:
                global_cache[key] = result
            return result

        result = default_model_name
        with lock:
            global_cache[key] = result
        return result
    except Exception as e:
        logger.warning(
            "Failed to inspect YOLOX model repository at '%s' (%s). Falling back to '%s'.",
            yolox_grpc_endpoint,
            type(e).__name__,
            default_model_name,
        )
        return default_model_name


@multiprocessing_cache(max_calls=100)  # Cache results first to avoid redundant retries from backoff
@backoff.on_predicate(backoff.expo, max_time=30)
def get_yolox_page_version(yolox_http_endpoint, default_version=YOLOX_PAGE_DEFAULT_VERSION):
    """
    Determines the YOLOX page elements model version by querying the endpoint.
    Falls back to a default version on failure.
    """
    try:
        yolox_version = get_model_name(yolox_http_endpoint, default_version)
        if not yolox_version:
            logger.warning(
                "Failed to obtain yolox-page-elements version from the endpoint. "
                f"Falling back to '{default_version}'."
            )
            return default_version

        return yolox_version
    except Exception:
        logger.warning(
            f"Failed to get yolox-page-elements version after 30 seconds. Falling back to '{default_version}'."
        )
        return default_version
