def test_viewer():
    model = load_model_from_path("mujoco_py/tests/test.xml")
    sim = MjSim(model)
    viewer = MjViewer(sim)
    for _ in range(100):
        sim.step()
        viewer.render()

def test_multiple_sims():
    # Ensure that creating new simulators still produces good renderings.
    xml = """
    <mujoco>
        <asset>
            <texture name="t1" width="32" height="32" type="2d" builtin="flat" />
            <material name="m1" texture="t1" />
        </asset>
        <worldbody>
            <light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
            <camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
            <geom name="g1" pos="0 0 0" type="box" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
        </worldbody>
    </mujoco>
    """

    model = load_model_from_xml(xml)
    random_state = np.random.RandomState(0)

    for i in range(3):
        sim = MjSim(model)
        sim.forward()
        modder = TextureModder(sim, random_state=random_state)
        for j in range(2):
            modder.rand_checker('g1')
            compare_imgs(
                sim.render(201, 205, camera_name="topcam"),
                'test_multiple_sims.loop%d_%d.png' % (i, j))

def test_concurrent_rendering():
    '''Best-effort testing that concurrent multi-threaded rendering works.
    The test has no guarantees around being deterministic, but if it fails
    you know something is wrong with concurrent rendering. If it passes,
    things are probably working.'''
    err = None
    def func(sim, event):
        event.wait()
        sim.data.qpos[:] = 0.0
        sim.forward()
        img1 = sim.render(width=40, height=40, camera_name="camera1")
        img2 = sim.render(width=40, height=40, camera_name="camera2")
        try:
            assert np.sum(img1[:]) == 23255
            assert np.sum(img2[:]) == 12007
        except Exception as e:
            nonlocal err
            err = e

    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.render(100, 100)
    event = Event()
    threads = []
    for _ in range(100):
        thread = Thread(target=func, args=(sim, event))
        threads.append(thread)
        thread.start()
    event.set()
    for thread in threads:
        thread.join()
    assert err is None, "Exception: %s" % (str(err))

def test_data_attribute_getters():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    assert_array_equal(sim.data.get_body_xpos("body1"), [0, 0, 1])
    with pytest.raises(ValueError):
        sim.data.get_body_xpos("body_foo")
    with pytest.raises(RuntimeError):
        sim.data.get_xpos("body1")
    assert len(sim.data.get_body_xquat("body1")) == 4
    assert_array_equal(sim.data.get_body_xmat("body1").shape, (3, 3))
    # At (0, 1, 1) since the geom is displaced in the body
    assert_array_equal(sim.data.get_body_xipos("body1"), [0, 1, 1])

    assert_array_equal(sim.data.get_site_xpos("site1"), [1, 0, 1])
    assert_array_equal(sim.data.get_site_xmat("site1").shape, (3, 3))
    assert_array_equal(sim.data.get_geom_xpos("geom1"), [0.5, 0.4, 0.3])
    assert_array_equal(sim.data.get_geom_xpos("geom2"), [0, 1, 1])
    assert_array_equal(sim.data.get_geom_xmat("geom2").shape, (3, 3))
    assert_array_equal(sim.data.get_light_xpos("light1"), [0, 0, 3])
    assert_array_equal(sim.data.get_light_xdir("light1"), [0, 0, -1])
    assert_array_equal(sim.data.get_camera_xpos("camera1"), [3, 0, 0])
    assert_array_equal(sim.data.get_camera_xmat("camera1").shape, (3, 3))

    assert_array_equal(sim.data.get_joint_xaxis("joint1"), [0, 0, 1])
    assert_array_equal(sim.data.get_joint_xanchor("joint1"), [0, 0, 1])

def test_high_res():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()
    img = sim.render(1000, 1000)
    img = np.array(Image.fromarray(img).resize(size=(200, 200)))
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering.freecam.png')

def test_high_res():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()
    img = sim.render(1000, 1000)
    img = scipy.misc.imresize(img, (200, 200, 3))
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering.freecam.png')

