from io import BytesIO
from pathlib import Path
from typing import List, Dict, Union, Optional
import numpy as np
from PIL import Image
from tdw.tdw_utils import TDWUtils
from tdw.output_data import OutputData, Magnebot, Images, CameraMatrices
from tdw.robot_data.robot_dynamic import RobotDynamic
from magnebot.arm import Arm


class MagnebotDynamic(RobotDynamic):
    """
    Dynamic data for the Magnebot.

    With a [`Magnebot` agent](magnebot.md):

    ```python
    from tdw.controller import Controller
    from tdw.tdw_utils import TDWUtils
    from magnebot import Magnebot

    m = Magnebot(robot_id=0, position={"x": 0.5, "y": 0, "z": -1})
    c = Controller()
    c.add_ons.append(m)
    c.communicate(TDWUtils.create_empty_room(12, 12))
    print(m.dynamic.transform.position)
    c.communicate({"$type": "terminate"})
    ```

    With a single-agent [`MagnebotController`](magnebot_controller.md):

    ```python
    from magnebot import MagnebotController

    m = MagnebotController()
    m.init_scene()
    print(m.magnebot.dynamic.transform.position)
    m.end()
    ```
    """

    def __init__(self, robot_id: int, resp: List[bytes], body_parts: List[int], frame_count: int, previous=None):
        super().__init__(robot_id=robot_id, resp=resp, body_parts=body_parts, previous=previous)

        """:field
        A dictionary of object IDs currently held by the Magnebot. Key = The arm. Value = a numpy array of object IDs.
        """
        self.held: Dict[Arm, np.array] = {Arm.left: np.array([]),
                                          Arm.right: np.array([])}
        """:field
        The position of the point at the top of the Magnebot's column.
        """
        self.top: np.array = np.array([0, 0, 0])

        """:field
        The images rendered by the robot as dictionary. Key = the name of the pass. Value = the pass as a numpy array.

        | Pass | Image | Description |
        | --- | --- | --- |
        | `"img"` | ![](images/pass_masks/img_0.jpg) | The rendered image. |
        | `"id"` | ![](images/pass_masks/id_0.png) | The object color segmentation pass. See `Magnebot.segmentation_color_to_id` and `Magnebot.objects_static` to map segmentation colors to object IDs. |
        | `"depth"` | ![](images/pass_masks/depth_0.png) | The depth values per pixel as a numpy array. Depth values are encoded into the RGB image; see `SceneState.get_depth_values()`. Use the camera matrices to interpret this data. |

        """
        self.images: Dict[str, np.array] = dict()
        """:field
        The [camera projection matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array.
        """
        self.projection_matrix: Optional[np.array] = None
        """:field
        The [camera matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array.
        """
        self.camera_matrix: Optional[np.array] = None
        """:field
        The current frame count. This is used for image filenames.
        """
        self.frame_count: int = frame_count
        # File extensions per pass.
        self.__image_extensions: Dict[str, str] = dict()

        got_magnebot_images = False
        avatar_id = str(robot_id)
        for i in range(0, len(resp) - 1):
            r_id = OutputData.get_data_type_id(resp[i])
            # Get the images captured by the avatar's camera.
            if r_id == "imag":
                images = Images(resp[i])
                # Get this robot's avatar and save the images.
                if images.get_avatar_id() == avatar_id:
                    got_magnebot_images = True
                    for j in range(images.get_num_passes()):
                        image_data = images.get_image(j)
                        pass_mask = images.get_pass_mask(j)
                        if pass_mask == "_depth":
                            image_data = TDWUtils.get_shaped_depth_pass(images=images, index=j)
                        # Remove the underscore from the pass mask such as: _img -> img
                        pass_name = pass_mask[1:]
                        # Save the image data.
                        self.images[pass_name] = image_data
                        # Record the file extension.
                        self.__image_extensions[pass_name] = images.get_extension(j)
            # Get the camera matrices for the avatar's camera.
            elif r_id == "cama":
                camera_matrices = CameraMatrices(resp[i])
                if camera_matrices.get_avatar_id() == avatar_id:
                    self.projection_matrix = camera_matrices.get_projection_matrix()
                    self.camera_matrix = camera_matrices.get_camera_matrix()
            # Get data for this Magnebot.
            elif r_id == "magn":
                magnebot = Magnebot(resp[i])
                if magnebot.get_id() == robot_id:
                    self.held[Arm.left] = magnebot.get_held_left()
                    self.held[Arm.right] = magnebot.get_held_right()
                    self.top = np.array(magnebot.get_top())
        # Update the frame count.
        if got_magnebot_images:
            self.frame_count += 1

    def save_images(self, output_directory: Union[str, Path]) -> None:
        """
        Save the ID pass (segmentation colors) and the depth pass to disk.
        Images will be named: `[frame_number]_[pass_name].[extension]`
        For example, the depth pass on the first frame will be named: `00000000_depth.png`

        The `img` pass is either a .jpg. The `id` and `depth` passes are .png files.

        :param output_directory: The directory that the images will be saved to.
        """

        if isinstance(output_directory, str):
            output_directory = Path(output_directory)
        if not output_directory.exists():
            output_directory.mkdir(parents=True)
        # The prefix is a zero-padded integer to ensure sequential images.
        prefix = TDWUtils.zero_padding(self.frame_count, 8)
        # Save each image.
        for pass_name in self.images:
            if self.images[pass_name] is None:
                continue
            # Get the filename, such as: `00000000_img.png`
            p = output_directory.joinpath(f"{prefix}_{pass_name}.{self.__image_extensions[pass_name]}")
            if pass_name == "depth":
                Image.fromarray(self.images[pass_name]).save(str(p.resolve()))
            else:
                with p.open("wb") as f:
                    f.write(self.images[pass_name])

    def get_pil_images(self) -> Dict[str, Image.Image]:
        """
        Convert each image pass from the robot camera to PIL images.

        :return: A dictionary of PIL images. Key = the pass name (img, id, depth); Value = The PIL image (can be None)
        """

        images = dict()
        for pass_name in self.images:
            if pass_name == "depth":
                images[pass_name] = Image.fromarray(self.images[pass_name])
            else:
                images[pass_name] = Image.open(BytesIO(self.images[pass_name]))
        return images

    def get_depth_values(self) -> np.array:
        """
        Convert the depth pass to depth values. Can be None if there is no depth image data.

        :return: A decoded depth pass as a numpy array of floats.
        """

        if "depth" in self.images:
            return TDWUtils.get_depth_values(self.images["depth"])
        else:
            return None

    def get_point_cloud(self) -> np.array:
        """
        Returns a point cloud from the depth pass. Can be None if there is no depth image data.

        :return: A decoded depth pass as a numpy array of floats.
        """

        if "depth" in self.images:
            return TDWUtils.get_point_cloud(depth=TDWUtils.get_depth_values(self.images["depth"]),
                                            camera_matrix=self.camera_matrix, far_plane=100, near_plane=1)
        else:
            return None

