# coding: utf-8
"""collection with :class:`~supervisely.annotation.obj_class.ObjClass` instances"""

# docs
from __future__ import annotations

from collections import defaultdict
from typing import Dict, Iterator, List, Optional

from supervisely import logger
from supervisely.annotation.obj_class import ObjClass
from supervisely.annotation.renamer import Renamer
from supervisely.collection.key_indexed_collection import KeyIndexedCollection
from supervisely.geometry.graph import GraphNodes
from supervisely.imaging.color import hex2rgb, rgb2hex
from supervisely.io.json import JsonSerializable


class ObjClassCollection(KeyIndexedCollection, JsonSerializable):
    """Collection with :class:`~supervisely.annotation.obj_class.ObjClass` instances. :class:`~supervisely.annotation.obj_class_collection.ObjClassCollection` object is immutable."""

    item_type = ObjClass

    def __init__(self, items: Optional[List[ObjClass]] = None):
        """
        :param items: List of :class:`~supervisely.annotation.obj_class.ObjClass` instances.
        :type items: list, optional
        :raises :class:`~supervisely.collection.key_indexed_collection.DuplicateKeyError`: if instance with given name already exist

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                # Create ObjClass (see class ObjClass for more information)
                class_lemon = sly.ObjClass('lemon', sly.Rectangle)
                class_kiwi = sly.ObjClass('kiwi', sly.Bitmap)

                class_arr = [class_lemon, class_kiwi]

                # Create ObjClassCollection from ObjClasses
                classes = sly.ObjClassCollection(class_arr)

                # Add items to ObjClassCollection
                class_potato = sly.ObjClass('potato', sly.Bitmap)

                # Remember that ObjClassCollection is immutable, and we need to assign new instance of ObjClassCollection to a new variable
                classes = classes.add(class_potato)

                # You can also add multiple items to collection
                class_cabbage = sly.ObjClass('cabbage', sly.Rectangle)
                class_carrot = sly.ObjClass('carrot', sly.Bitmap)
                class_turnip = sly.ObjClass('turnip', sly.Polygon)

                classes = classes.add_items([class_cabbage, class_carrot, class_turnip])

                # Has key, checks if given key exist in collection
                classes.has_key("cabbage")
                # Output: True

                # Intersection, finds intersection of given list of instances with collection items
                class_dog = sly.ObjClass('dog', sly.Rectangle)
                class_cat = sly.ObjClass('cat', sly.Rectangle)
                class_turtle = sly.ObjClass('turtle', sly.Rectangle)

                classes_animals = sly.ObjClassCollection([class_dog, class_cat, class_turtle])

                classes_intersections = classes.intersection(classes_animals)
                print(classes_intersections.to_json())
                # Output: []

                # Let's add the potato ObjClass from another collection and compare them again
                classes_animals = classes_animals.add(class_potato)

                classes_intersections = classes.intersection(classes_animals)
                print(classes_intersections.to_json())
                # Output: [
                #     {
                #         "title":"potato",
                #         "shape":"bitmap",
                #         "color":"#8A570F",
                #         "geometry_config":{},
                #         "hotkey":""
                #     }
                # ]

                # Difference, finds difference between collection and given list of ObjClass or ObjClassCollection
                class_car = sly.ObjClass('car', sly.Rectangle)
                class_bicycle = sly.ObjClass('bicycle', sly.Rectangle)

                classes_vehicles = sly.ObjClassCollection([class_car, class_bicycle])

                class_pedestrian = sly.ObjClass('pedestrian', sly.Rectangle)
                class_road = sly.ObjClass('road', sly.Rectangle)

                difference = classes_vehicles.difference([class_pedestrian, class_road])
                print(difference.to_json())
                # Output: [
                #     {
                #         "title":"car",
                #         "shape":"rectangle",
                #         "color":"#8A0F3B",
                #         "geometry_config":{},
                #         "hotkey":""
                #     },
                #     {
                #         "title":"bicycle",
                #         "shape":"rectangle",
                #         "color":"#0F8A1F",
                #         "geometry_config":{},
                #         "hotkey":""
                #     }
                # ]

                # Merge, merges collection and given list of ObjClasses
                c_1 = sly.ObjClassCollection([class_car, class_bicycle])
                c_2 = sly.ObjClassCollection([class_pedestrian, class_road])

                c_3 = c_1.merge(c_2)
                print(c_3.to_json())
                # Output: [
                #     {
                #         "title":"pedestrian",
                #         "shape":"rectangle",
                #         "color":"#8A0F27",
                #         "geometry_config":{},
                #         "hotkey":""
                #     },
                #     {
                #         "title":"road",
                #         "shape":"rectangle",
                #         "color":"#8A620F",
                #         "geometry_config":{},
                #         "hotkey":""
                #     },
                #     {
                #         "title":"car",
                #         "shape":"rectangle",
                #         "color":"#8A0F3B",
                #         "geometry_config":{},
                #         "hotkey":""
                #     },
                #     {
                #         "title":"bicycle",
                #         "shape":"rectangle",
                #         "color":"#0F8A1F",
                #         "geometry_config":{},
                #         "hotkey":""
                #     }
                # ]

                # Merge will raise ValueError if item name from given list is in collection but items in both are different
                class_bicycle_1 = sly.ObjClass('bicycle', sly.Rectangle)
                class_bicycle_2 = sly.ObjClass('bicycle', sly.Bitmap)

                classes_1 = sly.ObjClassCollection([class_bicycle_1])
                classes_2 = sly.ObjClassCollection([class_bicycle_2])

                test_merge = classes_1.merge(classes_2)
                # Output: ValueError: Error during merge for key 'bicycle': values are different

                # Let's try to create now a collection where ObjClasses have identical names
                class_cow = sly.ObjClass('cow', sly.Rectangle)
                class_chicken = sly.ObjClass('cow', sly.Rectangle)

                test_classes = sly.ObjClassCollection([class_cow, class_chicken])
                # Output: DuplicateKeyError: "Key 'cow' already exists"
        """
        super().__init__(items)

    def to_json(self) -> List[Dict]:
        """
        Convert the ObjClassCollection to a list of json dicts. Read more about `Supervisely format <https://docs.supervisely.com/data-organization/00_ann_format_navi>`_.

        :returns: List of dicts in json format
        :rtype: :class:`List[dict]`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                class_lemon = sly.ObjClass('lemon', sly.Rectangle)
                class_kiwi = sly.ObjClass('kiwi', sly.Bitmap)

                # Add ObjClasses to ObjClassCollection
                classes = sly.ObjClassCollection([class_lemon, class_kiwi])

                classes_json = classes.to_json()
                print(classes_json)
                # Output: [
                #      {
                #           "title": "lemon",
                #           "shape": "rectangle",
                #           "color": "#300F8A",
                #           "geometry_config": {},
                #           "hotkey": ""
                #      },
                #      {
                #           "title": "kiwi",
                #           "shape": "bitmap",
                #           "color": "#7C0F8A",
                #           "geometry_config": {},
                #           "hotkey": ""
                #      }
                # ]
        """
        return [obj_class.to_json() for obj_class in self]

    @classmethod
    def from_json(cls, data: List[Dict]) -> ObjClassCollection:
        """
        Convert a list with dicts in json format to ObjClassCollection. Read more about `Supervisely format <https://docs.supervisely.com/data-organization/00_ann_format_navi>`_.

        :param data: List with dicts in json format.
        :type data: List[dict]
        :returns: ObjClassCollection object
        :rtype: :class:`~supervisely.annotation.obj_class_collection.ObjClassCollection`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                data = [
                    {
                        "title": "lemon",
                        "shape": "rectangle",
                        "color": "#300F8A",
                        "hotkey": ""
                    },
                    {
                        "title": "kiwi",
                        "shape": "bitmap",
                        "color": "#7C0F8A",
                        "hotkey": ""
                    }
                ]

                classes = sly.ObjClassCollection.from_json(data)
        """
        obj_classes = [ObjClass.from_json(obj_class_json) for obj_class_json in data]
        return cls(obj_classes)

    def validate_classes_colors(self, logger: Optional[logger] = None) -> str or None:
        """
        Checks for unique colors in the ObjClassCollection.

        :param logger: Input logger.
        :type logger: logger, optional
        :returns: Notification if there are objects with the same colors, otherwise None
        :rtype: str or None

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                # Let's create 2 ObjClasses with the same color
                class_lemon = sly.ObjClass('lemon', sly.Rectangle, [0, 0, 0])
                class_kiwi = sly.ObjClass('kiwi', sly.Bitmap, [0, 0, 0])

                # Add them to ObjClassCollection
                classes = sly.ObjClassCollection([class_lemon, class_kiwi])

                print(classes.validate_classes_colors())
                # Output: Classes ['lemon', 'kiwi'] have the same RGB color = [0, 0, 0]

                # Now let's change colors of our ObjClasses
                class_lemon = sly.ObjClass('lemon', sly.Rectangle, [255, 0, 0])
                class_kiwi = sly.ObjClass('kiwi', sly.Bitmap, [0, 0, 255])

                classes = sly.ObjClassCollection([class_lemon, class_kiwi])

                print(classes.validate_classes_colors())
                # Output: None
        """
        color_names = defaultdict(list)
        for obj_class in self:
            hex = rgb2hex(obj_class.color)
            color_names[hex].append(obj_class.name)

        class_colors_notify = None
        for hex_color, class_names in color_names.items():
            if len(class_names) > 1:
                warn_str = "Classes {!r} have the same RGB color = {!r}".format(
                    class_names, hex2rgb(hex_color)
                )
                if logger is not None:
                    logger.warn(warn_str)
                if class_colors_notify is None:
                    class_colors_notify = ""
                class_colors_notify += warn_str + "\n\n"
        return class_colors_notify

    def __iter__(self) -> Iterator[ObjClass]:
        return next(self)

    def refresh_ids_from(self, classes: ObjClassCollection) -> None:
        """Updates ids of classes in the collection from given classes.

        :param classes: Collection with classes to update ids from.
        :type classes: :class:`~supervisely.annotation.obj_class_collection.ObjClassCollection`
        :returns: None
        :rtype: None
        """
        for new_class in classes:
            my_class = self.get(new_class.name)
            if my_class is None:
                continue
            my_class._set_id(new_class.sly_id)


