from __future__ import annotations

import gc
import os
import time
from typing import List, Tuple

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

class _SimpleTensor:

    def __init__(self, data):
        self._data = list(data)

    def cpu(self):
        return self

    def tolist(self):
        return list(self._data)

class _PredictionLike:
    """
    Minimal object that mimics the parts of MultiLabelClassificationPrediction
    used by prepare_multi_label_classification_response: attributes 'class_ids' and 'confidence'
    which support .cpu().tolist().
    """

    def __init__(self, class_ids, confidence):
        self.class_ids = _SimpleTensor(class_ids)
        self.confidence = _SimpleTensor(confidence)

def _invoke(preds, image_sizes, class_names, threshold):
    """
    Convenience to call the function under test. Kept to reduce repetition in tests.
    """
    return prepare_multi_label_classification_response(preds, image_sizes, class_names, threshold)

def test_basic_single_prediction():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[0, 1], confidence=[0.9, 0.2])]
    image_sizes = [(100, 200)]
    class_names = ['cat', 'dog']
    threshold = 0.5
    results = _invoke(preds, image_sizes, class_names, threshold)
    res = results[0]

def test_confidence_threshold_strict_greater():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[0], confidence=[0.5])]
    image_sizes = [(10, 20)]
    class_names = ['alpha']
    threshold = 0.5
    results = _invoke(preds, image_sizes, class_names, threshold)
    res = results[0]

def test_empty_inputs_return_empty_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    results = _invoke([], [], [], 0.1)

def test_more_image_sizes_than_predictions_ignored_excess_image_sizes():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[0], confidence=[1.0])]
    image_sizes = [(1, 1), (2, 2), (3, 3)]
    class_names = ['only']
    threshold = 0.0
    results = _invoke(preds, image_sizes, class_names, threshold)
    res = results[0]

def test_class_id_out_of_range_raises_index_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[10], confidence=[0.9])]
    image_sizes = [(5, 5)]
    class_names = ['a', 'b']
    threshold = 0.0
    with pytest.raises(IndexError):
        _invoke(preds, image_sizes, class_names, threshold)

def test_negative_class_id_uses_python_negative_indexing():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[-1, -2], confidence=[0.8, 0.4])]
    image_sizes = [(8, 16)]
    class_names = ['first', 'second', 'third']
    threshold = 0.5
    results = _invoke(preds, image_sizes, class_names, threshold)
    res = results[0]

def test_mismatched_class_ids_and_confidences_zip_behavior():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    preds = [_PredictionLike(class_ids=[0, 1, 2], confidence=[0.9, 0.1])]
    image_sizes = [(12, 24)]
    class_names = ['a', 'b', 'c']
    threshold = 0.2
    results = _invoke(preds, image_sizes, class_names, threshold)
    res = results[0]

def test_large_scale_multiple_predictions_efficiency_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    num_predictions = 100
    classes_per_prediction = 5
    class_names = [f'class_{i}' for i in range(classes_per_prediction)]
    threshold = 0.5
    preds = []
    image_sizes = []
    for i in range(num_predictions):
        confidences = [(i + j) % 10 / 10.0 for j in range(classes_per_prediction)]
        class_ids = list(range(classes_per_prediction))
        preds.append(_PredictionLike(class_ids=class_ids, confidence=confidences))
        image_sizes.append((i + 1, (i + 1) * 2))
    results = _invoke(preds, image_sizes, class_names, threshold)
    for idx in (0, num_predictions // 2, num_predictions - 1):
        res = results[idx]
        expected_height, expected_width = image_sizes[idx]
        for cname in class_names:
            pass
        for cname in res.predicted_classes:
            pass