import unittest
import numpy as np
import os
import psdr_cpu as psdr
from pypsdr.validate import RenderOptions, Scene


class TestIntegrator(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cwd = os.path.dirname(os.path.realpath(__file__))
        os.chdir(cwd)
        cls.direct = psdr.Direct()
        cls.path = psdr.Path2()
        cls.scene = Scene("../../../data/scenes/two_triangles/scene.xml")
        cls.sceneAD = psdr.SceneAD(cls.scene)
        cls.sampler = psdr.RndSampler(1, 1)
        cls.options = RenderOptions(1, 1, 100, 0, 0, 0)
        cls.pixel_idx = [263, 285]

    def test_direct(self):
        value = self.direct.Li(self.scene, self.sampler, self.pixel_idx)
        print(value)
        print(self.sampler.state)

    def test_path(self):
        value = self.path.Li(self.scene, self.sampler,
                             self.pixel_idx, self.options)
        print(value)
        print(self.sampler.state)

    def test_d_direct(self):
        self.direct.d_Li(self.sceneAD, self.pixel_idx,
                         self.sampler, 1, [1, 0, 0])
        sceneD = self.sceneAD.der
        sceneD.buildParamMap()
        params = sceneD.param_map2
        s = np.concatenate([np.array([kv[1]]).reshape(-1)
                           for kv in params.items()])
        print(s.sum())

    def test_d_path(self):
        self.path.d_Li(self.sceneAD, self.pixel_idx, self.sampler, 1,
                       [1, 0, 0], self.options)
        sceneD = self.sceneAD.der
        sceneD.buildParamMap()
        params = sceneD.param_map2
        s = np.concatenate([np.array([kv[1]]).reshape(-1)
                           for kv in params.items()])
        print(s.sum())


if __name__ == '__main__':
    test = TestIntegrator()
    test.setUpClass()
    test.test_direct()
    test.setUpClass()
    test.test_path()
