from typing import Mapping, Iterable

from image_prediction.exceptions import UnexpectedLabelFormat
from image_prediction.label_mapper.mapper import LabelMapper


class IndexLabelMapper(LabelMapper):
    def __init__(self, labels: Mapping[int, str]):
        self.__labels = labels

    def __validate_int_prediction_format(self, index_label: int) -> None:
        if not 0 <= index_label <= len(self.__labels):
            raise UnexpectedLabelFormat(
                f"Received index label  '{index_label}' that has no associated string label."
            )

    def __map_label(self, index_label: int) -> str:
        self.__validate_int_prediction_format(index_label)
        return self.__labels[index_label]

    def map_labels(self, index_labels: Iterable[int]) -> Iterable[str]:
        return map(self.__map_label, index_labels)