from __future__ import annotations

import gc
import os
import time

import cv2
import numpy as np
import pytest
import supervision as sv
from inference.core.workflows.core_steps.common.serializers import (
    MIN_POLYGON_POINT_COUNT, _attach_parent_metadata_to_detection_dict,
    mask_to_polygon, serialise_sv_detections)
from inference.core.workflows.execution_engine.constants import (
    AREA_CONVERTED_KEY_IN_INFERENCE_RESPONSE,
    AREA_CONVERTED_KEY_IN_SV_DETECTIONS, AREA_KEY_IN_INFERENCE_RESPONSE,
    AREA_KEY_IN_SV_DETECTIONS, BOUNDING_RECT_ANGLE_KEY_IN_INFERENCE_RESPONSE,
    BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS,
    BOUNDING_RECT_HEIGHT_KEY_IN_INFERENCE_RESPONSE,
    BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS,
    BOUNDING_RECT_RECT_KEY_IN_INFERENCE_RESPONSE,
    BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS,
    BOUNDING_RECT_WIDTH_KEY_IN_INFERENCE_RESPONSE,
    BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS, CLASS_ID_KEY, CLASS_NAME_KEY,
    CONFIDENCE_KEY)
from inference.core.workflows.execution_engine.constants import \
    DETECTED_CODE_KEY
from inference.core.workflows.execution_engine.constants import \
    DETECTED_CODE_KEY as DETECTED_CODE_KEY_CONST
from inference.core.workflows.execution_engine.constants import (
    DETECTION_ID_KEY, HEIGHT_KEY, IMAGE_DIMENSIONS_KEY,
    KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS,
    KEYPOINTS_KEY_IN_INFERENCE_RESPONSE, KEYPOINTS_XY_KEY_IN_SV_DETECTIONS,
    PARENT_COORDINATES_KEY, PARENT_DIMENSIONS_KEY, PARENT_ID_KEY,
    PARENT_ORIGIN_KEY, PATH_DEVIATION_KEY_IN_INFERENCE_RESPONSE,
    PATH_DEVIATION_KEY_IN_SV_DETECTIONS, POLYGON_KEY,
    POLYGON_KEY_IN_INFERENCE_RESPONSE, POLYGON_KEY_IN_SV_DETECTIONS,
    ROOT_PARENT_COORDINATES_KEY, ROOT_PARENT_DIMENSIONS_KEY,
    ROOT_PARENT_ID_KEY, ROOT_PARENT_ORIGIN_KEY,
    SMOOTHED_SPEED_KEY_IN_INFERENCE_RESPONSE,
    SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS,
    SMOOTHED_VELOCITY_KEY_IN_INFERENCE_RESPONSE,
    SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS, SPEED_KEY_IN_INFERENCE_RESPONSE,
    SPEED_KEY_IN_SV_DETECTIONS, TIME_IN_ZONE_KEY_IN_INFERENCE_RESPONSE,
    TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, TRACKER_ID_KEY,
    VELOCITY_KEY_IN_INFERENCE_RESPONSE, VELOCITY_KEY_IN_SV_DETECTIONS,
    WIDTH_KEY, X_KEY, Y_KEY)
from inference.core.workflows.execution_engine.entities.base import \
    ParentOrigin


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_basic_no_mask_serialisation_single_detection():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = [10, 20, 30, 60]
    mask = None
    confidence = 0.85
    class_id = 2
    tracker_id = None
    data = {CLASS_NAME_KEY: 'person', DETECTION_ID_KEY: 'det-123', IMAGE_DIMENSIONS_KEY: np.array([480, 640])}
    detections = sv.Detections(xyxy=np.array([xyxy]), mask=np.array([None], dtype=object), confidence=np.array([confidence]), class_id=np.array([class_id]), tracker_id=np.array([tracker_id], dtype=object), data=[data])
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_basic_no_mask_serialisation_single_detection', 'serialise_sv_detections', '7', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    pred = result['predictions'][0]

def test_mask_serialisation_generates_polygon_from_mask():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((20, 20), dtype=np.uint8)
    cv2.rectangle(mask, (5, 5), (10, 12), color=1, thickness=-1)
    xyxy = [5, 5, 10, 12]
    confidence = 0.99
    class_id = 1
    tracker_id = 7
    data = {CLASS_NAME_KEY: 'box', DETECTION_ID_KEY: 'box-1'}
    detections = sv.Detections(xyxy=np.array([xyxy]), mask=np.array([mask]), confidence=np.array([confidence]), class_id=np.array([class_id]), tracker_id=np.array([tracker_id]), data=[data])
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_mask_serialisation_generates_polygon_from_mask', 'serialise_sv_detections', '8', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    pred = out['predictions'][0]

