import pyredner
import numpy as np
import torch

cam = pyredner.Camera(position = torch.tensor([0.0, 0.0, -5.0]),
                      look_at = torch.tensor([0.0, 0.0, 0.0]),
                      up = torch.tensor([0.0, 1.0, 0.0]),
                      fov = torch.tensor([45.0]), # in degree
                      clip_near = 1e-2, # needs to > 0
                      resolution = (256, 256),
                      fisheye = False)

mat_grey = pyredner.Material(\
    diffuse_reflectance = \
        torch.tensor([0.5, 0.5, 0.5], device = pyredner.get_device()))
materials = [mat_grey]

shape_triangle = pyredner.Shape(\
    vertices = torch.tensor([[-1.7, 1.0, 0.0], [1.0, 1.0, 0.0], [-0.5, -1.0, 0.0]],
        device = pyredner.get_device()),
    indices = torch.tensor([[0, 1, 2]], dtype = torch.int32,
        device = pyredner.get_device()),
    uvs = None,
    normals = None,
    material_id = 0)

shape_light = pyredner.Shape(\
    vertices = torch.tensor([[-1.0, -1.0, -7.0],
                             [ 1.0, -1.0, -7.0],
                             [-1.0,  1.0, -7.0],
                             [ 1.0,  1.0, -7.0]], device = pyredner.get_device()),
    indices = torch.tensor([[0, 1, 2],[1, 3, 2]],
        dtype = torch.int32, device = pyredner.get_device()),
    uvs = None,
    normals = None,
    material_id = 0)

shapes = [shape_triangle, shape_light]
light = pyredner.AreaLight(shape_id = 1, 
                           intensity = torch.tensor([20.0,20.0,20.0]))
area_lights = [light]

scene = pyredner.Scene(cam, shapes, materials, area_lights)

scene_state_dict = scene.state_dict()
scene = pyredner.Scene.load_state_dict(scene_state_dict)

scene_args = pyredner.RenderFunction.serialize_scene(\
    scene = scene,
    num_samples = 16,
    max_bounces = 1)

render = pyredner.RenderFunction.apply
img = render(0, *scene_args)
pyredner.imwrite(img.cpu(), 'results/test_serialize/img.exr')

[((1526, 1577), 'pyredner.Scene', 'pyredner.Scene', (['cam', 'shapes', 'materials', 'area_lights'], {}), '(cam, shapes, materials, area_lights)\n', (1540, 1577), False, 'import pyredner\n'), ((1625, 1673), 'pyredner.Scene.load_state_dict', 'pyredner.Scene.load_state_dict', (['scene_state_dict'], {}), '(scene_state_dict)\n', (1655, 1673), False, 'import pyredner\n'), ((1688, 1775), 'pyredner.RenderFunction.serialize_scene', 'pyredner.RenderFunction.serialize_scene', ([], {'scene': 'scene', 'num_samples': '(16)', 'max_bounces': '(1)'}), '(scene=scene, num_samples=16,\n    max_bounces=1)\n', (1727, 1775), False, 'import pyredner\n'), ((82, 112), 'torch.tensor', 'torch.tensor', (['[0.0, 0.0, -5.0]'], {}), '([0.0, 0.0, -5.0])\n', (94, 112), False, 'import torch\n'), ((146, 175), 'torch.tensor', 'torch.tensor', (['[0.0, 0.0, 0.0]'], {}), '([0.0, 0.0, 0.0])\n', (158, 175), False, 'import torch\n'), ((204, 233), 'torch.tensor', 'torch.tensor', (['[0.0, 1.0, 0.0]'], {}), '([0.0, 1.0, 0.0])\n', (216, 233), False, 'import torch\n'), ((263, 283), 'torch.tensor', 'torch.tensor', (['[45.0]'], {}), '([45.0])\n', (275, 283), False, 'import torch\n'), ((1463, 1495), 'torch.tensor', 'torch.tensor', (['[20.0, 20.0, 20.0]'], {}), '([20.0, 20.0, 20.0])\n', (1475, 1495), False, 'import torch\n'), ((545, 566), 'pyredner.get_device', 'pyredner.get_device', ([], {}), '()\n', (564, 566), False, 'import pyredner\n'), ((728, 749), 'pyredner.get_device', 'pyredner.get_device', ([], {}), '()\n', (747, 749), False, 'import pyredner\n'), ((830, 851), 'pyredner.get_device', 'pyredner.get_device', ([], {}), '()\n', (849, 851), False, 'import pyredner\n'), ((1149, 1170), 'pyredner.get_device', 'pyredner.get_device', ([], {}), '()\n', (1168, 1170), False, 'import pyredner\n'), ((1261, 1282), 'pyredner.get_device', 'pyredner.get_device', ([], {}), '()\n', (1280, 1282), False, 'import pyredner\n')]