from __future__ import annotations

import gc
import os
import time
from unittest.mock import patch

import matplotlib.pyplot as plt
import numpy as np
import pytest
import torch
from inference.core.workflows.execution_engine.entities.base import (
    ImageParentMetadata, WorkflowImageData)
from inference.models.depth_anything_v2.depth_anything_v2_inference_models import \
    InferenceModelsDepthAnythingV2Adapter


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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        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 make_model_returning_tensor(tensor: torch.Tensor):
    return lambda inputs: (tensor,)

def test_basic_predict_normalization_and_image_generation():
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input_np = np.zeros((2, 2), dtype=np.uint8)
    predictions = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
    module_path = 'inference.models.depth_anything_v2.depth_anything_v2_inference_models'
    attr_to_patch = f'{module_path}.AutoModel.from_pretrained'
    with patch(attr_to_patch) as mock_from_pretrained:
        mock_from_pretrained.return_value = make_model_returning_tensor(predictions)
        adapter = InferenceModelsDepthAnythingV2Adapter(model_id='dummy-model')
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_basic_predict_normalization_and_image_generation', 'predict', '5', codeflash_loop_index, input_np)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    result = result_tuple[0]
    expected_normalized = (predictions.numpy() - 1.0) / 3.0
    returned_normalized = result['normalized_depth']
    image_obj = result['image']
    depth_for_viz = (returned_normalized * 255.0).astype(np.uint8)
    cmap = plt.get_cmap('viridis')
    expected_colored = (cmap(depth_for_viz)[:, :, :3] * 255).astype(np.uint8)
    returned_colored = image_obj.numpy_image

def test_predict_raises_on_constant_depth_map():
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input_np = np.zeros((4, 4), dtype=np.uint8)
    constant_tensor = torch.full((4, 4), 5.0, dtype=torch.float32)
    module_path = 'inference.models.depth_anything_v2.depth_anything_v2_inference_models'
    attr_to_patch = f'{module_path}.AutoModel.from_pretrained'
    with patch(attr_to_patch) as mock_from_pretrained:
        mock_from_pretrained.return_value = make_model_returning_tensor(constant_tensor)
        adapter = InferenceModelsDepthAnythingV2Adapter(model_id='dummy-model')
    with pytest.raises(ValueError) as excinfo:
        codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_raises_on_constant_depth_map', 'predict', '5_0', codeflash_loop_index, input_np)

def test_predict_handles_negative_values_and_normalizes_correctly():
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input_np = np.zeros((3, 3), dtype=np.uint8)
    predictions = torch.tensor([[-10.0, 0.0, 10.0], [-5.0, 5.0, 15.0], [20.0, -20.0, 2.5]], dtype=torch.float32)
    module_path = 'inference.models.depth_anything_v2.depth_anything_v2_inference_models'
    attr_to_patch = f'{module_path}.AutoModel.from_pretrained'
    with patch(attr_to_patch) as mock_from_pretrained:
        mock_from_pretrained.return_value = make_model_returning_tensor(predictions)
        adapter = InferenceModelsDepthAnythingV2Adapter(model_id='dummy-model')
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_predict_handles_negative_values_and_normalizes_correctly', 'predict', '5', codeflash_loop_index, input_np)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    result = result_tuple[0]
    returned_normalized = result['normalized_depth']
    p = predictions.numpy()
    expected = (p - p.min()) / (p.max() - p.min())

def test_large_scale_predict_performs_on_large_but_safe_tensor():
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input_np = np.zeros((600, 600), dtype=np.uint8)
    h, w = (600, 600)
    y = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None]
    x = np.linspace(0.0, 1.0, w, dtype=np.float32)[None, :]
    gradient = torch.from_numpy((x + y).astype(np.float32))
    gradient = gradient * 100.0 - 50.0
    module_path = 'inference.models.depth_anything_v2.depth_anything_v2_inference_models'
    attr_to_patch = f'{module_path}.AutoModel.from_pretrained'
    with patch(attr_to_patch) as mock_from_pretrained:
        mock_from_pretrained.return_value = make_model_returning_tensor(gradient)
        adapter = InferenceModelsDepthAnythingV2Adapter(model_id='dummy-model')
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_large_scale_predict_performs_on_large_but_safe_tensor', 'predict', '9', codeflash_loop_index, input_np)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    result = result_tuple[0]
    nd = result['normalized_depth']
    image_obj = result['image']
    img = image_obj.numpy_image

def test_parent_metadata_parent_id_is_present_and_non_empty():
    np.random.seed(42)
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input_np = np.zeros((2, 2), dtype=np.uint8)
    predictions = torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=torch.float32)
    module_path = 'inference.models.depth_anything_v2.depth_anything_v2_inference_models'
    attr_to_patch = f'{module_path}.AutoModel.from_pretrained'
    with patch(attr_to_patch) as mock_from_pretrained:
        mock_from_pretrained.return_value = make_model_returning_tensor(predictions)
        adapter = InferenceModelsDepthAnythingV2Adapter(model_id='dummy-model')
    codeflash_return_value = codeflash_wrap(adapter.predict, 'tests.inference.unit_tests.test_predict__unit_test_0', None, 'test_parent_metadata_parent_id_is_present_and_non_empty', 'predict', '5', codeflash_loop_index, input_np)
    codeflash_output = codeflash_return_value
    result_tuple = codeflash_output
    image_obj = result_tuple[0]['image']
    parent_meta = image_obj.parent_metadata