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):
    def test_loader(self):
        dir_path = os.path.dirname(os.path.realpath(__file__))
        os.chdir(dir_path)
        scene = vredner.Scene()
        scene.load_file('scene.xml')

    def test_render(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')
        integrator = vredner.Direct()
        img = torch.zeros(256, 256, 3)
        integrator.render(scene, options, vredner.float_ptr(
            img.data_ptr()))
        pyvredner.imwrite(img, "./img.exr")
    
    def test_d_render(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')
        integrator = vredner.Direct()
        img = torch.zeros(256, 256, 3)
        d_image = torch.ones(256, 256, 3)
        d_image[:, :, 1:3] = 0
        d_image1 = torch.zeros(256, 256, 3)
        d_scene = vredner.Scene()
        d_scene.load_file('scene.xml')
        d_scene.setZero()
        scene.shape_list[0].requires_grad = True
        scene.shape_list[0].setTranslation(np.array([1., 0., 0.]))
        integrator.d_render(scene, d_scene, options,
                    vredner.float_ptr(d_image.data_ptr()),
                    vredner.float_ptr(d_image1.data_ptr()))
        pyvredner.imwrite(d_image1, "./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.test_d_render()