from __future__ import annotations

import gc
import os
import time
from typing import List, Tuple
from unittest.mock import MagicMock, Mock

import pytest
from inference.core.entities.responses.inference import (
    InferenceResponseImage, MultiLabelClassificationInferenceResponse)
from inference.core.models.inference_models_adapters import \
    prepare_multi_label_classification_response
from inference_models import MultiLabelClassificationPrediction


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f'!######{test_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

def test_single_prediction_single_class_above_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test basic functionality with a single prediction, single class, and confidence above threshold.\n    This is the simplest case to verify core functionality works.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['cat']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_single_prediction_single_class_above_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_prediction_single_class_below_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that a class with confidence below threshold is not added to predicted_classes.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.3])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['dog']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_single_prediction_single_class_below_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_predictions_multiple_classes():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with multiple predictions and multiple classes per prediction.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9, 0.4, 0.7])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['cat', 'dog', 'bird']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_multiple_predictions_multiple_classes', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_multiple_predictions_multiple_images():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with multiple predictions for different images.\n    '
    mock_prediction1 = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction1.class_ids = Mock()
    mock_prediction1.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction1.confidence = Mock()
    mock_prediction1.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.8])))
    mock_prediction2 = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction2.class_ids = Mock()
    mock_prediction2.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[1])))
    mock_prediction2.confidence = Mock()
    mock_prediction2.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.6])))
    predictions = [mock_prediction1, mock_prediction2]
    image_sizes = [(480, 640), (720, 1280)]
    class_names = ['cat', 'dog']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_multiple_predictions_multiple_images', 'prepare_multi_label_classification_response', '15', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_image_dimensions_swapped_correctly():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that image dimensions are correctly assigned (height from first, width from second).\n    The function swaps them: width=image_size[1], height=image_size[0]\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
    predictions = [mock_prediction]
    image_sizes = [(1080, 1920)]
    class_names = ['object']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_image_dimensions_swapped_correctly', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_confidence_exactly_at_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test boundary condition where confidence equals the threshold.\n    The function uses > (not >=), so confidence equal to threshold should NOT be included.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.5])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['test']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_confidence_exactly_at_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_zero_confidence_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with threshold of 0.0, all positive confidences should be included.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.0, 0.001, 0.5])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['a', 'b', 'c']
    confidence_threshold = 0.0
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_zero_confidence_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_high_confidence_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with very high threshold (0.99), few classes should be included.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2, 3])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.99, 0.995, 0.98, 0.999])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['a', 'b', 'c', 'd']
    confidence_threshold = 0.99
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_high_confidence_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_all_classes_below_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test case where all classes have confidence below threshold.\n    predicted_classes should be empty but predictions dict should contain all classes.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.1, 0.2, 0.3])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['x', 'y', 'z']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_all_classes_below_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_all_classes_above_threshold():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test case where all classes have confidence above threshold.\n    All classes should be in both predicted_classes and predictions dict.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9, 0.8, 0.7])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['alpha', 'beta', 'gamma']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_all_classes_above_threshold', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_predictions_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with empty predictions list.\n    Should return empty results list.\n    '
    predictions = []
    image_sizes = []
    class_names = ['cat', 'dog']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_empty_predictions_list', 'prepare_multi_label_classification_response', '5', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_class_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with only one class available in class_names.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
    predictions = [mock_prediction]
    image_sizes = [(100, 200)]
    class_names = ['only_class']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_single_class_name', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_many_class_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with a large number of class names.\n    '
    class_names = [f'class_{i}' for i in range(100)]
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    class_ids = list(range(100))
    confidences = [0.9 if i % 2 == 0 else 0.3 for i in range(100)]
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=class_ids)))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=confidences)))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_many_class_names', 'prepare_multi_label_classification_response', '12', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_small_image_dimensions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with very small image dimensions (edge case).\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
    predictions = [mock_prediction]
    image_sizes = [(1, 1)]
    class_names = ['tiny']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_very_small_image_dimensions', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_very_large_image_dimensions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with very large image dimensions.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
    predictions = [mock_prediction]
    image_sizes = [(8192, 16384)]
    class_names = ['large']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_very_large_image_dimensions', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_class_id_zero_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that class_id 0 is handled correctly (edge case with index).\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.95])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['first_class', 'second_class']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_class_id_zero_handling', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_high_class_id_handling():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that high class IDs are handled correctly.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[99])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.95])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = [f'class_{i}' for i in range(100)]
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_high_class_id_handling', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_floating_point_precision():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with floating point values that are very close to threshold.\n    '
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0, 1, 2])))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.5000001, 0.4999999, 0.5])))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    class_names = ['a', 'b', 'c']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_floating_point_precision', 'prepare_multi_label_classification_response', '10', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_number_of_predictions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test scalability with a large number of predictions (500 images).\n    '
    predictions = []
    image_sizes = []
    for i in range(500):
        mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
        mock_prediction.class_ids = Mock()
        mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[i % 10])))
        mock_prediction.confidence = Mock()
        mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.5 + i % 50 * 0.01])))
        predictions.append(mock_prediction)
        image_sizes.append((480 + i % 100, 640 + i % 100))
    class_names = [f'class_{i}' for i in range(10)]
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_large_number_of_predictions', 'prepare_multi_label_classification_response', '6', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_number_of_classes_per_prediction():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with many classes per prediction (500 classes total).\n    '
    num_classes = 500
    class_names = [f'class_{i}' for i in range(num_classes)]
    mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
    class_ids = list(range(num_classes))
    confidences = [0.5 + i % 100 * 0.005 for i in range(num_classes)]
    mock_prediction.class_ids = Mock()
    mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=class_ids)))
    mock_prediction.confidence = Mock()
    mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=confidences)))
    predictions = [mock_prediction]
    image_sizes = [(480, 640)]
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_large_number_of_classes_per_prediction', 'prepare_multi_label_classification_response', '13', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_predicted = sum((1 for conf in confidences if conf > confidence_threshold))

