from ...entities import ImageEntity
from ...model.engine import DiffusionPipeline, Engine, PreTrainedModel
from ...model.lazy import LazyExternal
from ...model.vendor import TextGenerationVendor
from ...model.vision import BaseVisionModel

from typing import Any, Literal, cast

from PIL import Image
from torch import inference_mode

AutoImageProcessor = LazyExternal("transformers", "AutoImageProcessor")
AutoModelForImageClassification = LazyExternal(
    "transformers", "AutoModelForImageClassification"
)


# model predicts one of the 1000 ImageNet classes
class ImageClassificationModel(BaseVisionModel):
    def _load_model(
        self,
    ) -> PreTrainedModel | TextGenerationVendor | DiffusionPipeline:
        self._processor = cast(
            Any,
            cast(Any, AutoImageProcessor).from_pretrained(
                self._model_id,
                backend="torchvision",
            ),
        )
        assert self._model_id
        model = cast(
            PreTrainedModel,
            cast(Any, AutoModelForImageClassification).from_pretrained(
                self._model_id,
                device_map=self._device,
                tp_plan=Engine._get_tp_plan(self._settings.parallel),
                distributed_config=Engine._get_distributed_config(
                    self._settings.distributed_config
                ),
            ),
        )
        return model

    async def __call__(
        self,
        image_source: str | Image.Image,
        tensor_format: Literal["pt"] = "pt",
    ) -> ImageEntity:
        image = BaseVisionModel._get_image(image_source)
        inputs = self._processor(image, return_tensors=tensor_format)
        inputs.to(self._device)

        with inference_mode():
            logits = self._model(**inputs).logits

        label_index = logits.argmax(dim=1).item()
        return ImageEntity(label=self._model.config.id2label[label_index])
