# coding: utf-8

# isort: skip_file

from __future__ import annotations
from copy import deepcopy
from re import L
from typing import List, Union
import uuid
from supervisely.volume_annotation.volume_figure import VolumeFigure

from supervisely.project.project_meta import ProjectMeta
from supervisely._utils import take_with_default
from supervisely.io.fs import file_exists
from supervisely.video_annotation.key_id_map import KeyIdMap
from supervisely.volume_annotation.slice import Slice
from supervisely.volume_annotation.volume_tag_collection import VolumeTagCollection
from supervisely.volume_annotation.volume_object_collection import VolumeObjectCollection
from supervisely.volume_annotation.volume_object import VolumeObject
from supervisely.volume_annotation.volume_tag import VolumeTag
from supervisely.geometry.mask_3d import Mask3D
from supervisely.geometry.any_geometry import AnyGeometry
from supervisely.volume_annotation.plane import Plane
from supervisely.volume_annotation.constants import (
    NAME,
    TAGS,
    OBJECTS,
    KEY,
    VOLUME_ID,
    VOLUME_META,
    PLANES,
    SPATIAL_FIGURES,
)
from supervisely.io.fs import get_file_name
from supervisely.io.json import dump_json_file


class VolumeAnnotation:
    """Annotation for a 3D volume: objects, tags, planes, spatial figures. Immutable."""

    def __init__(
        self,
        volume_meta,
        objects=None,
        plane_sagittal=None,
        plane_coronal=None,
        plane_axial=None,
        tags=None,
        spatial_figures=None,
        key=None,
    ):
        """
        Annotation for a 3D volume.

        :param volume_meta: Volume metadata (dimensions, spacing, origin, etc.).
        :type volume_meta: dict
        :param objects: Collection of VolumeObjects.
        :type objects: :class:`~supervisely.volume_annotation.volume_object_collection.VolumeObjectCollection`, optional
        :param plane_sagittal: Sagittal plane. Default Plane created if not provided.
        :type plane_sagittal: :class:`~supervisely.volume_annotation.plane.Plane`, optional
        :param plane_coronal: Coronal plane. Default Plane created if not provided.
        :type plane_coronal: :class:`~supervisely.volume_annotation.plane.Plane`, optional
        :param plane_axial: Axial plane. Default Plane created if not provided.
        :type plane_axial: :class:`~supervisely.volume_annotation.plane.Plane`, optional
        :param tags: Collection of VolumeTags.
        :type tags: :class:`~supervisely.volume_annotation.volume_tag_collection.VolumeTagCollection`, optional
        :param spatial_figures: List of VolumeFigures (e.g. Mask3D).
        :type spatial_figures: List[:class:`~supervisely.volume_annotation.volume_figure.VolumeFigure`], optional
        :param key: UUID key. Auto-generated if not provided.
        :type key: uuid.UUID, optional

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)

                obj_class = sly.ObjClass('brain', sly.Rectangle)
                volume_obj = sly.VolumeObject(obj_class)
                vol_tag = sly.VolumeTag(sly.TagMeta('brain_tag', sly.TagValueType.ANY_STRING), value='human')
                volume_ann = sly.VolumeAnnotation(
                    volume_meta,
                    sly.VolumeObjectCollection([volume_obj]),
                    tags=sly.VolumeTagCollection([vol_tag]),
                )
        """
        self._volume_meta = volume_meta
        self._tags = take_with_default(tags, VolumeTagCollection())
        self._objects = take_with_default(objects, VolumeObjectCollection())
        self._key = take_with_default(key, uuid.uuid4())

        self._plane_sagittal = take_with_default(
            plane_sagittal,
            Plane(Plane.SAGITTAL, volume_meta=volume_meta),
        )
        self._plane_coronal = take_with_default(
            plane_coronal,
            Plane(Plane.CORONAL, volume_meta=volume_meta),
        )
        self._plane_axial = take_with_default(
            plane_axial,
            Plane(Plane.AXIAL, volume_meta=volume_meta),
        )

        self._spatial_figures = take_with_default(spatial_figures, [])
        self.validate_figures_bounds()

    @property
    def volume_meta(self) -> dict:
        """
        Volume meta data.

        :returns: Sagittal plane of the volume.
        :rtype: dict

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                volume_meta = vol_ann.volume_meta
        """

        return deepcopy(self._volume_meta)

    @property
    def plane_sagittal(self) -> Plane:
        """
        Sagital plane of the volume.

        :returns: Sagittal plane of the volume.
        :rtype: :class:`~supervisely.volume_annotation.plane.Plane`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                plane_sagittal = vol_ann.plane_sagittal
        """

        return self._plane_sagittal

    @property
    def plane_coronal(self) -> Plane:
        """
        Coronal plane of the volume.

        :returns: Coronal plane of the volume.
        :rtype: :class:`~supervisely.volume_annotation.plane.Plane`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                plane_coronal = vol_ann.plane_coronal
        """

        return self._plane_coronal

    @property
    def plane_axial(self) -> Plane:
        """
        Axial plane of the volume.

        :returns: Axial plane of the volume.
        :rtype: :class:`~supervisely.volume_annotation.plane.Plane`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                plane_axial = vol_ann.plane_axial
        """

        return self._plane_axial

    @property
    def objects(self) -> VolumeObjectCollection:
        """
        VolumeAnnotation objects.

        :returns: VolumeObjectCollection object
        :rtype: :class:`~supervisely.volume_annotation.volume_object_collection.VolumeObjectCollection`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/Users/Downloads/volumes/Demo volumes_ds1_CTChest.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)

                # VolumeObjectCollection
                obj_class_heart = sly.ObjClass('heart', sly.Rectangle)
                volume_obj_heart = sly.VolumeObject(obj_class_heart)
                objects = sly.VolumeObjectCollection([volume_obj_heart])
                volume_ann = sly.VolumeAnnotation(volume_meta, objects)

                print(volume_ann.objects.to_json())
                # Output: [
                #     {
                #         "key": "2b5d70baa5a74d06a525b950b5f2b756",
                #         "classTitle": "heart",
                #         "tags": []
                #     }
                # ]
        """

        return self._objects

    @property
    def tags(self) -> VolumeTagCollection:
        """
        VolumeTag objects.

        :returns: VolumeTagCollection object
        :rtype: :class:`~supervisely.volume_annotation.volume_tag_collection.VolumeTagCollection`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                tags = vol_ann.tags
        """

        return self._tags

    @property
    def spatial_figures(self) -> List[VolumeFigure]:
        """
        Get a list of spatial figures.

        :returns: List of spatial figures from :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation` object.
        :rtype: List[:class:`~supervisely.volume_annotation.volume_figure.VolumeFigure`]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                spatial_figures = vol_ann.spatial_figures
        """

        return self._spatial_figures

    @property
    def figures(self) -> List[VolumeFigure]:
        """
        VolumeFigure objects.

        :returns: List of VolumeFigure objects from VolumeAnnotation object.
        :rtype: List[:class:`~supervisely.volume_annotation.volume_figure.VolumeFigure`]

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                figures = vol_ann.figures
        """

        all_figures = []
        for plane in [self.plane_sagittal, self.plane_coronal, self.plane_axial]:
            all_figures.extend(plane.figures)
        return all_figures

    def key(self) -> str:
        """
        Volume annotation key value.

        :returns: Key value of VolumeAnnotation object.
        :rtype: str

        :Usage Example:

            .. code-block:: python

                import supervisely as sly
                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                vol_ann = sly.VolumeAnnotation(volume_meta)
                key = vol_ann.key()
        """
        return self._key

    def validate_figures_bounds(self):
        """
        Checks if all slices in each plane contains figures.

        :raises :class:`~supervisely.video_annotation.video_figure.OutOfImageBoundsException`: if figure is out of slices images bounds
        :returns: None
        :rtype: None

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                plane_axial = sly.Plane(sly.Plane.AXIAL, frames, volume_meta=volume_meta)
                volume_ann = sly.VolumeAnnotation(volume_meta, objects, plane_axial=plane_axial)
                volume_ann.validate_figures_bounds()
        """

        self.plane_sagittal.validate_figures_bounds()
        self.plane_coronal.validate_figures_bounds()
        self.plane_axial.validate_figures_bounds()

    def is_empty(self) -> bool:
        """
        Check whether volume annotation contains objects or tags, or not.

        :returns: True if volume annotation is empty, False otherwise.
        :rtype: bool

        :Usage exmaple:

         .. code-block:: python

            import supervisely as sly
            path = "/home/admin/work/volumes/vol_01.nrrd"
            volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
            vol_ann = sly.VolumeAnnotation(volume_meta)

            is_empty = vol_ann.is_empty()
        """

        if len(self.objects) == 0 and len(self.tags) == 0:
            return True
        else:
            return False

    def clone(
        self,
        volume_meta=None,
        objects=None,
        plane_sagittal=None,
        plane_coronal=None,
        plane_axial=None,
        tags=None,
        spatial_figures=None,
    ):
        """
        Makes a copy of VolumeAnnotation with new fields, if fields are given, otherwise it will use fields of the original VolumeAnnotation.

        :param volume_meta: Metadata of the volume.
        :type volume_meta: dict
        :param objects: VolumeObjectCollection object.
        :type objects: :class:`~supervisely.volume_annotation.volume_object_collection.VolumeObjectCollection`, optional
        :param plane_sagittal: Sagittal plane of the volume.
        :type plane_sagittal: Plane object, optional
        :param plane_coronal: Coronal plane of the volume.
        :type plane_coronal: Plane object, optional
        :param plane_axial: Axial plane of the volume.
        :type plane_axial: Plane object, optional
        :param tags: VolumeTagCollection object.
        :type tags: :class:`~supervisely.volume_annotation.volume_tag_collection.VolumeTagCollection`, optional
        :param spatial_figures: List of spatial figures associated with the volume.
        :type spatial_figures: List[:class:`~supervisely.volume_annotation.volume_figure.VolumeFigure`], optional

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)

                obj_class_heart = sly.ObjClass('heart', sly.Rectangle)
                video_obj_heart = sly.VolumeObject(obj_class_heart)
                new_objects = sly.VolumeObjectCollection([volume_obj_heart])
                new_volume_ann = volume_ann.clone(objects=new_objects)
        """

        return VolumeAnnotation(
            volume_meta=take_with_default(volume_meta, self.volume_meta),
            objects=take_with_default(objects, self.objects),
            plane_sagittal=take_with_default(plane_sagittal, self.plane_sagittal),
            plane_coronal=take_with_default(plane_coronal, self.plane_coronal),
            plane_axial=take_with_default(plane_axial, self.plane_axial),
            tags=take_with_default(tags, self.tags),
            spatial_figures=take_with_default(spatial_figures, self.spatial_figures),
        )

    @classmethod
    def from_json(
        cls,
        data: dict,
        project_meta: ProjectMeta,
        key_id_map: KeyIdMap = None,
        spatial_geometry_paths: Union[list, dict] = None,
    ):
        """
        Convert a json dict to VolumeAnnotation.

        :param data: Volume annotation in json format as a dict.
        :type data: dict
        :param project_meta: Input ProjectMeta object.
        :type project_meta: :class:`~supervisely.project.project_meta.ProjectMeta`
        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`, optional
        :param spatial_geometry_paths: Optional. Can be either:
            - a list of file paths to spatial geometry files, where each file name should match either the figure's id or the hex value of its key,
            - or a dict mapping figure ids (or keys) to their corresponding geometry file paths.
            Used to load 3D geometry for spatial figures.
        :type spatial_geometry_paths: list or dict, optional
        :returns: VolumeAnnotation object
        :rtype: :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                from supervisely.video_annotation.key_id_map import KeyIdMap

                meta = sly.ProjectMeta()
                key_id_map = KeyIdMap()

                ann_json = {
                    "key": "56107223943346e5900fc256b8dcd7f0",
                    "objects": [],
                    "planes": [
                        { "name": "sagittal", "normal": { "x": 1, "y": 0, "z": 0 }, "slices": [] },
                        { "name": "coronal", "normal": { "x": 0, "y": 1, "z": 0 }, "slices": [] },
                        { "name": "axial", "normal": { "x": 0, "y": 0, "z": 1 }, "slices": [] }
                    ],
                    "spatialFigures": [],
                    "tags": [],
                    "volumeMeta": {
                        "ACS": "RAS",
                        "channelsCount": 1,
                        "dimensionsIJK": { "x": 512, "y": 512, "z": 139 },
                        "directions": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
                        "intensity": { "max": 3071.0, "min": -3024.0 },
                        "origin": [-194.238403081894, -217.5384061336518, -347.7500000000001],
                        "rescaleIntercept": 0,
                        "rescaleSlope": 1,
                        "spacing": [0.7617189884185793, 0.7617189884185793, 2.5],
                        "windowCenter": 23.5,
                        "windowWidth": 6095.0
                    }
                }

                ann = sly.VolumeAnnotation.from_json(ann_json, meta, key_id_map)
        """

        volume_key = uuid.UUID(data[KEY]) if KEY in data else uuid.uuid4()
        if key_id_map is not None:
            key_id_map.add_video(volume_key, data.get(VOLUME_ID, None))

        volume_meta = data[VOLUME_META]

        tags = VolumeTagCollection.from_json(data[TAGS], project_meta.tag_metas, key_id_map)
        objects = VolumeObjectCollection.from_json(data[OBJECTS], project_meta, key_id_map)

        plane_sagittal = None
        plane_coronal = None
        plane_axial = None
        for plane_json in data[PLANES]:
            if plane_json[NAME] == Plane.SAGITTAL:
                plane_sagittal = Plane.from_json(
                    plane_json,
                    Plane.SAGITTAL,
                    objects,
                    volume_meta=volume_meta,
                    key_id_map=key_id_map,
                )
            elif plane_json[NAME] == Plane.CORONAL:
                plane_coronal = Plane.from_json(
                    plane_json,
                    Plane.CORONAL,
                    objects,
                    volume_meta=volume_meta,
                    key_id_map=key_id_map,
                )
            elif plane_json[NAME] == Plane.AXIAL:
                plane_axial = Plane.from_json(
                    plane_json,
                    Plane.AXIAL,
                    objects,
                    volume_meta=volume_meta,
                    key_id_map=key_id_map,
                )
            else:
                raise RuntimeError(f"Unknown plane name {plane_json[NAME]}")

        spatial_figures = []
        for figure_json in data.get(SPATIAL_FIGURES, []):
            figure = VolumeFigure.from_json(
                figure_json,
                objects,
                plane_name=None,
                slice_index=None,
                key_id_map=key_id_map,
            )
            if spatial_geometry_paths:
                figure_id = figure_json["id"]
                if isinstance(spatial_geometry_paths, list):
                    spatial_geometry_paths = {
                        get_file_name(path): path for path in spatial_geometry_paths
                    }
                    geometry_path = spatial_geometry_paths.get(
                        figure.key().hex
                    ) or spatial_geometry_paths.get(str(figure_id))
                elif isinstance(spatial_geometry_paths, dict):
                    geometry_path = spatial_geometry_paths.get(figure_id, None)
                else:
                    raise ValueError(
                        f"spatial_geometry_paths should be either a list or a dict. Got: {type(spatial_geometry_paths)}"
                    )

                if geometry_path is not None:
                    if not file_exists(geometry_path):
                        raise RuntimeError(
                            f"Geometry file {geometry_path} for figure {figure_id} does not exist."
                        )
                    mask3d = Mask3D.create_from_file(geometry_path)
                    mask3d.sly_id = figure_id
                    figure._set_3d_geometry(mask3d)

            spatial_figures.append(figure)

        return cls(
            volume_meta=volume_meta,
            objects=objects,
            plane_sagittal=plane_sagittal,
            plane_coronal=plane_coronal,
            plane_axial=plane_axial,
            tags=tags,
            spatial_figures=spatial_figures,
            key=volume_key,
        )

    def to_json(self, key_id_map: KeyIdMap = None) -> dict:
        """
        Convert the VolumeAnnotation to a json dict.

        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`, optional
        :returns: Volume annotation in json format as a dict.
        :rtype: dict

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                from supervisely.video_annotation.key_id_map import KeyIdMap

                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)

                print(volume_ann.to_json())
                # Output: {
                # {
                #     "key": "56107223943346e5900fc256b8dcd7f0",
                #     "objects": [],
                #     "planes": [
                #         { "name": "sagittal", "normal": { "x": 1, "y": 0, "z": 0 }, "slices": [] },
                #         { "name": "coronal", "normal": { "x": 0, "y": 1, "z": 0 }, "slices": [] },
                #         { "name": "axial", "normal": { "x": 0, "y": 0, "z": 1 }, "slices": [] }
                #     ],
                #     "spatialFigures": [],
                #     "tags": [],
                #     "volumeMeta": {
                #         "ACS": "RAS",
                #         "channelsCount": 1,
                #         "dimensionsIJK": { "x": 512, "y": 512, "z": 139 },
                #         "directions": [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0],
                #         "intensity": { "max": 3071.0, "min": -3024.0 },
                #         "origin": [-194.238403081894, -217.5384061336518, -347.7500000000001],
                #         "rescaleIntercept": 0,
                #         "rescaleSlope": 1,
                #         "spacing": [0.7617189884185793, 0.7617189884185793, 2.5],
                #         "windowCenter": 23.5,
                #         "windowWidth": 6095.0
                #     }
                # }
        """

        res_json = {
            VOLUME_META: self.volume_meta,
            KEY: self.key().hex,
            TAGS: self.tags.to_json(key_id_map),
            OBJECTS: self.objects.to_json(key_id_map),
            PLANES: [
                self.plane_sagittal.to_json(),
                self.plane_coronal.to_json(),
                self.plane_axial.to_json(),
            ],
            SPATIAL_FIGURES: [figure.to_json(key_id_map) for figure in self.spatial_figures],
        }

        if key_id_map is not None:
            volume_id = key_id_map.get_video_id(self.key())
            if volume_id is not None:
                res_json[VOLUME_ID] = volume_id

        return res_json

    def dump_json(self, path: str, key_id_map: KeyIdMap = None) -> None:
        """
        Save the VolumeAnnotation to a json file.

        :param path: Path to the json file.
        :type path: str
        :param key_id_map: KeyIdMap object.
        :type key_id_map: :class:`~supervisely.video_annotation.key_id_map.KeyIdMap`, optional
        :returns: None
        :rtype: None

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/home/admin/work/volumes/vol_01.nrrd"
                volume, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)
                volume_ann.dump_json("/home/admin/work/volumes/vol_01.json")
        """
        simple_dict = self.to_json(key_id_map)
        dump_json_file(simple_dict, path)

    def add_objects(
        self, objects: Union[List[VolumeObject], VolumeObjectCollection]
    ) -> VolumeAnnotation:
        """
        Add new objects to a VolumeAnnotation object.

        :param objects: New volume objects.
        :type objects: List[VolumeObject] or VolumeObjectCollection
        :returns: A :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation` object containing the original and new volume objects.
        :rtype: :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation`

        :Usage Example:

            .. code-block:: python

                import os
                from dotenv import load_dotenv

                import supervisely as sly

                path = "/vol_01.nrrd"
                _, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)
                obj_class_heart = sly.ObjClass('heart', sly.Mask3D)
                volume_obj_heart = sly.VolumeObject(obj_class_heart)
                volume_ann = volume_ann.add_objects([volume_obj_heart])

        """

        sf_figures = []
        for volume_object in objects:
            if volume_object.obj_class.geometry_type in (Mask3D, AnyGeometry):
                if isinstance(volume_object.figure.geometry, Mask3D):
                    sf_figures.append(volume_object.figure)

        collection = self.objects.add_items(objects)
        new_ann = self.clone(objects=collection)
        new_ann.spatial_figures.extend(sf_figures)
        return new_ann

    def remove_objects(self, keys: Union[List[uuid.UUID], uuid.UUID]) -> VolumeAnnotation:
        """
        Remove annotation objects from a VolumeAnnotation object.

        :param keys: List of object keys or single object key.
        :type keys: List[uuid.UUID] or uuid.UUID
        :returns: A VolumeAnnotation object containing the original volume objects without the removed objects.
        :rtype: :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/vol_01.nrrd"
                _, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)
                obj_class_heart = sly.ObjClass('heart', sly.Mask3D)
                volume_obj_heart = sly.VolumeObject(obj_class_heart)
                obj_class_tumor = sly.ObjClass('tumor', sly.Mask3D)
                volume_obj_tumor = sly.VolumeObject(obj_class_tumor)
                volume_ann = volume_ann.add_objects([volume_obj_heart, volume_obj_tumor])
                volume_ann = volume_ann.remove_objects(volume_obj_heart.key())

        """
        if not isinstance(keys, list):
            keys = [keys]

        spatial_figures_to_del = []
        for ann_object in self.objects:
            if ann_object.key() in keys and ann_object.obj_class.geometry_type in (
                Mask3D,
                AnyGeometry,
            ):
                for figure in self.spatial_figures:
                    if figure.parent_object == ann_object:
                        spatial_figures_to_del.append(figure)

        for spatial_figure in spatial_figures_to_del:
            self.spatial_figures.remove(spatial_figure)

        collection = self.objects.remove_items(keys)
        new_ann = self.clone(objects=collection)
        return new_ann

    def add_tags(self, tags: Union[List[VolumeTag], VolumeTagCollection]) -> VolumeAnnotation:
        """
        Add new VolumeTags to a VolumeAnnotation object.

        :param tags: New VolumeTags.
        :type tags: List[:class:`~supervisely.volume_annotation.volume_tag.VolumeTag`] or :class:`~supervisely.volume_annotation.volume_tag_collection.VolumeTagCollection`
        :returns: A :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation` object containing the original and new volume tags.
        :rtype: :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/vol_01.nrrd"
                _, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)
                brain_meta = sly.TagMeta('brain_tag', sly.TagValueType.ANY_STRING)
                vol_tag = sly.VolumeTag(brain_meta, value='human')
                volume_ann = volume_ann.add_tags([vol_tag])

        """
        collection = self.tags.add_items(tags)
        return self.clone(tags=collection)

    def remove_tags(self, keys: Union[List[uuid.UUID], uuid.UUID]) -> VolumeAnnotation:
        """
        Remove VolumeTags from a VolumeAnnotation object.

        :param keys: List of VolumeTag keys or single tag key.
        :type keys: List[uuid.UUID] or uuid.UUID
        :returns: A :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation` object containing the original VolumeTags without the removed VolumeTags.
        :rtype: :class:`~supervisely.volume_annotation.volume_annotation.VolumeAnnotation`

        :Usage Example:

            .. code-block:: python

                import supervisely as sly

                path = "/vol_01.nrrd"
                _, volume_meta = sly.volume.read_nrrd_serie_volume_np(path)
                volume_ann = sly.VolumeAnnotation(volume_meta)
                brain_meta = sly.TagMeta('brain_tag', sly.TagValueType.ANY_STRING)
                vol_tag = sly.VolumeTag(brain_meta, value='human')
                volume_ann = volume_ann.add_tags([vol_tag])
                volume_ann = volume_ann.remove_tags(vol_tag.key())

        """
        if not isinstance(keys, list):
            keys = [keys]
        collection = self.tags.remove_items(keys)
        return self.clone(tags=collection)