def test_ignore_mujoco_warnings():
    # Two boxes on a plane need more than 1 contact (nconmax)
    xml = '''
    <mujoco>
      <size nconmax="1"/>
      <worldbody>
        <geom type="plane" size="1 1 0.1"/>
        <body pos="1 0 1"> <joint type="free"/> <geom size="1"/> </body>
        <body pos="0 1 1"> <joint type="free"/> <geom size="1"/> </body>
      </worldbody>
    </mujoco>
    '''
    model = load_model_from_xml(xml)
    sim = MjSim(model)

    sim.reset()
    with ignore_mujoco_warnings():
        # This should raise an exception due to the mujoco warning callback,
        # but it's suppressed by the context manager.
        sim.step()

    sim.reset()
    with pytest.raises(Exception):
        # test to make sure previous warning callback restored.
        sim.step()

def test_viewercontext():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()
    renderer = cymj.MjRenderContext(sim, offscreen=True)
    renderer.add_marker(type=const.GEOM_SPHERE,
                        size=np.ones(3) * 0.1,
                        pos=np.zeros(3),
                        mat=np.eye(3).flatten(),
                        rgba=np.ones(4),
                        label="mark")

def test_mj_sim_basics():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model, nsubsteps=2)

    sim.reset()
    sim.step()
    sim.reset()
    sim.forward()

def test_read_depth_buffer():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()
    ctx = MjRenderContext(sim, offscreen=True, opengl_backend='glfw')

    buf = np.zeros((11, 100), dtype=np.float32)
    assert buf.sum() == 0, f'{buf.sum()}'

    ctx.render(buf.shape[1], buf.shape[0], 0)
    ctx.read_pixels_depth(buf)
    assert buf.sum() != 0, f'{buf.sum()} {buf.max()}'

def test_glfw_context():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    render_context = MjRenderContext(sim, offscreen=True, opengl_backend='glfw')
    assert len(sim.render_contexts) == 1
    assert sim.render_contexts[0] is render_context
    assert isinstance(render_context.opengl_context, GlfwContext)

    compare_imgs(sim.render(201, 205, camera_name="topcam"), 'test_glfw_context.png')
    assert len(sim.render_contexts) == 1
    assert sim.render_contexts[0] is render_context

def test_joint_qpos_qvel_ops():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    # Test setting one with a list
    sim.data.set_joint_qpos("joint1", [1, 2, 3, 1, 0, 0, 0])
    # And the other with an np.ndarray
    sim.data.set_joint_qvel("joint1", np.array([1, 2, 3, 0.1, 0.1, 0.1]))
    sim.forward()
    assert_array_equal(sim.data.get_joint_qpos(
        "joint1"), [1, 2, 3, 1, 0, 0, 0])
    assert_array_equal(sim.data.get_joint_qvel(
        "joint1"), [1, 2, 3, 0.1, 0.1, 0.1])

    def __init__(self, model_path, frame_skip):

        if model_path.startswith("/"):
            fullpath = model_path
        else:
            fullpath = os.path.join(os.path.dirname(__file__), "assets", model_path)
        if not path.exists(fullpath):
            raise IOError("File %s does not exist" % fullpath)
        self.frame_skip = frame_skip
        self.model = load_model_from_path(fullpath)
        self.sim = MjSim(self.model)
        self.data = self.sim.data

        self.metadata = {
            'render.modes': ['human', 'rgb_array'],
            'video.frames_per_second': int(np.round(1.0 / self.dt))
        }
        self.mujoco_render_frames = False

        self.init_qpos = self.data.qpos.ravel().copy()
        self.init_qvel = self.data.qvel.ravel().copy()
        observation, _reward, done, _info = self._step(np.zeros(self.model.nu))
        assert not done
        self.obs_dim = np.sum([o.size for o in observation]) if type(observation) is tuple else observation.size

        bounds = self.model.actuator_ctrlrange.copy()
        low = bounds[:, 0]
        high = bounds[:, 1]
        self.action_space = spaces.Box(low, high)

        high = np.inf*np.ones(self.obs_dim)
        low = -high
        self.observation_space = spaces.Box(low, high)

        self._seed()