def test_polygon_in_data_overrides_mask_and_serialises_in_inference_response_form():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask = np.zeros((10, 10), dtype=np.uint8)
    cv2.rectangle(mask, (1, 1), (2, 2), color=1, thickness=-1)
    provided_polygon = np.array([[1.2, 1.8], [2.5, 1.1], [2.2, 2.9], [1.0, 2.0]])
    xyxy = [1, 1, 2, 2]
    data = {CLASS_NAME_KEY: 'poly', DETECTION_ID_KEY: 'poly-1', POLYGON_KEY_IN_SV_DETECTIONS: provided_polygon}
    detections = sv.Detections(xyxy=np.array([xyxy]), mask=np.array([mask]), confidence=np.array([0.5]), class_id=np.array([0]), tracker_id=np.array([None], dtype=object), data=[data])
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_polygon_in_data_overrides_mask_and_serialises_in_inference_response_form', 'serialise_sv_detections', '6', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    pred = out['predictions'][0]
    poly_out = pred[POLYGON_KEY_IN_INFERENCE_RESPONSE]

def test_mask_to_polygon_returns_none_for_empty_mask():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    empty_mask = np.zeros((8, 8), dtype=np.uint8)
    polygon = mask_to_polygon(empty_mask)

def test_mask_to_polygon_pads_short_contours_to_minimum_points():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    single_point_mask = np.zeros((6, 6), dtype=np.uint8)
    single_point_mask[3, 4] = 1
    contour = mask_to_polygon(single_point_mask)
    if contour.shape[0] > 1:
        pass

def test_empty_detections_list_results_in_empty_predictions_and_none_image_dimensions():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    detections = sv.Detections(xyxy=np.empty((0, 4)), mask=np.empty((0,), dtype=object), confidence=np.empty((0,)), class_id=np.empty((0,), dtype=int), tracker_id=np.empty((0,), dtype=object), data=[])
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_empty_detections_list_results_in_empty_predictions_and_none_image_dimensions', 'serialise_sv_detections', '1', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output

def test_attach_parent_metadata_with_numpy_values_sets_parent_origin_fields():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    detection_dict = {}
    parent_coords = np.array([12.4, 34.6])
    parent_dims = np.array([200.2, 300.7])
    data = {PARENT_COORDINATES_KEY: parent_coords, PARENT_DIMENSIONS_KEY: parent_dims}
    _attach_parent_metadata_to_detection_dict(detection_dict=detection_dict, data=data, coordinates_key=PARENT_COORDINATES_KEY, dimensions_key=PARENT_DIMENSIONS_KEY, origin_key=PARENT_ORIGIN_KEY)
    origin = detection_dict[PARENT_ORIGIN_KEY]

def test_keypoints_and_motion_and_area_fields_serialise_correct_types():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    xyxy = [0, 0, 10, 10]
    mask = None
    confidence = 0.7
    class_id = 5
    tracker_id = None
    kp_class_id = [1, 2]
    kp_class_name = ['left_eye', 'right_eye']
    kp_confidence = [0.9, 0.8]
    kp_xy = [[1.0, 2.0], [3.0, 4.0]]
    velocity = np.array([0.1, -0.2])
    smoothed_velocity = np.array([0.0, 0.0])
    speed = np.array(3.5)
    smoothed_speed = np.array(2.5)
    area = 123.45
    area_converted = 1.2345
    data = {CLASS_NAME_KEY: 'kp', DETECTION_ID_KEY: 'kp-1', KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS: kp_class_id, KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS: kp_class_name, KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS: kp_confidence, KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: kp_xy, VELOCITY_KEY_IN_SV_DETECTIONS: velocity, SPEED_KEY_IN_SV_DETECTIONS: speed, SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS: smoothed_velocity, SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS: smoothed_speed, AREA_KEY_IN_SV_DETECTIONS: area, AREA_CONVERTED_KEY_IN_SV_DETECTIONS: area_converted}
    detections = sv.Detections(xyxy=np.array([xyxy]), mask=np.array([mask], dtype=object), confidence=np.array([confidence]), class_id=np.array([class_id]), tracker_id=np.array([tracker_id], dtype=object), data=[data])
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_keypoints_and_motion_and_area_fields_serialise_correct_types', 'serialise_sv_detections', '17', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    pred = out['predictions'][0]
    kp_list = pred[KEYPOINTS_KEY_IN_INFERENCE_RESPONSE]

def test_large_number_of_detections_serialise_quickly_and_completely():
    np.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    n = 1000
    x1 = np.linspace(0, 100, n)
    y1 = np.linspace(0, 200, n)
    x2 = x1 + 10
    y2 = y1 + 20
    xyxy = np.vstack([x1, y1, x2, y2]).T.astype(float)
    confidences = np.linspace(0.0, 1.0, n)
    class_ids = np.zeros(n, dtype=int)
    tracker_ids = np.array([None] * n, dtype=object)
    data_list = [{CLASS_NAME_KEY: 'bulk', DETECTION_ID_KEY: f'bulk-{i}'} for i in range(n)]
    detections = sv.Detections(xyxy=xyxy, mask=np.empty((n,), dtype=object), confidence=confidences, class_id=class_ids, tracker_id=tracker_ids, data=data_list)
    codeflash_return_value = codeflash_wrap(serialise_sv_detections, 'tests.inference.unit_tests.test_serialise_sv_detections__unit_test_0', None, 'test_large_number_of_detections_serialise_quickly_and_completely', 'serialise_sv_detections', '11', codeflash_loop_index, detections)
    codeflash_output = codeflash_return_value
    out = codeflash_output
    for idx in (0, n // 2, n - 1):
        pred = out['predictions'][idx]