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, unique

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


class SemanticSegmentationModel(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, AutoModelForSemanticSegmentation).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
                ),
            ),
        )
        cast(Any, model).eval()
        return model

    async def __call__(
        self,
        image_source: str | Image.Image,
        tensor_format: Literal["pt"] = "pt",
    ) -> list[str]:
        image = BaseVisionModel._get_image(image_source)
        inputs = self._processor(images=image, return_tensors=tensor_format)
        inputs.to(self._device)
        with inference_mode():
            logits = self._model(**inputs).logits
        # shape (height, width) with class indices
        mask = logits.argmax(dim=1)[0]
        labels_tensor = unique(mask)
        labels = [
            self._model.config.id2label[idx.item()] for idx in labels_tensor
        ]
        return labels
