from __future__ import annotations

import gc
import os
import time

import numpy as np
import pytest
import supervision as sv
from inference.core.workflows.core_steps.common.utils import \
    add_inference_keypoints_to_sv_detections
from inference.core.workflows.execution_engine.constants import (
    KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE,
    KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE,
    KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE,
    KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_KEY_IN_INFERENCE_RESPONSE, KEYPOINTS_XY_KEY_IN_SV_DETECTIONS,
    X_KEY, Y_KEY)


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_detection_two_keypoints_basic():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1
    detections = sv.Detections(np.zeros((n, 4), dtype=float))
    kp1 = {KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: 'eye', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: 1, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: 0.9, X_KEY: 10.0, Y_KEY: 20.0}
    kp2 = {KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: 'nose', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: 2, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: 0.8, X_KEY: 15.0, Y_KEY: 25.0}
    inference_prediction = [{KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: [kp1, kp2]}]
    codeflash_return_value = codeflash_wrap(add_inference_keypoints_to_sv_detections, 'tests.inference.unit_tests.test_add_inference_keypoints_to_sv_detections__unit_test_0', None, 'test_single_detection_two_keypoints_basic', 'add_inference_keypoints_to_sv_detections', '5', codeflash_loop_index, inference_prediction, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    padded_xy = detections[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]
    padded_conf = detections[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS]
    padded_cls = detections[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS]
    padded_names = detections[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS]

def test_mismatched_lengths_raises_value_error():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    detections = sv.Detections(np.zeros((1, 4), dtype=float))
    inference_prediction = [{KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: []}, {KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: []}]
    with pytest.raises(ValueError):
        codeflash_return_value = codeflash_wrap(add_inference_keypoints_to_sv_detections, 'tests.inference.unit_tests.test_add_inference_keypoints_to_sv_detections__unit_test_0', None, 'test_mismatched_lengths_raises_value_error', 'add_inference_keypoints_to_sv_detections', '2_0', codeflash_loop_index, inference_prediction, detections)

def test_all_detections_with_no_keypoints_results_in_zero_sized_kp_axes():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 3
    detections = sv.Detections(np.zeros((n, 4), dtype=float))
    inference_prediction = [{}, {}, {}]
    codeflash_return_value = codeflash_wrap(add_inference_keypoints_to_sv_detections, 'tests.inference.unit_tests.test_add_inference_keypoints_to_sv_detections__unit_test_0', None, 'test_all_detections_with_no_keypoints_results_in_zero_sized_kp_axes', 'add_inference_keypoints_to_sv_detections', '3', codeflash_loop_index, inference_prediction, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    padded_xy = detections[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]
    padded_conf = detections[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS]
    padded_cls = detections[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS]
    padded_names = detections[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS]

def test_variable_keypoint_counts_are_padded_correctly_and_order_preserved():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 3
    detections = sv.Detections(np.zeros((n, 4), dtype=float))
    infer = [{KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: [{KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: 'A', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: 10, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: 0.5, X_KEY: 1.0, Y_KEY: 2.0}]}, {KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: []}, {KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: [{KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: 'B', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: 20, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: 0.6, X_KEY: 3.0, Y_KEY: 4.0}, {KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: 'C', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: 30, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: 0.7, X_KEY: 5.0, Y_KEY: 6.0}]}]
    codeflash_return_value = codeflash_wrap(add_inference_keypoints_to_sv_detections, 'tests.inference.unit_tests.test_add_inference_keypoints_to_sv_detections__unit_test_0', None, 'test_variable_keypoint_counts_are_padded_correctly_and_order_preserved', 'add_inference_keypoints_to_sv_detections', '3', codeflash_loop_index, infer, detections)
    padded_xy = detections[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]
    padded_conf = detections[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS]
    padded_cls = detections[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS]
    padded_names = detections[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS]

def test_large_scale_padding_and_values_integrity():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    detections = sv.Detections(np.zeros((n, 4), dtype=float))
    inference_prediction = []
    for i in range(n):
        k_count = i % 5
        kps = []
        for k in range(k_count):
            kps.append({KEYPOINTS_CLASS_NAME_KEY_IN_INFERENCE_RESPONSE: f'name_{i}_{k}', KEYPOINTS_CLASS_ID_KEY_IN_INFERENCE_RESPONSE: i * 10 + k, KEYPOINTS_CONFIDENCE_KEY_IN_INFERENCE_RESPONSE: float(i % 100) / 100.0, X_KEY: float(i + k + 0.1), Y_KEY: float(i + k + 0.2)})
        inference_prediction.append({KEYPOINTS_KEY_IN_INFERENCE_RESPONSE: kps})
    codeflash_return_value = codeflash_wrap(add_inference_keypoints_to_sv_detections, 'tests.inference.unit_tests.test_add_inference_keypoints_to_sv_detections__unit_test_0', None, 'test_large_scale_padding_and_values_integrity', 'add_inference_keypoints_to_sv_detections', '4', codeflash_loop_index, inference_prediction, detections)
    padded_xy = detections[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]
    padded_conf = detections[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS]
    padded_cls = detections[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS]
    padded_names = detections[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS]
    sample_indices = [0, 1, 4, 5, 999]
    for idx in sample_indices:
        k_count = idx % 5
        for k in range(k_count):
            expected_x = float(idx + k + 0.1)
            expected_y = float(idx + k + 0.2)
            expected_conf = float(idx % 100) / 100.0
            expected_cls = idx * 10 + k
        for k in range(k_count, 4):
            pass