from unittest.mock import Mock

import tensorflow as tf

from deep_learning.models.yolo import YoloModelBuilder, box_loss, build_yolo_preprocessor


def test_build_yolo_preprocessor():
    preprocessor = build_yolo_preprocessor()
    image = tf.zeros((1, 32, 24, 3), dtype=tf.uint8)

    output = preprocessor(image)

    assert output.shape == (1, 448, 448, 3)


def test_yolo_model_builder_builds_training_model():
    """验证 YOLO 模型构建器会输出训练模型需要的 box 和 class 预测。"""
    artifact = YoloModelBuilder(
        image_size=448,
        grid_size=6,
        num_labels=91
    ).build_training_artifact()
    model = artifact.model
    model.summary()
    inputs = tf.zeros((2, 448, 448, 3), dtype=tf.float32)

    outputs = model(inputs, training=False)

    assert "box" in outputs
    assert "class" in outputs
    assert outputs["box"].shape == (2, 6, 6, 5)
    assert outputs["class"].shape == (2, 6, 6, 91)


def test_yolo_model_builder_compiles_training_model():
    """验证 YOLO 模型构建器会使用 box 和 class 两路损失编译模型。"""
    model = Mock()
    builder = YoloModelBuilder(
        image_size=448,
        grid_size=6,
        num_labels=91
    )

    builder.compile_training_model(model)

    model.compile.assert_called_once()
    _, kwargs = model.compile.call_args
    assert kwargs["optimizer"].__class__.__name__ == "Adam"
    assert kwargs["loss"] == {
        "box": box_loss,
        "class": "sparse_categorical_crossentropy"
    }