def make_renamed_classes(
    src_obj_classes: ObjClassCollection,
    renamer: Renamer,
    skip_missing: Optional[bool] = False,
) -> ObjClassCollection:
    """Returns a new ObjClassCollection with renamed classes.

    :param src_obj_classes: ObjClassCollection to rename.
    :type src_obj_classes: Input ObjClassCollection
    :param renamer: :class:`~supervisely.annotation.renamer.Renamer` object, which will handle renaming process.
    :type renamer: :class:`~supervisely.annotation.renamer.Renamer`
    :param skip_missing: If True, missing classes will be skipped, otherwise KeyError will be raised.
    :type skip_missing: Optional[bool]
    :returns: New ObjClassCollection with renamed classes.
    :rtype: :class:`~supervisely.annotation.obj_class_collection.ObjClassCollection`
    :raises KeyError: If skip_missing is False and some classes could not be renamed.
    """
    renamed_classes = []
    for src_cls in src_obj_classes:
        renamed_name = renamer.rename(src_cls.name)
        if renamed_name is not None:
            renamed_classes.append(src_cls.clone(name=renamed_name))
        elif not skip_missing:
            raise KeyError(
                "Object class name {} could not be mapped to a destination name.".format(
                    src_cls.name
                )
            )
    return ObjClassCollection(items=renamed_classes)