def test_mixed_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test large scale with both many predictions and many classes.\n    '
    num_predictions = 100
    num_classes = 100
    class_names = [f'class_{i}' for i in range(num_classes)]
    predictions = []
    image_sizes = []
    for pred_idx in range(num_predictions):
        mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
        class_ids = list(range(num_classes))
        confidences = [0.5 + (pred_idx + class_idx) % 50 * 0.01 for class_idx in range(num_classes)]
        mock_prediction.class_ids = Mock()
        mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=class_ids)))
        mock_prediction.confidence = Mock()
        mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=confidences)))
        predictions.append(mock_prediction)
        image_sizes.append((480, 640))
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_mixed_large_scale', 'prepare_multi_label_classification_response', '8', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_predictions_with_varied_confidence_distributions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with a large number of predictions having varied confidence distributions.\n    '
    num_predictions = 200
    class_names = ['cat', 'dog', 'bird', 'fish', 'hamster']
    predictions = []
    image_sizes = []
    for pred_idx in range(num_predictions):
        mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
        class_ids = [0, 1, 2, 3, 4]
        if pred_idx % 3 == 0:
            confidences = [0.9, 0.1, 0.8, 0.2, 0.95]
        elif pred_idx % 3 == 1:
            confidences = [0.4, 0.4, 0.4, 0.4, 0.4]
        else:
            confidences = [0.95, 0.95, 0.95, 0.95, 0.95]
        mock_prediction.class_ids = Mock()
        mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=class_ids)))
        mock_prediction.confidence = Mock()
        mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=confidences)))
        predictions.append(mock_prediction)
        image_sizes.append((480, 640))
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_large_scale_predictions_with_varied_confidence_distributions', 'prepare_multi_label_classification_response', '7', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_preservation_of_class_information():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that in large scale scenarios, class_id and confidence are preserved correctly.\n    '
    num_predictions = 150
    class_names = [f'class_{i}' for i in range(50)]
    predictions = []
    image_sizes = []
    for pred_idx in range(num_predictions):
        mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
        class_ids = list(range(50))
        confidences = [0.5 + (pred_idx * 0.001 + i * 0.001) % 0.5 for i in range(50)]
        mock_prediction.class_ids = Mock()
        mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=class_ids)))
        mock_prediction.confidence = Mock()
        mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=confidences)))
        predictions.append(mock_prediction)
        image_sizes.append((480, 640))
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_large_scale_preservation_of_class_information', 'prepare_multi_label_classification_response', '7', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for pred_idx, response in enumerate(result):
        for class_name, class_info in response.predictions.items():
            expected_class_id = int(class_name.split('_')[1])

def test_large_scale_image_dimensions_variety():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with large number of predictions having varied image dimensions.\n    '
    num_predictions = 250
    predictions = []
    image_sizes = []
    for i in range(num_predictions):
        mock_prediction = Mock(spec=MultiLabelClassificationPrediction)
        mock_prediction.class_ids = Mock()
        mock_prediction.class_ids.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0])))
        mock_prediction.confidence = Mock()
        mock_prediction.confidence.cpu = Mock(return_value=Mock(tolist=Mock(return_value=[0.9])))
        predictions.append(mock_prediction)
        height = 360 + i % 100 * 10
        width = 480 + i % 100 * 20
        image_sizes.append((height, width))
    class_names = ['object']
    confidence_threshold = 0.5
    codeflash_return_value = codeflash_wrap(prepare_multi_label_classification_response, 'tests.inference.unit_tests.test_prepare_multi_label_classification_response__unit_test_1', None, 'test_large_scale_image_dimensions_variety', 'prepare_multi_label_classification_response', '7', codeflash_loop_index, predictions, image_sizes, class_names, confidence_threshold)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for idx, response in enumerate(result):
        expected_height, expected_width = image_sizes[idx]