import unittest
import numpy as np
import vredner
import pyvredner
import torch
import os

options = vredner.RenderOptions(13,     # random seed
                                1024,    # spp
                                5,      # max bounces
                                0,      # sppse0
                                0,      # sppe
                                False)


class TestLoader(unittest.TestCase):
    """ change to the scene directory """
    @classmethod
    def setUpClass(cls):
        dir_path = os.path.dirname(os.path.realpath(__file__))
        scene_dir = os.path.join(dir_path, 'kitty')
        os.chdir(scene_dir)
        cls.scene_file = os.path.join(scene_dir, 'tar.xml')

    def test_loader(self):
        vredner.set_verbose(True)
        scene = vredner.Scene(self.scene_file)
        print(scene)
    
    def test_old_loader(self):
        vredner.set_verbose(True)
        # old loader
        scene, integrator = pyvredner.load_mitsuba(self.scene_file)
        integrator = vredner.Direct()
        cam = scene.camera
        c_scene = scene.c_obj()
        print(c_scene)

    def test_render(self):
        vredner.set_verbose(True)
        scene = vredner.Scene(self.scene_file)
        integrator = vredner.Direct()
        img = integrator.renderC(scene, options).reshape(256, 256, 3)
        pyvredner.imwrite(torch.from_numpy(img), "./img.exr")
    
    def test_forward_render(self):
        vredner.set_verbose(True)
        # old loader
        scene1, integrator = pyvredner.load_mitsuba("./scene.xml")
        integrator = vredner.Direct()
        cam = scene1.camera
        c_scene1 = scene1.c_obj()
        img1 = torch.zeros(cam.resolution[1], cam.resolution[0], 3)
        integrator.render(c_scene1, options, vredner.float_ptr(
            img1.data_ptr()))
        pyvredner.imwrite(img1, "./old_loader.exr")

        #new loader
        scene = vredner.Scene("./scene.xml")
        image = integrator.renderC(scene, options).reshape(256, 256, 3)
        pyvredner.imwrite(torch.from_numpy(image), "./new_loader.exr")
        
    def test_d_render(self):
        vredner.set_verbose(True)
        scene = vredner.Scene(self.scene_file)
        d_scene = vredner.Scene(self.scene_file)
        d_scene.setZero()
        integrator = vredner.Direct()
        d_image = torch.ones(256, 256, 3)
        d_image[:, :, 1:3] = 0
        scene.shape_list[0].requires_grad = True
        scene.shape_list[0].setTranslation(np.array([1., 0., 0.]))
        g_image = integrator.renderD(
            scene, d_scene, options, d_image.reshape((-1, 3)).numpy())\
            .reshape((256, 256, 3))
        pyvredner.imwrite(torch.from_numpy(g_image), "./grad.exr")

    """ currently pybind copy the nested properties of a readwrite object """

    def test_shape(self):
        vredner.set_verbose(True)
        dir_path = os.path.dirname(os.path.realpath(__file__))
        os.chdir(dir_path)
        scene = vredner.Scene()
        scene.load_file('scene.xml')
        shape = scene.shape_list[0]
        print(shape.vertices[0])
        shape.vertices[0] = np.array([1, 1, 1])
        print(shape.vertices[0])


if __name__ == '__main__':
    test = TestLoader()
    test.setUpClass()
    test.test_d_render()
