import numpy as np
import pyrfuniverse.attributes as attr
import base64


class CameraAttr(attr.BaseAttr):
    """
    Camera attribute class, which can capture many kinds of screenshot
    of the scene in RFUniverse.
    """

    def parse_message(self, data: dict):
        """
        Parse messages. This function is called by internal function.

        Returns:
            Dict: A dict containing useful information of this class.

            self.data['width']: The width of image.

            self.data['height']: The height of image.

            self.data['fov']: The field of view of camera.

            self.data['rgb']: The bytes of rgb image.

            self.data['normal']: The bytes of normal image.

            self.data['id_map']: The bytes of instance segmentation mask image.

            self.data['depth']: The bytes of depth image.

            self.data['depth_exr']: The bytes of depth image in exr format.

            self.data['amodal_mask']: The bytes of amodal mask image.

            self.data['heat_map']: The bytes of heat map image.

            self.data['2d_bounding_box']: The 2d bouding box of objects in camera (image) coordinate.

            self.data['3d_bounding_box']: The 3d bounding box of objects in world coordinate.
        """
        super().parse_message(data)
        if "rgb" in self.data:
            self.data["rgb"] = base64.b64decode(self.data["rgb"])
        if "normal" in self.data:
            self.data["normal"] = base64.b64decode(self.data["normal"])
        if "id_map" in self.data:
            self.data["id_map"] = base64.b64decode(self.data["id_map"])
        if "depth" in self.data:
            self.data["depth"] = base64.b64decode(self.data["depth"])
        if "depth_exr" in self.data:
            self.data["depth_exr"] = base64.b64decode(self.data["depth_exr"])
        if "amodal_mask" in self.data:
            self.data["amodal_mask"] = base64.b64decode(self.data["amodal_mask"])
        if "heat_map" in self.data:
            self.data["heat_map"] = base64.b64decode(self.data["heat_map"])

    def AlignView(self):
        """
        Make the camera in RFUniverse align the current view in GUI.
        """
        self._send_data("AlignView")

    def GetRGB(
        self,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the camera RGB image.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data("GetRGB", intrinsic_matrix, int(width), int(height), float(fov))

    def GetNormal(
        self,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the normal image in world coordinate.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "GetNormal", intrinsic_matrix, int(width), int(height), float(fov)
        )

    def GetID(
        self,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the instance segmentation mask image. The color for each pixel is computed from object ID, see `pyrfuniverse.utils.rfuniverse_utility.GetColorFromID` for more details.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data("GetID", intrinsic_matrix, int(width), int(height), float(fov))

    def GetDepth(
        self,
        zero_dis: float,
        one_dis: float,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the depth image from camera. Since eacg pixel of depth image returned from this function is 8-bit, user should limit the depth range (`zero_dis` and `one_dis`) for more accurate results.

        Args:
            zero_dis: The minimum distance in calculation.
            one_dis: The maximum distance in calculation.
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "GetDepth",
            float(zero_dis),
            float(one_dis),
            intrinsic_matrix,
            int(width),
            int(height),
            float(fov),
        )

    def GetDepthEXR(
        self,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the depth image from camera. This function returns EXR format image bytes and each pixel is 32-bit.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "GetDepthEXR", intrinsic_matrix, int(width), int(height), float(fov)
        )

    def GetAmodalMask(
        self,
        target_id: int,
        width: int = 512,
        height: int = 512,
        fov: float = 60,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the amodal mask image for target object.

        Args:
            target_id: The target object ID.
            width: Int, the width of image.
            height: Int, the height of image.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "GetAmodalMask",
            int(target_id),
            intrinsic_matrix,
            int(width),
            int(height),
            float(fov),
        )

    def StartHeatMapRecord(self, targets_id: list):
        targets_id = [int(i) for i in targets_id]

        self._send_data("StartHeatMapRecord", targets_id)

    def EndHeatMapRecord(self):
        self._send_data("EndHeatMapRecord")

    def GetHeatMap(
        self,
        width: int = 512,
        height: int = 512,
        radius: int = 50,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the heat map image.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            radius: The radius of heat map.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "GetHeatMap",
            int(radius),
            intrinsic_matrix,
            int(width),
            int(height),
            float(fov),
        )

    def Get2DBBox(
        self,
        width: int = 512,
        height: int = 512,
        fov: float = 60.0,
        intrinsic_matrix: np.ndarray = None,
    ):
        """
        Get the 2d bounding box of objects in current camera view.

        Args:
            width: Int, the width of image.
            height: Int, the height of image.
            radius: The radius of heat map.
            fov: Float, the field of view for camera.
            intrinsic_matrix: A ndarray of shape 3*3, representing the camera intrinsic matrix. When this parameter is passed, `width`, `height` and `fov` will be ignroed.
        """
        self._send_data(
            "Get2DBBox", intrinsic_matrix, int(width), int(height), float(fov)
        )

    def Get3DBBox(self):
        """
        Get the 3d bounding box of objects in world coordinate.
        """
        self._send_data("Get3DBBox")