from io import BytesIO
from pathlib import Path
from typing import List, Dict, Union, Optional
import numpy as np
from PIL import Image
from tdw.tdw_utils import TDWUtils
from tdw.output_data import OutputData, Magnebot, Images, CameraMatrices
from tdw.robot_data.robot_dynamic import RobotDynamic
from magnebot.arm import Arm


class MagnebotDynamic(RobotDynamic):
    """
    Dynamic data for the Magnebot.

    With a [`Magnebot` agent](magnebot.md):

    ```python
    from tdw.controller import Controller
    from tdw.tdw_utils import TDWUtils
    from magnebot import Magnebot

    m = Magnebot(robot_id=0, position={"x": 0.5, "y": 0, "z": -1})
    c = Controller()
    c.add_ons.append(m)
    c.communicate(TDWUtils.create_empty_room(12, 12))
    print(m.dynamic.transform.position)
    c.communicate({"$type": "terminate"})
    ```

    With a single-agent [`MagnebotController`](magnebot_controller.md):

    ```python
    from magnebot import MagnebotController

    m = MagnebotController()
    m.init_scene()
    print(m.magnebot.dynamic.transform.position)
    m.end()
    ```
    """

    def __init__(self, robot_id: int, resp: List[bytes], body_parts: List[int], frame_count: int, previous=None):
        super().__init__(robot_id=robot_id, resp=resp, body_parts=body_parts, previous=previous)

        """:field
        A dictionary of object IDs currently held by the Magnebot. Key = The arm. Value = a numpy array of object IDs.
        """
        self.held: Dict[Arm, np.array] = {Arm.left: np.array([]),
                                          Arm.right: np.array([])}
        """:field
        The position of the point at the top of the Magnebot's column.
        """
        self.top: np.array = np.array([0, 0, 0])

        """:field
        The images rendered by the robot as dictionary. Key = the name of the pass. Value = the pass as a numpy array.

        | Pass | Image | Description |
        | --- | --- | --- |
        | `"img"` | ![](images/pass_masks/img_0.jpg) | The rendered image. |
        | `"id"` | ![](images/pass_masks/id_0.png) | The object color segmentation pass. See `Magnebot.segmentation_color_to_id` and `Magnebot.objects_static` to map segmentation colors to object IDs. |
        | `"depth"` | ![](images/pass_masks/depth_0.png) | The depth values per pixel as a numpy array. Depth values are encoded into the RGB image; see `SceneState.get_depth_values()`. Use the camera matrices to interpret this data. |

        """
        self.images: Dict[str, np.array] = dict()
        """:field
        The [camera projection matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array.
        """
        self.projection_matrix: Optional[np.array] = None
        """:field
        The [camera matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array.
        """
        self.camera_matrix: Optional[np.array] = None
        """:field
        The current frame count. This is used for image filenames.
        """
        self.frame_count: int = frame_count
        # File extensions per pass.
        self.__image_extensions: Dict[str, str] = dict()

        got_magnebot_images = False
        avatar_id = str(robot_id)
        for i in range(0, len(resp) - 1):
            r_id = OutputData.get_data_type_id(resp[i])
            # Get the images captured by the avatar's camera.
            if r_id == "imag":
                images = Images(resp[i])
                # Get this robot's avatar and save the images.
                if images.get_avatar_id() == avatar_id:
                    got_magnebot_images = True
                    for j in range(images.get_num_passes()):
                        image_data = images.get_image(j)
                        pass_mask = images.get_pass_mask(j)
                        if pass_mask == "_depth":
                            image_data = TDWUtils.get_shaped_depth_pass(images=images, index=j)
                        # Remove the underscore from the pass mask such as: _img -> img
                        pass_name = pass_mask[1:]
                        # Save the image data.
                        self.images[pass_name] = image_data
                        # Record the file extension.
                        self.__image_extensions[pass_name] = images.get_extension(j)
            # Get the camera matrices for the avatar's camera.
            elif r_id == "cama":
                camera_matrices = CameraMatrices(resp[i])
                if camera_matrices.get_avatar_id() == avatar_id:
                    self.projection_matrix = camera_matrices.get_projection_matrix()
                    self.camera_matrix = camera_matrices.get_camera_matrix()
            # Get data for this Magnebot.
            elif r_id == "magn":
                magnebot = Magnebot(resp[i])
                if magnebot.get_id() == robot_id:
                    self.held[Arm.left] = magnebot.get_held_left()
                    self.held[Arm.right] = magnebot.get_held_right()
                    self.top = np.array(magnebot.get_top())
        # Update the frame count.
        if got_magnebot_images:
            self.frame_count += 1

    def save_images(self, output_directory: Union[str, Path]) -> None:
        """
        Save the ID pass (segmentation colors) and the depth pass to disk.
        Images will be named: `[frame_number]_[pass_name].[extension]`
        For example, the depth pass on the first frame will be named: `00000000_depth.png`

        The `img` pass is either a .jpg. The `id` and `depth` passes are .png files.

        :param output_directory: The directory that the images will be saved to.
        """

        if isinstance(output_directory, str):
            output_directory = Path(output_directory)
        if not output_directory.exists():
            output_directory.mkdir(parents=True)
        # The prefix is a zero-padded integer to ensure sequential images.
        prefix = TDWUtils.zero_padding(self.frame_count, 8)
        # Save each image.
        for pass_name in self.images:
            if self.images[pass_name] is None:
                continue
            # Get the filename, such as: `00000000_img.png`
            p = output_directory.joinpath(f"{prefix}_{pass_name}.{self.__image_extensions[pass_name]}")
            if pass_name == "depth":
                Image.fromarray(self.images[pass_name]).save(str(p.resolve()))
            else:
                with p.open("wb") as f:
                    f.write(self.images[pass_name])

    def get_pil_images(self) -> Dict[str, Image.Image]:
        """
        Convert each image pass from the robot camera to PIL images.

        :return: A dictionary of PIL images. Key = the pass name (img, id, depth); Value = The PIL image (can be None)
        """

        images = dict()
        for pass_name in self.images:
            if pass_name == "depth":
                images[pass_name] = Image.fromarray(self.images[pass_name])
            else:
                images[pass_name] = Image.open(BytesIO(self.images[pass_name]))
        return images

    def get_depth_values(self) -> np.array:
        """
        Convert the depth pass to depth values. Can be None if there is no depth image data.

        :return: A decoded depth pass as a numpy array of floats.
        """

        if "depth" in self.images:
            return TDWUtils.get_depth_values(self.images["depth"])
        else:
            return None

    def get_point_cloud(self) -> np.array:
        """
        Returns a point cloud from the depth pass. Can be None if there is no depth image data.

        :return: A decoded depth pass as a numpy array of floats.
        """

        if "depth" in self.images:
            return TDWUtils.get_point_cloud(depth=TDWUtils.get_depth_values(self.images["depth"]),
                                            camera_matrix=self.camera_matrix, far_plane=100, near_plane=1)
        else:
            return None