def test_textures():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.premod.png')

    random_state = np.random.RandomState(0)
    modder = TextureModder(sim, random_state=random_state)
    modder.whiten_materials()
    modder.whiten_materials(['g1', 'g2'])

    modder.set_rgb('g1', (255, 0, 0))
    modder.set_rgb('g2', (0, 255, 0))
    modder.set_rgb('g3', (0, 0, 255))
    modder.set_rgb('g4', (255, 0, 255))
    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.rgb.png')

    modder.set_checker('g1', (255, 0, 0), (0, 255, 0))
    modder.set_gradient('g2', (0, 255, 0), (0, 0, 255), vertical=True)
    modder.set_gradient('g3', (255, 255, 0), (0, 0, 255), vertical=False)
    modder.set_noise('g4', (0, 0, 255), (255, 0, 0), 0.1)
    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.variety.png')

    modder.rand_checker('g1')
    modder.rand_gradient('g2')
    modder.rand_noise('g3')
    modder.rand_rgb('g4')
    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.rand_specific.png')

    modder.rand_all('g1')
    modder.rand_all('g2')
    modder.rand_all('g3')
    modder.rand_all('g4')
    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.rand_all.png')

    modder.rand_checker('g1')
    modder.rand_checker('g2')
    modder.rand_checker('g3')
    modder.rand_checker('g4')
    mat_modder = MaterialModder(sim, random_state=random_state)
    mat_modder.rand_texrepeat('g1')
    mat_modder.rand_texrepeat('g2')
    mat_modder.rand_texrepeat('g3')
    mat_modder.rand_texrepeat('g4')
    compare_imgs(sim.render(201, 205, camera_name="topcam"),
                 'test_textures.rand_texrepeat.png')

