import mujoco
from mujoco import viewer
import sys
import numpy as np
import time
import threading


class MjBasicRenderer:
    def __new__(cls, *args, **kwargs):
        return super().__new__(cls)

    def __init__(self, mj_model=None, mj_data=None):
        # keyboard flag
        self.render_paused = True
        self.exit_flag = False
        # init param
        self.mj_model = mj_model
        self.mj_data = mj_data
        self.renderer = "viewer"  # default
        self.viewer = None
        self._image = None

        # Set up mujoco viewer
        self.image_renderer = mujoco.Renderer(self.mj_model)

    def __del__(self):
        pass

    def _init_renderer(self):
        """Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3),
        another renderer with "mujoco_viewer"
        """

        def key_callback(keycode):
            if keycode == 32:  # space
                self.render_paused = not self.render_paused
            elif keycode == 256:  # escape
                self.exit_flag = not self.exit_flag

        if self.renderer == "viewer":
            # This function does not block, allowing user code to continue execution.
            self.viewer = viewer.launch_passive(
                self.mj_model,
                self.mj_data,
                key_callback=key_callback,
                show_left_ui=False,
                show_right_ui=False,
            )
            self.set_renderer_config()
        else:
            raise ValueError("Invalid renderer for some reason.")

    def render(self):
        """mujoco render"""
        if self.viewer is not None and self.render_paused is True:
            if self.viewer.is_running() and self.exit_flag is False:
                self.viewer: viewer.Handle
                self.viewer.sync()
            else:
                self.viewer.close()

    def set_renderer_config(self):
        """Setup mujoco global config while using viewer as renderer.
        It should be noted that the render thread need locked.
        """
        self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
        self.viewer.cam.azimuth -= 0.005
        with self.viewer.lock():
            self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(
                self.mj_data.time % 2
            )


try:
    import cv2
except ImportError:
    print("Could not import cv2, please install it to enable camera viewer.")


class MjMultiRenderer(MjBasicRenderer):

    # __slots__=('mj_model','mj_data','renderer','enable_camera_viewer')
    def __new__(cls, *args, **kwargs):
        return super().__new__(cls)

    def __init__(
        self,
        mj_model=None,
        mj_data=None,
        renderer=None,
        enable_camera_viewer=False,
        enable_depth=False,
    ):
        super().__init__(mj_model, mj_data)
        self._depth = None
        self.renderer = renderer
        self._init_renderer()
        self.enable_camera_viewer = enable_camera_viewer
        if self.enable_camera_viewer:
            self.enable_depth = enable_depth
            self._init_window()
        else:
            self.enable_depth = False
            print("No Camera View")

    def __del__(self):
        self.close()

    def _init_renderer(self):
        """
        Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3)
        """
        if self.renderer == "unity":
            # TODO: Support unity renderer.
            raise ValueError("Unity renderer init failed for no supporting reason")
        elif self.renderer == "viewer":
            super()._init_renderer()
            print("mujoco viewer init !")
        else:
            raise ValueError("renderer init failed for some reason.")

    def _init_window(self, name="Camera view"):
        if not self.enable_depth:
            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
        else:
            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
            cv2.namedWindow("Camera depth view", cv2.WINDOW_NORMAL)

    def render(self):
        """render mujoco"""
        if self.renderer == "viewer":
            super().render()
        elif self.renderer == "unity":
            # TODO: Support unity renderer.
            raise ValueError("Unity renderer not supported now.")
        else:
            raise ValueError("Invalid renderer for some reason.")

    def render(self):
        """mujoco render"""
        if self.viewer is not None and self.render_paused is True:
            if self.viewer.is_running() and self.exit_flag is False:
                self.viewer: viewer.Handle
                self.viewer.sync()
            else:
                self.viewer.close()

    def camera_render(self, cam=None):
        if self.enable_camera_viewer:
            if not self.enable_depth:
                rgb, depth = self.render_from_camera(cam)
                rgb = cv2.resize(rgb, (1920, 1600))
                cv2.imshow("Camera view", rgb)
                cv2.waitKey(1)

            else:
                rgb, depth = self.render_from_camera(cam)
                cv2.imshow("Camera view", rgb)
                cv2.imshow("Camera depth view", depth)
                cv2.waitKey(1)

        else:
            print("camera info disable")
            return

    def render_from_camera(self, cam=None):
        self.image_renderer.update_scene(self.mj_data, camera=cam)
        if self.enable_depth is True:
            self.image_renderer.enable_depth_rendering()
            org = self.image_renderer.render()
            depth = org[:, :]
            self.image_renderer.disable_depth_rendering()
            org = self.image_renderer.render()
            image = org[:, :, ::-1]
        else:
            org = self.image_renderer.render()
            image = org[:, :, ::-1]
            depth = np.zeros([240, 320])
        return image, depth

    def close(self):
        """close the environment."""
        if self.enable_camera_viewer and self.viewer.is_running() == False:
            cv2.destroyAllWindows()
        self.viewer.close()
        # sys.exit(0)

    # def get_cam_intrinsic(self, fovy=45.0, width=320, height=240):
    #     aspect = width * 1.0 / height
    #     fovx = np.degrees(2 * np.arctan(aspect * np.tan(np.radians(fovy / 2))))

    #     cx = 0.5 * width
    #     cy = 0.5 * height
    #     fx = cx / np.tan(fovx * np.pi / 180 * 0.5)
    #     fy = cy / np.tan(fovy * np.pi / 180 * 0.5)

    #     K = np.array([[fx, 0, cx],
    #                   [0, fy, cy],
    #                   [0, 0, 1]], dtype=np.float32)