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')

[(43, 'pyredner.Scene', 'pyredner.Scene', False, 'import pyredner\n'), (46, 'pyredner.Scene.load_state_dict', 'pyredner.Scene.load_state_dict', False, 'import pyredner\n'), (48, 'pyredner.RenderFunction.serialize_scene', 'pyredner.RenderFunction.serialize_scene', False, 'import pyredner\n'), (5, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (6, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (7, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (8, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (40, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (15, 'pyredner.get_device', 'pyredner.get_device', False, 'import pyredner\n'), (20, 'pyredner.get_device', 'pyredner.get_device', False, 'import pyredner\n'), (22, 'pyredner.get_device', 'pyredner.get_device', False, 'import pyredner\n'), (31, 'pyredner.get_device', 'pyredner.get_device', False, 'import pyredner\n'), (33, 'pyredner.get_device', 'pyredner.get_device', False, 'import pyredner\n')]