# Copyright (C) 2025 Embedl AB

"""Tests for the calibration data module in the Embedl Hub SDK."""

from pathlib import Path

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

from embedl_hub.core.quantization.calibration_data import (
    load_calibration_dataset,
    load_onnx_calibration_data,
    load_tflite_calibration_data,
)


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)


def test_load_calibration_dataset_single_input(tmp_path: Path):
    """Test load_calibration_dataset with a single input model."""
    data_path = tmp_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))

    dataset = load_calibration_dataset(data_path, ["input1"])
    assert isinstance(dataset, dict)
    assert "input1" in dataset
    assert len(dataset["input1"]) == 3
    sample = dataset["input1"][0]
    assert sample.shape == (1, 3, 16, 16)


def test_load_calibration_dataset_multiple_input(tmp_path: Path):
    """Test load_calibration_dataset with a multiple input model."""
    input_names = ["input1", "input2"]
    data_path = tmp_path / "cal_data"
    data_path.mkdir()
    for name in input_names:
        input_dir = data_path / name
        input_dir.mkdir()
        for i in range(2):
            np.save(
                input_dir / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16)
            )

    dataset = load_calibration_dataset(data_path, input_names)

    assert isinstance(dataset, dict)
    assert "input1" in dataset
    assert "input2" in dataset
    assert len(dataset["input1"]) == 2
    assert len(dataset["input2"]) == 2
    input1_sample = dataset["input1"][0]
    assert input1_sample.shape == (1, 3, 16, 16)
    input2_sample = dataset["input2"][0]
    assert input2_sample.shape == (1, 3, 16, 16)


def test_load_calibration_dataset_invalid_data_path(tmp_path: Path):
    """Test error when data_path is not a directory."""
    # Test with non-existent path
    non_existent_path = tmp_path / "non_existent"
    with pytest.raises(ValueError, match="Invalid data path"):
        load_calibration_dataset(non_existent_path, ["input1"])


def test_load_calibration_dataset_missing_input_dir_for_multi_input(
    tmp_path: Path,
):
    """Test error when an input directory is missing for a multi-input model."""
    data_path = tmp_path / "cal_data"
    data_path.mkdir()
    # Only create dir for input1, input2 directory is missing entirely

    with pytest.raises(FileNotFoundError, match="Input directory not found"):
        load_calibration_dataset(data_path, ["input1", "input2"])


def test_load_calibration_dataset_empty_input_dir_for_multi_input(
    tmp_path: Path,
):
    """Test error when an input directory exists but contains no .npy files."""
    data_path = tmp_path / "cal_data"
    data_path.mkdir()
    (data_path / "input1").mkdir()  # Create dir but leave it empty
    (data_path / "input2").mkdir()  # Create dir but leave it empty

    with pytest.raises(FileNotFoundError, match="No .npy files found"):
        load_calibration_dataset(data_path, ["input1", "input2"])


def test_load_calibration_dataset_no_npy_files_found(tmp_path: Path):
    """Test error when no .npy files are found in the data directory."""
    data_path = tmp_path / "cal_data"
    data_path.mkdir()

    with pytest.raises(FileNotFoundError, match="No .npy files found"):
        load_calibration_dataset(data_path, ["input1"])


def test_load_calibration_dataset_mismatched_sample_counts(tmp_path: Path):
    """Test error when inputs have different numbers of samples."""
    input_names = ["input1", "input2"]
    data_path = tmp_path / "cal_data"
    data_path.mkdir()

    # Create different numbers of samples for each input
    input1_dir = data_path / "input1"
    input1_dir.mkdir()
    for i in range(3):
        np.save(input1_dir / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16))

    input2_dir = data_path / "input2"
    input2_dir.mkdir()
    for i in range(2):  # Different count
        np.save(input2_dir / f"sample_{i}.npy", np.random.rand(1, 3, 16, 16))

    with pytest.raises(
        ValueError, match="All inputs must have the same number of samples"
    ):
        load_calibration_dataset(data_path, input_names)


def test_load_onnx_calibration_data_single_input(tmp_path: Path):
    """Test load_onnx_calibration_data with a single input model."""
    model_path = tmp_path / "single_input.onnx"
    _create_onnx_model(model_path, ["input1"])

    data_path = tmp_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))

    dataset = load_onnx_calibration_data(model_path, data_path)

    assert "input1" in dataset
    assert len(dataset["input1"]) == 5


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

    data_path = tmp_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)
            )

    dataset = load_onnx_calibration_data(model_path, data_path)

    assert "input1" in dataset
    assert "input2" in dataset
    assert len(dataset["input1"]) == 5
    assert len(dataset["input2"]) == 5


@pytest.fixture
def simple_tflite_model(tmp_path: Path):
    """Create a simple TensorFlow Lite model for testing."""
    # Create a simple TensorFlow model
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(3, 16, 16), name='input1'),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(10, name='output'),
        ]
    )

    # Convert to TensorFlow Lite
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    tflite_model = converter.convert()

    # Save the model
    model_path = tmp_path / "simple_model.tflite"
    with open(model_path, "wb") as f:
        f.write(tflite_model)

    return model_path


def test_load_tflite_calibration_data_single_input(
    tmp_path: Path, simple_tflite_model
):
    """Test load_tflite_calibration_data with a single input model."""
    # Create interpreter
    interpreter = tf.lite.Interpreter(model_path=str(simple_tflite_model))
    interpreter.allocate_tensors()

    # Create calibration data
    data_path = tmp_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))

    # Test the function
    dataset = load_tflite_calibration_data(interpreter, data_path)

    assert isinstance(dataset, list)
    assert len(dataset) == 3
    assert isinstance(dataset[0], dict)
    assert "input1" in dataset[0]
    assert dataset[0]["input1"].shape == (1, 3, 16, 16)


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