from typing import Iterable

from image_prediction.formatter.formatter import Formatter


class Snake2CamelCaseKeyFormatter(Formatter):

    @staticmethod
    def __format_key(key):

        if isinstance(key, str):
            head, *tail = key.split("_")
            return head + "".join(map(lambda s: s.title(), tail))
        else:
            return key

    def __format(self, data):

        # If we wanted to do this properly, we would need handlers for all expected types and dispatch based
        # on a type comparison. This is too much engineering for the limited use-case of this class though.
        if isinstance(data, Iterable) and not isinstance(data, dict) and not isinstance(data, str):
            f = map(self.__format, data)
            return type(data)(f) if not isinstance(data, map) else f

        if not isinstance(data, dict):
            return data

        keys_formatted = list(map(self.__format_key, data))

        return dict(zip(keys_formatted, map(self.__format, data.values())))

    def format(self, data):
        return self.__format(data)