Dynamic data for the Magnebot.      With a [`Magnebot` agent](magnebot.md):      ```python     from tdw.controller import Controller     from tdw.tdw_utils import TDWUtils     from magnebot import Magnebot      m = Magnebot(robot_id=0, position={"x": 0.5, "y": 0, "z": -1})     c = Controller()     c.add_ons.append(m)     c.communicate(TDWUtils.create_empty_room(12, 12))     print(m.dynamic.transform.position)     c.communicate({"$type": "terminate"})     ```      With a single-agent [`MagnebotController`](magnebot_controller.md):      ```python     from magnebot import MagnebotController      m = MagnebotController()     m.init_scene()     print(m.magnebot.dynamic.transform.position)     m.end()     ``` :field         A dictionary of object IDs currently held by the Magnebot. Key = The arm. Value = a numpy array of object IDs. :field         The position of the point at the top of the Magnebot's column. :field         The images rendered by the robot as dictionary. Key = the name of the pass. Value = the pass as a numpy array.          | Pass | Image | Description |         | --- | --- | --- |         | `"img"` | ![](images/pass_masks/img_0.jpg) | The rendered image. |         | `"id"` | ![](images/pass_masks/id_0.png) | The object color segmentation pass. See `Magnebot.segmentation_color_to_id` and `Magnebot.objects_static` to map segmentation colors to object IDs. |         | `"depth"` | ![](images/pass_masks/depth_0.png) | The depth values per pixel as a numpy array. Depth values are encoded into the RGB image; see `SceneState.get_depth_values()`. Use the camera matrices to interpret this data. | :field         The [camera projection matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array. :field         The [camera matrix](https://github.com/threedworld-mit/tdw/blob/master/Documentation/api/output_data.md#cameramatrices) of the Magnebot's camera as a numpy array. :field         The current frame count. This is used for image filenames. # File extensions per pass. # Get the images captured by the avatar's camera. # Get this robot's avatar and save the images. # Remove the underscore from the pass mask such as: _img -> img # Save the image data. # Record the file extension. # Get the camera matrices for the avatar's camera. # Get data for this Magnebot. # Update the frame count. Save the ID pass (segmentation colors) and the depth pass to disk.         Images will be named: `[frame_number]_[pass_name].[extension]`         For example, the depth pass on the first frame will be named: `00000000_depth.png`          The `img` pass is either a .jpg. The `id` and `depth` passes are .png files.          :param output_directory: The directory that the images will be saved to. # The prefix is a zero-padded integer to ensure sequential images. # Save each image. # Get the filename, such as: `00000000_img.png` Convert each image pass from the robot camera to PIL images.          :return: A dictionary of PIL images. Key = the pass name (img, id, depth); Value = The PIL image (can be None) Convert the depth pass to depth values. Can be None if there is no depth image data.          :return: A decoded depth pass as a numpy array of floats. Returns a point cloud from the depth pass. Can be None if there is no depth image data.          :return: A decoded depth pass as a numpy array of floats.