import yep
import unittest
from pyvredner import shape
import vredner
import pyvredner
import os
import os.path
import numpy as np
import torch

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

curr_dir = os.path.dirname(os.path.realpath(__file__))
options = vredner.RenderOptions(13,     # random seed
                                1024,    # spp
                                6,      # max bounces
                                128,      # sppe
                                1280,      # sppse0
                                False)


class TestIntegrator(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        vredner.set_verbose(True)
        dir_path = os.path.dirname(os.path.realpath(__file__))
        scene_dir = os.path.join(dir_path, 'cbox')
        os.chdir(scene_dir)
        cls.scene_file = os.path.join(scene_dir, 'scene.xml')
        cls.integrator = vredner.Path2()
        cls.res = [256, 256]

    def test_renderC(self):
        vredner.set_verbose(True)
        # old loader
        # scene1, integrator = pyvredner.load_mitsuba("./scene.xml")
        # c_scene1 = scene1.c_obj()
        # new loader
        scene = vredner.Scene("./scene.xml")
        image = self.integrator.renderC(
            scene, options).reshape(self.res[1], self.res[0], 3)
        pyvredner.imwrite(torch.from_numpy(image), "./forward.exr")

    def test_renderD(self):
        vredner.set_verbose(True)
        scene = vredner.Scene(self.scene_file)
        d_scene = vredner.Scene(self.scene_file)
        d_scene.setZero()
        d_image = np.ones((self.res[1], self.res[0], 3))
        d_image[:, :, 1:3] = 0

        shape = scene.shape_list[0]
        scene.shape_list[0].requires_grad = True
        scene.shape_list[0].setTranslation(np.array([1., 0., 0.]))

        img = self.integrator.renderD(
            scene, d_scene, options, d_image.reshape((-1, 3)))
        img = img.reshape((self.res[1], self.res[0], 3))
        pyvredner.imwrite(torch.from_numpy(img), "./grad.exr")
        print(np.array(d_scene.shape_list[0].vertices).sum())

    def test_boundary(self):
        vredner.set_verbose(True)
        scene = vredner.Scene(self.scene_file)
        d_scene = vredner.Scene(self.scene_file)
        d_scene.setZero()
        d_image = torch.ones((self.res[1], self.res[0], 3))
        d_image[:, :, 1:3] = 0

        shape = scene.shape_list[0]
        scene.shape_list[0].requires_grad = True
        scene.shape_list[0].setTranslation(np.array([1., 0., 0.]))

        img = torch.zeros(self.res[1], self.res[0], 3)
        # integrator = vredner.DirectEdgeIntegrator(scene)
        integrator = vredner.IndirectEdgeIntegrator(scene)
        # integrator = vredner.DirectEdgeIntegrator(scene)
        integrator.d_render(scene, d_scene, options,
                            vredner.float_ptr(d_image.data_ptr()),
                            vredner.float_ptr(img.data_ptr()))
        # img = self.integrator.renderD(
        #     scene, d_scene, options, d_image.reshape((-1, 3)))
        # img = img.reshape((self.res[1], self.res[0], 3))
        pyvredner.imwrite(img, "./boundary.exr")
        print(np.array(d_scene.shape_list[0].vertices).sum())

    def test_fd(self):
        scene = vredner.Scene(self.scene_file)
        delta = 0.01
        image1 = self.integrator.renderC(
            scene, options).reshape(self.res[1], self.res[0], 3)

        vertices = np.array(scene.shape_list[0].vertices)
        for v in vertices:
            v += np.array([delta, 0., 0.])
        scene.shape_list[0].vertices = vertices

        image2 = self.integrator.renderC(
            scene, options).reshape(self.res[1], self.res[0], 3)
        fd = (image2 - image1)/delta
        pyvredner.imwrite(torch.from_numpy(fd), "./fd.exr")


if __name__ == '__main__':
    test = TestIntegrator()
    test.setUpClass()
    # yep.start('file_name.prof')
    test.test_renderC()
    # test.test_renderD()
    # test.test_boundary()
    # test.test_fd()
    # yep.stop()
    # test.test_fd()
