# coding: utf-8

# docs
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

from supervisely.api.module_api import ApiField
from supervisely.collection.key_indexed_collection import KeyIndexedCollection
from supervisely.sly_logger import logger
from supervisely.video_annotation.frame import Frame
from supervisely.video_annotation.key_id_map import KeyIdMap
from supervisely.video_annotation.video_figure import VideoFigure
from supervisely.video_annotation.video_object_collection import VideoObjectCollection


class FrameCollection(KeyIndexedCollection):
    """
    Collection of Frame instances indexed by frame index. Immutable.

    :Usage Example:

        .. code-block:: python

            import supervisely as sly

            # Create two frames for collection
            fr_index_1 = 7
            frame_1 = sly.Frame(fr_index_1)
            fr_index_2 = 10
            frame_2 = sly.Frame(fr_index_2)

            # Create FrameCollection
            fr_collection = sly.FrameCollection([frame_1, frame_2])
            print(fr_collection.to_json())
            # Output: [
            #     {
            #         "index": 7,
            #         "figures": []
            #     },
            #     {
            #         "index": 10,
            #         "figures": []
            #     }
            # ]

            # Add item to FrameCollection
            frame_3 = sly.Frame(12)

            # Remember that TagCollection is immutable, and we need to assign new instance of TagCollection to a new variable
            new_fr_collection = fr_collection.add(frame_3)
            print(new_fr_collection.to_json())
            # Output: [
            #     {
            #         "index": 7,
            #         "figures": []
            #     },
            #     {
            #         "index": 10,
            #         "figures": []
            #     },
            #     {
            #         "index": 12,
            #         "figures": []
            #     }
            # ]

            # You can also add multiple items to collection
            frame_3 = sly.Frame(12)
            frame_4 = sly.Frame(15)

            # Remember that TagCollection is immutable, and we need to assign new instance of TagCollection to a new variable
            new_fr_collection = fr_collection.add_items([frame_3, frame_4])
            print(new_fr_collection.to_json())
            # Output: [
            #     {
            #         "index": 7,
            #         "figures": []
            #     },
            #     {
            #         "index": 10,
            #         "figures": []
            #     },
            #     {
            #         "index": 12,
            #         "figures": []
            #     },
            #     {
            #         "index": 15,
            #         "figures": []
            #     }
            # ]

            # Has key, checks if given key exist in collection
            fr_collection.has_key(7)
            # Output: True

            # Find intersection of given list of instances with collection items
            frame_1 = sly.Frame(7)
            frame_2 = sly.Frame(10)
            fr_collection = sly.FrameCollection([frame_1, frame_2])
            frame_3 = sly.Frame(12)
            frames_intersections = fr_collection.intersection([frame_3])
            print(frames_intersections.to_json())
            # Output: []

            frames_intersections = fr_collection.intersection([frame_2])
            print(frames_intersections.to_json())
            # Output: [
            #     {
            #         "index": 10,
            #         "figures": []
            #     }
            # ]

            # Note, two frames with the same index values are not equal
            frame_4 = sly.Frame(10)
            frames_intersections = fr_collection.intersection([frame_4])
            # Output:
            # ValueError: Different values for the same key 10

            # Find difference between collection and given list of Frames
            frames_difference = fr_collection.difference([frame_2])
            print(frames_difference.to_json())
            # Output: [
            #     {
            #         "index": 7,
            #         "figures": []
            #     }
            # ]

            # Merge collection and given list of FrameCollection
            frame_3 = sly.Frame(12)
            frame_4 = sly.Frame(15)
            over_collection = sly.FrameCollection([frame_3, frame_4])
            merged_collection = fr_collection.merge(over_collection)
            print(merged_collection.to_json())
            # Output: [
            #     {
            #         "index": 12,
            #         "figures": []
            #     },
            #     {
            #         "index": 15,
            #         "figures": []
            #     },
            #     {
            #         "index": 7,
            #         "figures": []
            #     },
            #     {
            #         "index": 10,
            #         "figures": []
            #     }
            # ]
    """

    item_type = Frame

    def get(self, key: str, default: Optional[Any] = None) -> Frame:
        """
        Get item from collection with given key(name) and set a default if item does not exist.

        :param key: Name of Frame in collection.
        :type key: str
        :param default: The value that is returned if there is no key in the collection.
        :type default:  Optional[Any]
        :returns: Frame, Slice or PointcloudEpisodeFrame object
        :rtype: :class:`~supervisely.collection.key_indexed_collection.KeyObject`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                frame_index = 7
                geometry = sly.Rectangle(0, 0, 100, 100)
                class_car = sly.ObjClass('car', sly.Rectangle)
                object_car = sly.VideoObject(class_car)
                figure_car = sly.VideoFigure(object_car, geometry, frame_index)

                frame = sly.Frame(frame_index, figures=[figure_car])
                frame_collection = sly.FrameCollection([frame])

                item = frame_collection.get(frame_index)
                print(item.to_json())
                # Output: {
                #     "figures": [
                #         {
                #         "geometry": {
                #             "points": {
                #             "exterior": [
                #                 [0, 0],
                #                 [100, 100]
                #             ],
                #             "interior": []
                #             }
                #         },
                #         "geometryType": "rectangle",
                #         "key": "713968a7d5384709bc5d4e63cd4535f2",
                #         "objectKey": "3342e68eff3b44dcb75712499265be55"
                #         }
                #     ],
                #     "index": 7
                # }
        """
        return super().get(key, default)

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

        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`, optional
        :returns: List of dicts in json format
        :rtype: List[dict]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                frame_1 = sly.Frame(7)
                frame_2 = sly.Frame(10)

                # Create FrameCollection
                fr_collection = sly.FrameCollection([frame_1, frame_2])
                print(fr_collection.to_json())
                # Output: [
                #     {
                #         "index": 7,
                #         "figures": []
                #     },
                #     {
                #         "index": 10,
                #         "figures": []
                #     }
                # ]
        """
        return [frame.to_json(key_id_map) for frame in self]

    @classmethod
    def from_json(
        cls,
        data: List[Dict],
        objects: VideoObjectCollection,
        frames_count: Optional[int] = None,
        key_id_map: Optional[KeyIdMap] = None,
        skip_corrupted: Optional[bool] = False,
    ) -> FrameCollection:
        """
        Convert a list of json dicts to FrameCollection. 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]
        :param objects: VideoObjectCollection object.
        :type objects: :class:`~supervisely.video_annotation.video_object_collection.VideoObjectCollection`
        :param frames_count: Number of frames in video.
        :type frames_count: int, optional
        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`, optional
        :returns: FrameCollection object.
        :rtype: :class:`~supervisely.video_annotation.frame_collection.FrameCollection`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                fr_collection_json = [
                    {"index": 7, "figures": []},
                    {"index": 10, "figures": []}
                ]

                objects = []
                fr_collection = sly.FrameCollection.from_json(fr_collection_json, objects)
        """
        frames = []
        for frame_json in data:
            try:
                frame = cls.item_type.from_json(frame_json, objects, frames_count, key_id_map)
                frames.append(frame)
            except Exception as e:
                if skip_corrupted:
                    logger.warning(f"Skipping corrupted frame: {e}", exc_info=True)
                    continue
                else:
                    raise e
        return cls(frames)

    def __str__(self):
        return "Frames:\n" + super(FrameCollection, self).__str__()

    @property
    def figures(self) -> List[VideoFigure]:
        """
        Get figures from all frames in collection.

        :returns: List of figures from all frames in collection
        :rtype: List[:class:`~supervisely.video_annotation.video_figure.VideoFigure`]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                fr_index_1 = 7
                geometry = sly.Rectangle(0, 0, 100, 100)
                obj_class_car = sly.ObjClass('car', sly.Rectangle)
                video_object_car = sly.VideoObject(obj_class_car)
                video_figure_car = sly.VideoFigure(video_object_car, geometry, fr_index_1)
                frame_1 = sly.Frame(fr_index_1, figures=[video_figure_car])

                fr_index_2 = 10
                geometry = sly.Rectangle(0, 0, 500, 600)
                obj_class_bus = sly.ObjClass('bus', sly.Rectangle)
                video_object_bus = sly.VideoObject(obj_class_bus)
                video_figure_bus = sly.VideoFigure(video_object_bus, geometry, fr_index_2)
                frame_2 = sly.Frame(fr_index_2, figures=[video_figure_bus])

                fr_collection = sly.FrameCollection([frame_1, frame_2])
                figures = fr_collection.figures
        """
        figures_array = []
        for frame in self:
            figures_array.extend(frame.figures)
        return figures_array

    def get_figures_and_keys(self, key_id_map: KeyIdMap) -> Tuple[List[Dict], List[str]]:
        """
        Get figures from all frames in collection in json format, keys from all figures in frames in collection.

        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`
        :returns: Figures from all frames in collection in json format, keys from all figures in frames in collection
        :rtype: Tuple[List[dict], List[str]]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                from supervisely.video_annotation.key_id_map import KeyIdMap

                key_id_map = KeyIdMap()

                fr_index_1 = 7
                geometry = sly.Rectangle(0, 0, 100, 100)
                obj_class_car = sly.ObjClass('car', sly.Rectangle)
                video_object_car = sly.VideoObject(obj_class_car)
                video_figure_car = sly.VideoFigure(video_object_car, geometry, fr_index_1)
                frame_1 = sly.Frame(fr_index_1, figures=[video_figure_car])

                fr_index_2 = 10
                geometry = sly.Rectangle(0, 0, 500, 600)
                obj_class_bus = sly.ObjClass('bus', sly.Rectangle)
                video_object_bus = sly.VideoObject(obj_class_bus)
                video_figure_bus = sly.VideoFigure(video_object_bus, geometry, fr_index_2)
                frame_2 = sly.Frame(fr_index_2, figures=[video_figure_bus])

                fr_collection = sly.FrameCollection([frame_1, frame_2])
                figures, keys = fr_collection.get_figures_and_keys(key_id_map)
                print(keys) # [UUID('0ac041b2-314e-4f6b-9d38-704b341fb383'), UUID('88aa1cb3-b1e3-480f-8ace-6346c9a9daba')]

                print(figures)
                # Output: [
                #     {
                #         "key": "a8cae05d6b8c4a67b18004130941fdec",
                #         "objectKey": "cc9a9475d360481c9753f8ac3c63f8b7",
                #         "geometryType": "rectangle",
                #         "geometry": {
                #             "points": {
                #                 "exterior": [
                #                     [
                #                         0,
                #                         0
                #                     ],
                #                     [
                #                         100,
                #                         100
                #                     ]
                #                 ],
                #                 "interior": []
                #             }
                #         },
                #         "meta": {
                #             "frame": 7
                #         }
                #     },
                #     {
                #         "key": "6e00287acc4644dfb21d67406534080b",
                #         "objectKey": "cad78d53ffc84e69a28f5f8941be9021",
                #         "geometryType": "rectangle",
                #         "geometry": {
                #             "points": {
                #                 "exterior": [
                #                     [
                #                         0,
                #                         0
                #                     ],
                #                     [
                #                         600,
                #                         500
                #                     ]
                #                 ],
                #                 "interior": []
                #             }
                #         },
                #         "meta": {
                #             "frame": 10
                #         }
                #     }
                # ]
        """
        keys = []
        figures_json = []
        for frame in self:
            for figure in frame.figures:
                keys.append(figure.key())
                figure_json = figure.to_json(key_id_map)
                figure_json[ApiField.META] = figure.get_meta()
                figures_json.append(figure_json)
        return figures_json, keys
