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


import base64
import io
import logging
import warnings
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 cv2
import numpy as np
import packaging
import pandas as pd
import torch
import torchvision
from PIL import Image

from nv_ingest_api.internal.primitives.nim import ModelInterface
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

logger = logging.getLogger(__name__)

# yolox-page-elements-v1 and v2 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-elements-v1 contants
YOLOX_PAGE_V1_NUM_CLASSES = 4
YOLOX_PAGE_V1_FINAL_SCORE = {"table": 0.48, "chart": 0.48}
YOLOX_PAGE_V1_CLASS_LABELS = [
    "table",
    "chart",
    "title",
]

# yolox-page-elements-v2 contants
YOLOX_PAGE_V2_NUM_CLASSES = 4
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_NUM_CLASSES = 10
YOLOX_GRAPHIC_CONF_THRESHOLD = 0.01
YOLOX_GRAPHIC_IOU_THRESHOLD = 0.25
YOLOX_GRAPHIC_MIN_SCORE = 0.1
YOLOX_GRAPHIC_FINAL_SCORE = 0.0
YOLOX_GRAPHIC_NIM_MAX_IMAGE_SIZE = 512_000

# TODO(Devin): Legacy items aren't working right for me. Double check these.
LEGACY_YOLOX_GRAPHIC_IMAGE_PREPROC_HEIGHT = 1024
LEGACY_YOLOX_GRAPHIC_IMAGE_PREPROC_WIDTH = 1024
YOLOX_GRAPHIC_IMAGE_PREPROC_HEIGHT = 1024
YOLOX_GRAPHIC_IMAGE_PREPROC_WIDTH = 1024

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_NUM_CLASSES = 5
YOLOX_TABLE_CONF_THRESHOLD = 0.01
YOLOX_TABLE_IOU_THRESHOLD = 0.25
YOLOX_TABLE_MIN_SCORE = 0.1
YOLOX_TABLE_FINAL_SCORE = 0.0
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,
        image_preproc_width: Optional[int] = None,
        image_preproc_height: Optional[int] = None,
        nim_max_image_size: Optional[int] = None,
        num_classes: Optional[int] = None,
        conf_threshold: Optional[float] = None,
        iou_threshold: Optional[float] = None,
        min_score: Optional[float] = None,
        final_score: Optional[float] = None,
        class_labels: Optional[List[str]] = None,
    ):
        """
        Initialize the YOLOX model interface.
        Parameters
        ----------
        """
        self.image_preproc_width = image_preproc_width
        self.image_preproc_height = image_preproc_height
        self.nim_max_image_size = nim_max_image_size
        self.num_classes = num_classes
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        self.min_score = min_score
        self.final_score = final_score
        self.class_labels = class_labels

    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]:
            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 = 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 model")
            # Resize images for model input (Yolox expects 1024x1024).
            resized_images = [
                resize_image(image, (self.image_preproc_width, self.image_preproc_height)) for image in data["images"]
            ]
            # Chunk the resized images, the original images, and their shapes.
            resized_chunks = chunk_list_geometrically(resized_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 r_chunk, orig_chunk, shapes in zip(resized_chunks, original_chunks, shape_chunks):
                # Reorder axes from (B, H, W, C) to (B, C, H, W) as expected by the model.
                input_array = np.einsum("bijk->bkij", r_chunk).astype(np.float32)
                batched_inputs.append(input_array)
                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)
                # Convert the numpy array to a PIL Image.
                image_pil = Image.fromarray(image)
                original_size = image_pil.size

                # Save the image to a buffer and encode to base64.
                buffered = io.BytesIO()
                image_pil.save(buffered, format="PNG")
                image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")

                # 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.
        """
        original_image_shapes = kwargs.get("original_image_shapes", [])

        if protocol == "http":
            # For http, the output already has postprocessing applied. Skip to table/chart expansion.
            results = output

        elif protocol == "grpc":
            # For grpc, apply the same NIM postprocessing.
            pred = postprocess_model_prediction(
                output,
                self.num_classes,
                self.conf_threshold,
                self.iou_threshold,
                class_agnostic=False,
            )
            results = postprocess_results(
                pred,
                original_image_shapes,
                self.image_preproc_width,
                self.image_preproc_height,
                self.class_labels,
                min_score=self.min_score,
            )

        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.
    """

    def __init__(self, yolox_model_name: str = "nemoretriever-page-elements-v2"):
        """
        Initialize the yolox-page-elements model interface.
        """
        if yolox_model_name.endswith("-v1"):
            num_classes = YOLOX_PAGE_V1_NUM_CLASSES
            final_score = YOLOX_PAGE_V1_FINAL_SCORE
            class_labels = YOLOX_PAGE_V1_CLASS_LABELS
        else:
            num_classes = YOLOX_PAGE_V2_NUM_CLASSES
            final_score = YOLOX_PAGE_V2_FINAL_SCORE
            class_labels = YOLOX_PAGE_V2_CLASS_LABELS

        super().__init__(
            image_preproc_width=YOLOX_PAGE_IMAGE_PREPROC_WIDTH,
            image_preproc_height=YOLOX_PAGE_IMAGE_PREPROC_HEIGHT,
            nim_max_image_size=YOLOX_PAGE_NIM_MAX_IMAGE_SIZE,
            num_classes=num_classes,
            conf_threshold=YOLOX_PAGE_CONF_THRESHOLD,
            iou_threshold=YOLOX_PAGE_IOU_THRESHOLD,
            min_score=YOLOX_PAGE_MIN_SCORE,
            final_score=final_score,
            class_labels=class_labels,
        )

    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, **kwargs):
        original_image_shapes = kwargs.get("original_image_shapes", [])

        expected_final_score_keys = [x for x in self.class_labels if x != "title"]
        if (not isinstance(self.final_score, dict)) or (
            sorted(self.final_score.keys()) != sorted(expected_final_score_keys)
        ):
            raise ValueError(
                "yolox-page-elements-v2 requires a dictionary of thresholds per each class: "
                f"{expected_final_score_keys}"
            )

        # 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 "table" in annotation_dict:
                new_dict["table"] = [bb for bb in annotation_dict["table"] if bb[4] >= self.final_score["table"]]
            if "chart" in annotation_dict:
                new_dict["chart"] = [bb for bb in annotation_dict["chart"] if bb[4] >= self.final_score["chart"]]
            if "infographic" in annotation_dict:
                new_dict["infographic"] = [
                    bb for bb in annotation_dict["infographic"] if bb[4] >= self.final_score["infographic"]
                ]
            if "title" in annotation_dict:
                new_dict["title"] = annotation_dict["title"]
            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.
    """

    def __init__(self, yolox_version: Optional[str] = None):
        """
        Initialize the yolox-graphic-elements model interface.
        """
        if yolox_version and (
            packaging.version.Version(yolox_version) >= packaging.version.Version("1.2.0-rc5")  # gtc release
        ):
            image_preproc_width = YOLOX_GRAPHIC_IMAGE_PREPROC_WIDTH
            image_preproc_height = YOLOX_GRAPHIC_IMAGE_PREPROC_HEIGHT
        else:
            image_preproc_width = LEGACY_YOLOX_GRAPHIC_IMAGE_PREPROC_WIDTH
            image_preproc_height = LEGACY_YOLOX_GRAPHIC_IMAGE_PREPROC_HEIGHT

        super().__init__(
            image_preproc_width=image_preproc_width,
            image_preproc_height=image_preproc_height,
            nim_max_image_size=YOLOX_GRAPHIC_NIM_MAX_IMAGE_SIZE,
            num_classes=YOLOX_GRAPHIC_NUM_CLASSES,
            conf_threshold=YOLOX_GRAPHIC_CONF_THRESHOLD,
            iou_threshold=YOLOX_GRAPHIC_IOU_THRESHOLD,
            min_score=YOLOX_GRAPHIC_MIN_SCORE,
            final_score=YOLOX_GRAPHIC_FINAL_SCORE,
            class_labels=YOLOX_GRAPHIC_CLASS_LABELS,
        )

    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):
        """
        Initialize the yolox-graphic-elements model interface.
        """
        super().__init__(
            image_preproc_width=YOLOX_TABLE_IMAGE_PREPROC_HEIGHT,
            image_preproc_height=YOLOX_TABLE_IMAGE_PREPROC_HEIGHT,
            nim_max_image_size=YOLOX_TABLE_NIM_MAX_IMAGE_SIZE,
            num_classes=YOLOX_TABLE_NUM_CLASSES,
            conf_threshold=YOLOX_TABLE_CONF_THRESHOLD,
            iou_threshold=YOLOX_TABLE_IOU_THRESHOLD,
            min_score=YOLOX_TABLE_MIN_SCORE,
            final_score=YOLOX_TABLE_FINAL_SCORE,
            class_labels=YOLOX_TABLE_CLASS_LABELS,
        )

    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 postprocess_model_prediction(prediction, num_classes, conf_thre=0.7, nms_thre=0.45, class_agnostic=False):
    # Convert numpy array to torch tensor
    prediction = torch.from_numpy(prediction.copy())

    # Compute box corners
    box_corner = prediction.new(prediction.shape)
    box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
    box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
    prediction[:, :, :4] = box_corner[:, :, :4]

    output = [None for _ in range(len(prediction))]

    for i, image_pred in enumerate(prediction):
        # If no detections, continue to the next image
        if not image_pred.size(0):
            continue

        # Ensure image_pred is 2D
        if image_pred.ndim == 1:
            image_pred = image_pred.unsqueeze(0)

        # Get score and class with highest confidence
        class_conf, class_pred = torch.max(image_pred[:, 5 : 5 + num_classes], 1, keepdim=True)

        # Confidence mask
        squeezed_conf = class_conf.squeeze(dim=1)
        conf_mask = image_pred[:, 4] * squeezed_conf >= conf_thre

        # Apply confidence mask
        detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1)
        detections = detections[conf_mask]

        if not detections.size(0):
            continue

        # Apply Non-Maximum Suppression (NMS)
        if class_agnostic:
            nms_out_index = torchvision.ops.nms(
                detections[:, :4],
                detections[:, 4] * detections[:, 5],
                nms_thre,
            )
        else:
            nms_out_index = torchvision.ops.batched_nms(
                detections[:, :4],
                detections[:, 4] * detections[:, 5],
                detections[:, 6],
                nms_thre,
            )
        detections = detections[nms_out_index]

        # Append detections to output
        output[i] = detections

    return output


def postprocess_results(
    results, original_image_shapes, image_preproc_width, image_preproc_height, class_labels, min_score=0.0
):
    """
    For each item (==image) in results, computes annotations in the form

     {"table": [[0.0107, 0.0859, 0.7537, 0.1219, 0.9861], ...],
      "figure": [...],
      "title": [...]
      }
    where each list of 5 floats represents a bounding box in the format [x1, y1, x2, y2, confidence]

    Keep only bboxes with high enough confidence.
    """
    out = []

    for original_image_shape, result in zip(original_image_shapes, results):
        annotation_dict = {label: [] for label in class_labels}

        if result is None:
            out.append(annotation_dict)
            continue

        try:
            result = result.cpu().numpy()
            scores = result[:, 4] * result[:, 5]
            result = result[scores > min_score]

            # ratio is used when image was padded
            ratio = min(
                image_preproc_width / original_image_shape[0],
                image_preproc_height / original_image_shape[1],
            )
            bboxes = result[:, :4] / ratio

            bboxes[:, [0, 2]] /= original_image_shape[1]
            bboxes[:, [1, 3]] /= original_image_shape[0]
            bboxes = np.clip(bboxes, 0.0, 1.0)

            labels = result[:, 6]
            scores = scores[scores > min_score]
        except Exception as e:
            raise ValueError(f"Error in postprocessing {result.shape} and {original_image_shape}: {e}")

        for box, score, label in zip(bboxes, scores, labels):
            # TODO(Devin): Sometimes we get back unexpected class labels?
            if (label < 0) or (label >= len(class_labels)):
                logger.warning(f"Invalid class label {label} found in postprocessing")
                continue
            else:
                class_name = class_labels[int(label)]

            annotation_dict[class_name].append([round(float(x), 4) for x in np.concatenate((box, [score]))])

        out.append(annotation_dict)

    return out


def resize_image(image, target_img_size):
    w, h, _ = np.array(image).shape

    if target_img_size is not None:  # Resize + Pad
        r = min(target_img_size[0] / w, target_img_size[1] / h)
        image = cv2.resize(
            image,
            (int(h * r), int(w * r)),
            interpolation=cv2.INTER_LINEAR,
        ).astype(np.uint8)
        image = np.pad(
            image,
            ((0, target_img_size[0] - image.shape[0]), (0, target_img_size[1] - image.shape[1]), (0, 0)),
            mode="constant",
            constant_values=114,
        )

    return image


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(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(chart_bboxes[found_title_idxs], r_x=1.05, r_y=1.1)
    chart_bboxes[no_found_title_idxs] = expand_boxes(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 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(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(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):
        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 get_yolox_model_name(yolox_http_endpoint, default_model_name="nemoretriever-page-elements-v2"):
    try:
        yolox_model_name = get_model_name(yolox_http_endpoint, default_model_name)
        if not yolox_model_name:
            logger.warning(
                "Failed to obtain yolox-page-elements model name from the endpoint. "
                f"Falling back to '{default_model_name}'."
            )
            yolox_model_name = default_model_name
    except Exception:
        logger.warning(
            "Failed to get yolox-page-elements version after 30 seconds. " f"Falling back to '{default_model_name}'."
        )
        yolox_model_name = default_model_name

    return yolox_model_name
