import importlib
import pathlib
import pkgutil

from ..enums import AIRuntime, HardwareAcceleration, ModelFormat, Platform, Quantization
from ._common import ModelAttributes

# Discover all modules and subpackages under this package
__all__ = []
submodules = [
	importlib.import_module(f"{__name__}.{module_info.name}")
	for module_info in pkgutil.iter_modules([str(pathlib.Path(__file__).parent)])
	if module_info.name != "_common"
]

capabilities: dict[
	AIRuntime,
	dict[Platform, dict[Quantization, set[HardwareAcceleration]]],
] = {module.runtime: module.capabilities for module in submodules}
"""
Capabilities of each AI runtime on each platform for each quantization and hardware acceleration.
"""

format_distinguishes: dict[ModelFormat, list[ModelAttributes]] = {
	module.format: module.format_distinguishes for module in submodules
}
"""
A dictionary mapping model formats to lists of attributes that the format distinguishes.
"""

format: dict[AIRuntime, ModelFormat] = {
	module.runtime: module.format for module in submodules
}
"""
A dictionary mapping AI runtimes to model formats.
"""

__all__ = ["capabilities", "format_distinguishes", "format", "ModelAttributes"]
