import base64
import cv2
import numpy as np
import pyrfuniverse.attributes as attr
import pyrfuniverse.utils.active_depth_generate as active_depth


class ActiveLightSensorAttr(attr.CameraAttr):
    """
    Class of IR-based depth sensor, which can simulate the noise of
    real-world depth camera and produce depth image in similar pattern.
    """

    def __init__(self, env, id: int, data=None):
        super().__init__(env, id, data)
        self.main_intrinsic_matrix = np.eye(3)
        self.ir_intrinsic_matrix = np.eye(3)

    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.\n
            self.data['active_depth']: IR-based depth, shape = (w,h)
        """
        super().parse_message(data)
        if "ir_left" in self.data and "ir_right" in self.data:
            self.data["ir_left"] = base64.b64decode(self.data["ir_left"])
            self.data["ir_right"] = base64.b64decode(self.data["ir_right"])

            image_left = np.frombuffer(self.data["ir_left"], dtype=np.uint8)
            image_left = cv2.imdecode(image_left, cv2.IMREAD_COLOR)[..., 2]
            image_right = np.frombuffer(self.data["ir_right"], dtype=np.uint8)
            image_right = cv2.imdecode(image_right, cv2.IMREAD_COLOR)[..., 2]
            left_extrinsic_matrix = np.array(
                [
                    [0.0, -1.0, 0.0, -0.0175],
                    [0.0, 0.0, -1.0, 0.0],
                    [1.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0],
                ]
            )
            right_extrinsic_matrix = np.array(
                [
                    [0.0, -1.0, 0.0, -0.072],
                    [0.0, 0.0, -1.0, 0.0],
                    [1.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0],
                ]
            )
            main_extrinsic_matrix = np.array(
                [
                    [0.0, -1.0, 0.0, 0.0],
                    [0.0, 0.0, -1.0, 0.0],
                    [1.0, 0.0, 0.0, 0.0],
                    [0.0, 0.0, 0.0, 1.0],
                ]
            )
            self.data["active_depth"] = active_depth.calc_main_depth_from_left_right_ir(
                image_left,
                image_right,
                left_extrinsic_matrix,
                right_extrinsic_matrix,
                main_extrinsic_matrix,
                self.ir_intrinsic_matrix,
                self.ir_intrinsic_matrix,
                self.main_intrinsic_matrix,
                lr_consistency=False,
                main_cam_size=(
                    self.main_intrinsic_matrix[0, 2] * 2,
                    self.main_intrinsic_matrix[1, 2] * 2,
                ),
                ndisp=128,
                use_census=True,
                register_depth=True,
                census_wsize=7,
                use_noise=False,
            )
            self.data["active_depth"][self.data["active_depth"] > 8.0] = 0
            self.data["active_depth"][self.data["active_depth"] < 0.1] = 0

    def GetActiveDepth(
        self,
        main_intrinsic_matrix_local: np.ndarray,
        ir_intrinsic_matrix_local: np.ndarray,
    ):
        """
        Get IR-based depth image.

        Args:
            main_intrinsic_matrix_local: np.ndarray The intrinsic matrix of main camera.
            ir_intrinsic_matrix_local: np.ndarray The intrinsic matrix of IR-based camera.
        """
        self.main_intrinsic_matrix = main_intrinsic_matrix_local
        self.ir_intrinsic_matrix = ir_intrinsic_matrix_local
        self._send_data("GetActiveDepth", self.ir_intrinsic_matrix)
