# Copyright (C) 2025 Embedl AB

"""Tests for the quantization module in the Embedl Hub SDK."""

from pathlib import Path

import numpy as np
import onnx
import pytest
from onnx import TensorProto, helper

from embedl_hub.core.quantization.quantize import (
    _load_calibration_data,
    collect_calibration_data,
    get_input_shapes,
)


def _create_onnx_model(model_path: Path, input_names: list[str]):
    """Helper function to create a dummy ONNX model."""
    inputs = [
        helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 3, 16, 16])
        for name in input_names
    ]
    output = helper.make_tensor_value_info(
        "output", TensorProto.FLOAT, [1, 10]
    )
    nodes = [
        helper.make_node("Add", [name, name], [f"add_{name}"])
        for name in input_names
    ]
    # Dummy node to consume all inputs
    final_node = helper.make_node(
        "Sum", [f"add_{name}" for name in input_names], ["output"]
    )

    graph_def = helper.make_graph(
        nodes + [final_node], "test-model", inputs, [output]
    )
    model_def = helper.make_model(graph_def, producer_name="test")
    onnx.save(model_def, model_path)


@pytest.fixture
def temp_path(tmp_path: Path) -> Path:
    """Fixture to create a temporary directory for tests."""
    return tmp_path


def test_get_input_shapes_single_input(temp_path: Path):
    """Test get_input_shapes with a single input model."""
    model_path = temp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    input_shapes = get_input_shapes(model_path)
    assert isinstance(input_shapes, dict)
    assert "input1" in input_shapes
    assert input_shapes["input1"] == (1, 3, 16, 16)


def test_get_input_shapes_multiple_input(temp_path: Path):
    """Test get_input_shapes with a multiple input model."""
    model_path = temp_path / "multi_input.onnx"
    input_names = ["input1", "input2"]
    _create_onnx_model(model_path, input_names)

    input_shapes = get_input_shapes(model_path)
    assert isinstance(input_shapes, dict)
    assert "input1" in input_shapes
    assert "input2" in input_shapes
    assert input_shapes["input1"] == (1, 3, 16, 16)
    assert input_shapes["input2"] == (1, 3, 16, 16)


def test_collect_calibration_data_no_data_path(temp_path: Path):
    """Test collect_calibration_data with no data path (random data)."""
    model_path = temp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    num_samples, data = collect_calibration_data(
        model_path=model_path, data_path=None, num_samples=5
    )

    assert num_samples == 1  # Random data returns 1 sample
    assert "input1" in data
    assert len(data["input1"]) == 1
    assert data["input1"][0].shape == (1, 3, 16, 16)


def test_load_calibration_data_single_input(temp_path: Path):
    """Test _load_calibration_data with a single input model."""
    model_path = temp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    data_path = temp_path / "cal_data"
    data_path.mkdir()
    for i in range(5):
        np.save(data_path / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16))

    num_samples, data = _load_calibration_data(model_path, data_path, 3)

    assert "input1" in data
    assert len(data["input1"]) == 3
    assert num_samples == 3


def test_load_calibration_data_multi_input(temp_path: Path):
    """Test _load_calibration_data with a multi-input model."""
    model_path = temp_path / "multi_input.onnx"
    input_names = ["input1", "input2"]
    _create_onnx_model(model_path, input_names)

    data_path = temp_path / "cal_data"
    data_path.mkdir()
    for name in input_names:
        input_dir = data_path / name
        input_dir.mkdir()
        for i in range(5):
            np.save(
                input_dir / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16)
            )

    num_samples, data = _load_calibration_data(model_path, data_path, 2)

    assert "input1" in data
    assert "input2" in data
    assert len(data["input1"]) == 2
    assert len(data["input2"]) == 2
    assert num_samples == 2


def test_load_calibration_data_num_samples_greater_than_available(
    temp_path: Path,
):
    """Test _load_calibration_data when num_samples is more than available."""
    model_path = temp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    data_path = temp_path / "cal_data"
    data_path.mkdir()
    for i in range(3):
        np.save(data_path / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16))

    num_samples, data = _load_calibration_data(model_path, data_path, 10)

    assert "input1" in data
    assert len(data["input1"]) == 3
    assert num_samples == 3


def test_collect_calibration_data_with_data_path(temp_path: Path):
    """Test collect_calibration_data with actual data path."""
    model_path = temp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    data_path = temp_path / "cal_data"
    data_path.mkdir()
    for i in range(4):
        np.save(data_path / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16))

    num_samples, data = collect_calibration_data(
        model_path=model_path, data_path=data_path, num_samples=2
    )

    assert num_samples == 2
    assert "input1" in data
    assert len(data["input1"]) == 2
    assert data["input1"][0].shape == (1, 3, 16, 16)


if __name__ == "__main__":
    pytest.main([__file__])