def test_rendering():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    img, depth = sim.render(200, 200, depth=True)
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering.freecam.png')

    depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
    depth = np.asarray(depth * 255, dtype=np.uint8)
    assert depth.shape == (200, 200)

    # Unfortunately mujoco 2.0 renders slightly different depth image on mac and on linux here
    if "darwin" in sys.platform.lower():
        compare_imgs(depth, 'test_rendering.freecam.depth-darwin.png')
    else:
        compare_imgs(depth, 'test_rendering.freecam.depth.png')

    img = sim.render(100, 100, camera_name="camera1")
    assert img.shape == (100, 100, 3)
    compare_imgs(img, 'test_rendering.camera1.png')

    img = sim.render(200, 100, camera_name="camera1")
    assert img.shape == (100, 200, 3)
    compare_imgs(img, 'test_rendering.camera1.narrow.png')

    render_context = sim.render_contexts[0]
    render_context.add_marker(size=np.array([.4, .5, .6]),
                              pos=np.array([.4, .5, .6]),
                              rgba=np.array([.7, .8, .9, 1.0]),
                              label="mark")
    img = sim.render(200, 200, camera_name="camera1")
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering_markers.camera1.png')

    def test_ray(self):
        ''' Test raycasting and exclusions '''
        sim = MjSim(load_model_from_xml(self.xml))
        sim.forward()

        # Include all geoms
        self.check_rays(sim,
                        [0.9, 0.1, 0.9, 0.1, 0.9, 0.1, -1.0],
                        ['A', 'A', 'B', 'B', 'C', 'C', None])

        # Include static geoms, but exclude worldbody (which contains 'A')
        self.check_rays(sim,
                        [2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
                        ['B', 'B', 'B', 'B', 'C', 'C', None],
                        exclude_body=0)

        # Include static geoms, and exclude body 1 (which contains 'C')
        self.check_rays(sim,
                        [0.9, 0.1, 0.9, 0.1, -1.0, -1.0, -1.0],
                        ['A', 'A', 'B', 'B', None, None, None],
                        exclude_body=1)

        # Include static geoms, and exclude body 2 (which contains 'B')
        self.check_rays(sim,
                        [0.9, 0.1, 2.9, 1.9, 0.9, 0.1, -1.0],
                        ['A', 'A', 'C', 'C', 'C', 'C', None],
                        exclude_body=2)

        # Exclude static geoms ('A' is the only static geom)
        self.check_rays(sim,
                        [2.9, 1.9, 0.9, 0.1, 0.9, 0.1, -1.0],
                        ['B', 'B', 'B', 'B', 'C', 'C', None],
                        include_static_geoms=False)

        # Exclude static geoms, and exclude body 1 ('C')
        self.check_rays(sim,
                        [2.9, 1.9, 0.9, 0.1, -1.0, -1.0, -1.0],
                        ['B', 'B', 'B', 'B', None, None, None],
                        include_static_geoms=False, exclude_body=1)

        # Exclude static geoms, and exclude body 2 (which contains 'B')
        self.check_rays(sim,
                        [4.9, 3.9, 2.9, 1.9, 0.9, 0.1, -1.0],
                        ['C', 'C', 'C', 'C', 'C', 'C', None],
                        include_static_geoms=False, exclude_body=2)

def test_rendering():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    img, depth = sim.render(200, 200, depth=True)
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering.freecam.png')

    depth = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
    depth = np.asarray(depth * 255, dtype=np.uint8)
    assert depth.shape == (200, 200)
    compare_imgs(depth, 'test_rendering.freecam.depth.png')

    img = sim.render(100, 100, camera_name="camera1")
    assert img.shape == (100, 100, 3)
    compare_imgs(img, 'test_rendering.camera1.png')

    img = sim.render(200, 100, camera_name="camera1")
    assert img.shape == (100, 200, 3)
    compare_imgs(img, 'test_rendering.camera1.narrow.png')

    render_context = sim.render_contexts[0]
    render_context.add_marker(size=np.array([.4, .5, .6]),
                              pos=np.array([.4, .5, .6]),
                              rgba=np.array([.7, .8, .9, 1.0]),
                              label="mark")
    img = sim.render(200, 200, camera_name="camera1")
    assert img.shape == (200, 200, 3)
    compare_imgs(img, 'test_rendering_markers.camera1.png')

def test_mocap_ops():
    model = load_model_from_xml(BASIC_MODEL_XML)
    sim = MjSim(model)
    sim.forward()

    assert_array_equal(sim.data.get_body_xpos("mocap1"), [1, 0, 0])
    assert_array_equal(sim.data.get_mocap_pos("mocap1"), [1, 0, 0])
    assert_array_equal(sim.data.get_mocap_quat("mocap1"), [1, 0, 0, 0])
    new_pos = [2, 1, 1]
    new_quat = [0.707107, 0.707107, 0, 0]
    sim.data.set_mocap_pos("mocap1", new_pos)
    sim.data.set_mocap_quat("mocap1", new_quat)
    sim.forward()
    assert_array_equal(sim.data.get_mocap_pos("mocap1"), new_pos)
    assert_array_almost_equal(sim.data.get_mocap_quat("mocap1"), new_quat)
    assert_array_equal(sim.data.get_body_xpos("mocap1"), new_pos)
    assert_array_almost_equal(sim.data.get_body_xquat("mocap1"), new_quat)
    assert_array_almost_equal(sim.data.get_body_xmat("mocap1"),
                              [[1, 0, 0], [0, 0, -1], [0, 1, 0]])

def test_sim_save():
    model = load_model_from_xml(BASIC_MODEL_XML)
    assert model.nkey == 0
    sim = MjSim(model)

    with StringIO() as f:
        sim.save(f)

        f.seek(0)
        loaded_model = load_model_from_xml(f.read())

        assert loaded_model.nkey == 1

    with BytesIO() as f:
        sim.save(f, format='mjb')

        f.seek(0)
        loaded_model = load_model_from_mjb(f.read())
        assert loaded_model.nkey == 1

def test_mj_warning_raises():
    ''' Test that MuJoCo warnings cause exceptions. '''
    # Two boxes on a plane need more than 1 contact (nconmax)
    xml = '''
    <mujoco>
      <size nconmax="1"/>
      <worldbody>
        <geom type="plane" size="1 1 0.1"/>
        <body pos="1 0 1"> <joint type="free"/> <geom size="1"/> </body>
        <body pos="0 1 1"> <joint type="free"/> <geom size="1"/> </body>
      </worldbody>
    </mujoco>
    '''
    model = load_model_from_xml(xml)
    sim = MjSim(model)

    sim.reset()
    with pytest.raises(Exception):
        # This should raise an exception due to the mujoco warning callback
        sim.step()