from __future__ import annotations

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

from supervisely.pointcloud_annotation.pointcloud_episode_frame import (
    PointcloudEpisodeFrame,
)
from supervisely.pointcloud_annotation.pointcloud_episode_object_collection import (
    PointcloudEpisodeObjectCollection,
)
from supervisely.pointcloud_annotation.pointcloud_figure import PointcloudFigure
from supervisely.video_annotation.frame_collection import FrameCollection
from supervisely.video_annotation.key_id_map import KeyIdMap


class PointcloudEpisodeFrameCollection(FrameCollection):
    """
    Collection with :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame.PointcloudEpisodeFrame` instances. :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame_collection.PointcloudEpisodeFrameCollection` object is immutable.

    :Usage Example:

        .. code-block:: python

            import supervisely as sly
            from supervisely.geometry.cuboid_3d import Cuboid3d, Vector3d
            from supervisely.pointcloud_annotation.pointcloud_episode_object_collection import PointcloudEpisodeObjectCollection

            # Create pointcloud object
            obj_class_car = sly.ObjClass('car', Cuboid3d)
            pointcloud_obj_car = sly.PointcloudEpisodeObject(obj_class_car)
            objects = PointcloudEpisodeObjectCollection([pointcloud_obj_car])

            # Create two figures
            frame_index_1 = 7
            position_1, rotation_1, dimension_1 = Vector3d(-3.4, 28.9, -0.7), Vector3d(0., 0, -0.03), Vector3d(1.8, 3.9, 1.6)
            cuboid_1 = Cuboid3d(position_1, rotation_1, dimension_1)

            frame_index_2 = 10
            position_2, rotation_2, dimension_2 = Vector3d(-3.4, 28.9, -0.7), Vector3d(0., 0, -0.03), Vector3d(1.8, 3.9, 1.6)
            cuboid_2 = Cuboid3d(position_2, rotation_2, dimension_2)

            figure_1 = sly.PointcloudFigure(pointcloud_obj_car, cuboid_1, frame_index=frame_index_1)
            figure_2 = sly.PointcloudFigure(pointcloud_obj_car, cuboid_2, frame_index=frame_index_2)

            # Create two frames for collection
            frame_1 = sly.PointcloudEpisodeFrame(frame_index_1, figures=[figure_1])
            frame_2 = sly.PointcloudEpisodeFrame(frame_index_2, figures=[figure_2])

            # Create PointcloudEpisodeFrameCollection
            pcd_episodes_fr_collection = sly.PointcloudEpisodeFrameCollection([frame_1, frame_2])

            print(pcd_episodes_fr_collection.to_json())
            # Output: [
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "c9fb727a9b53432fa0316d0a5b6043bc",
            #             "objectKey": "0fc681681b1f4b12909ccf685c53b43e"
            #         }
            #         ],
            #         "index": 7
            #     },
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "c3df69c90fa14cf284906ebad1c360ae",
            #             "objectKey": "0fc681681b1f4b12909ccf685c53b43e"
            #         }
            #         ],
            #         "index": 10
            #     }
            # ]

            # Add item to PointcloudEpisodeFrameCollection
            frame_index_3 = 13
            frame_3 = sly.PointcloudEpisodeFrame(frame_index_3)
            # Remember that PointcloudEpisodeFrameCollection is immutable, and we need to assign new instance of PointcloudEpisodeFrameCollection to a new variable
            new_pcd_episodes_fr_collection = pcd_episodes_fr_collection.add(frame_3)
            print(new_pcd_episodes_fr_collection.to_json())
            # Output: [
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "d16c2bf646664007bd22bcf2996710ed",
            #             "objectKey": "e09431246cc24231a9d76b9ba55ce4e7"
            #         }
            #         ],
            #         "index": 7
            #     },
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "9f915e2d4342408ab0d4517f485f9950",
            #             "objectKey": "e09431246cc24231a9d76b9ba55ce4e7"
            #         }
            #         ],
            #         "index": 10
            #     },
            #     { "figures": [], "index": 13 }
            # ]

            # You can also add multiple items to collection
            frame_3 = sly.PointcloudEpisodeFrame(12)
            frame_4 = sly.PointcloudEpisodeFrame(15)
            # Remember that PointcloudEpisodeFrameCollection is immutable, and we need to assign new instance of PointcloudEpisodeFrameCollection to a new variable
            new_pcd_episodes_fr_collection = pcd_episodes_fr_collection.add_items([frame_3, frame_4])
            print(new_pcd_episodes_fr_collection.to_json())
            # Output: [
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "df14bb10e76c4b4093a137f4ea01e3ba",
            #             "objectKey": "a83a21ef7b734bb9949ddf132b42d1f0"
            #         }
            #         ],
            #         "index": 7
            #     },
            #     {
            #         "figures": [
            #         {
            #             "geometry": {
            #             "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
            #             "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
            #             "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
            #             },
            #             "geometryType": "cuboid_3d",
            #             "key": "6f86a1859da34415a37896ec2c972b47",
            #             "objectKey": "a83a21ef7b734bb9949ddf132b42d1f0"
            #         }
            #         ],
            #         "index": 10
            #     },
            #     { "figures": [], "index": 12 },
            #     { "figures": [], "index": 15 }
            # ]

            # Has key, checks if given key exist in point cloud episodes frame collection
            pcd_episodes_fr_collection.has_key(10) # True
            pcd_episodes_fr_collection.has_key(22) # False

            # Find intersection of given list of instances with collection items
            frame_1 = sly.PointcloudEpisodeFrame(1)
            frame_2 = sly.PointcloudEpisodeFrame(2)
            pcd_episodes_fr_collection = sly.PointcloudEpisodeFrameCollection([frame_1, frame_2])

            frame_3 = sly.PointcloudEpisodeFrame(3)

            frames_intersections = pcd_episodes_fr_collection.intersection([frame_3])
            print(frames_intersections.to_json())
            # Output: []

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

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

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

            # Merge collection and given list of PointcloudEpisodeFrameCollection
            frame_3 = sly.PointcloudEpisodeFrame(3)
            frame_4 = sly.PointcloudEpisodeFrame(4)
            over_collection = sly.PointcloudEpisodeFrameCollection([frame_3, frame_4])
            merged_collection = pcd_episodes_fr_collection.merge(over_collection)
            print(merged_collection.to_json())
            # Output: [
            #     { "index": 3, "figures": [] },
            #     { "index": 4, "figures": [] },
            #     { "index": 1, "figures": [] },
            #     { "index": 2, "figures": [] }
            # ]

    """

    item_type = PointcloudEpisodeFrame

    def get(self, key: str, default: Optional[Any] = None) -> PointcloudEpisodeFrame:
        """
        Get a PointcloudEpisodeFrame by its key and set default value if it does not exist.

        :param key: Key of the :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame.PointcloudEpisodeFrame`.
        :type key: str
        :param default: Default value to return if the key is not found (default: None).
        :type default: Optional[Any]
        :returns: :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame.PointcloudEpisodeFrame` object.
        :rtype: :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame.PointcloudEpisodeFrame`

        :Usage Example:

            .. code-block:: python

                frame_1 = sly.PointcloudEpisodeFrame(1)
                frame_2 = sly.PointcloudEpisodeFrame(2)

                pcd_episodes_fr_collection = sly.PointcloudEpisodeFrameCollection([frame_1, frame_2])
                print(pcd_episodes_fr_collection.get(2).to_json())
                # Output: {'index': 2, 'figures': []}
        """

        return super().get(key, default)

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

    @classmethod
    def from_json(
        cls,
        data: List[Dict],
        objects: PointcloudEpisodeObjectCollection,
        frames_count: Optional[int] = None,
        key_id_map: Optional[KeyIdMap] = None,
    ) -> PointcloudEpisodeFrameCollection:
        """
        Convert a list of json dicts to PointcloudEpisodeFrameCollection. 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: Pointcloud episode objects collection.
        :type objects: :class:`~supervisely.pointcloud_annotation.pointcloud_episode_object_collection.PointcloudEpisodeObjectCollection`
        :param frames_count: Number of frames in pointcloud episodes.
        :type frames_count: int
        :param key_id_map: Key ID map.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`
        :returns: Pointcloud episode frames collection.
        :rtype: :class:`~supervisely.pointcloud_annotation.pointcloud_episode_frame_collection.PointcloudEpisodeFrameCollection`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                fr_collection_json = [
                    {
                        "figures": [
                            {
                                "geometry": {
                                    "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
                                    "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
                                    "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
                                },
                                "geometryType": "cuboid_3d",
                                "key": "c9fb727a9b53432fa0316d0a5b6043bc",
                                "objectKey": "0fc681681b1f4b12909ccf685c53b43e"
                            }
                        ],
                        "index": 7
                    },
                    {
                        "figures": [
                            {
                                "geometry": {
                                    "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
                                    "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
                                    "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
                                },
                                "geometryType": "cuboid_3d",
                                "key": "c3df69c90fa14cf284906ebad1c360ae",
                                "objectKey": "0fc681681b1f4b12909ccf685c53b43e"
                            }
                        ],
                        "index": 10
                    }
                ]

                objects = []
                fr_collection = sly.PointcloudEpisodeFrameCollection.from_json(fr_collection_json, objects)
        """

        return super().from_json(data, objects, frames_count=frames_count, key_id_map=key_id_map)

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

        :returns: List of pointcloud figures from all frames in collection.
        :rtype: List[:class:`~supervisely.pointcloud_annotation.pointcloud_figure.PointcloudFigure`]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                from supervisely.geometry.cuboid_3d import Cuboid3d, Vector3d
                from supervisely.pointcloud_annotation.pointcloud_episode_object_collection import PointcloudEpisodeObjectCollection

                # Create pointcloud object
                obj_class_car = sly.ObjClass('car', Cuboid3d)
                pointcloud_obj_car = sly.PointcloudEpisodeObject(obj_class_car)
                objects = PointcloudEpisodeObjectCollection([pointcloud_obj_car])

                # Create two figures
                frame_index_1 = 7
                position_1, rotation_1, dimension_1 = Vector3d(-3.4, 28.9, -0.7), Vector3d(0., 0, -0.03), Vector3d(1.8, 3.9, 1.6)
                cuboid_1 = Cuboid3d(position_1, rotation_1, dimension_1)

                frame_index_2 = 10
                position_2, rotation_2, dimension_2 = Vector3d(-3.4, 28.9, -0.7), Vector3d(0., 0, -0.03), Vector3d(1.8, 3.9, 1.6)
                cuboid_2 = Cuboid3d(position_2, rotation_2, dimension_2)

                figure_1 = sly.PointcloudFigure(pointcloud_obj_car, cuboid_1, frame_index=frame_index_1)
                figure_1 = sly.PointcloudFigure(pointcloud_obj_car, cuboid_2, frame_index=frame_index_2)

                # Create two frames for collection
                frame_1 = sly.PointcloudEpisodeFrame(frame_index_1, figures=[figure_1])
                frame_2 = sly.PointcloudEpisodeFrame(frame_index_2, figures=[figure_2])

                # Create PointcloudEpisodeFrameCollection
                pcd_episodes_fr_collection = sly.PointcloudEpisodeFrameCollection([frame_1, frame_2])

                print([figure.to_json() for figure in pcd_episodes_fr_collection.figures])
                # Output: [
                #     {
                #         "geometry": {
                #         "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
                #         "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
                #         "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
                #         },
                #         "geometryType": "cuboid_3d",
                #         "key": "e35386df7c994485ba7186e1a8c67361",
                #         "objectKey": "c43ac88f2cb74852817390ca3b749ffd"
                #     },
                #     {
                #         "geometry": {
                #         "dimensions": { "x": 1.8, "y": 3.9, "z": 1.6 },
                #         "position": { "x": -3.4, "y": 28.9, "z": -0.7 },
                #         "rotation": { "x": 0.0, "y": 0, "z": -0.03 }
                #         },
                #         "geometryType": "cuboid_3d",
                #         "key": "1b09a434c20f47faab3f0d7efdb4cbbf",
                #         "objectKey": "c43ac88f2cb74852817390ca3b749ffd"
                #     }
                # ]
        """

        return super().figures
