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, 8, 43, 59), 'pyredner.Scene', 'pyredner.Scene', ({(43, 23, 43, 26): 'cam', (43, 28, 43, 34): 'shapes', (43, 36, 43, 45): 'materials', (43, 47, 43, 58): 'area_lights'}, {}), '(cam, shapes, materials, area_lights)', False, 'import pyredner\n'), ((46, 8, 46, 56), 'pyredner.Scene.load_state_dict', 'pyredner.Scene.load_state_dict', ({(46, 39, 46, 55): 'scene_state_dict'}, {}), '(scene_state_dict)', False, 'import pyredner\n'), ((48, 13, 51, 20), 'pyredner.RenderFunction.serialize_scene', 'pyredner.RenderFunction.serialize_scene', (), '', False, 'import pyredner\n'), ((5, 33, 5, 63), 'torch.tensor', 'torch.tensor', ({(5, 46, 5, 62): '[0.0, 0.0, -5.0]'}, {}), '([0.0, 0.0, -5.0])', False, 'import torch\n'), ((6, 32, 6, 61), 'torch.tensor', 'torch.tensor', ({(6, 45, 6, 60): '[0.0, 0.0, 0.0]'}, {}), '([0.0, 0.0, 0.0])', False, 'import torch\n'), ((7, 27, 7, 56), 'torch.tensor', 'torch.tensor', ({(7, 40, 7, 55): '[0.0, 1.0, 0.0]'}, {}), '([0.0, 1.0, 0.0])', False, 'import torch\n'), ((8, 28, 8, 48), 'torch.tensor', 'torch.tensor', ({(8, 41, 8, 47): '[45.0]'}, {}), '([45.0])', False, 'import torch\n'), ((40, 39, 40, 69), 'torch.tensor', 'torch.tensor', ({(40, 52, 40, 68): '[20.0, 20.0, 20.0]'}, {}), '([20.0, 20.0, 20.0])', False, 'import torch\n'), ((15, 47, 15, 68), 'pyredner.get_device', 'pyredner.get_device', ({}, {}), '()', False, 'import pyredner\n'), ((20, 17, 20, 38), 'pyredner.get_device', 'pyredner.get_device', ({}, {}), '()', False, 'import pyredner\n'), ((22, 17, 22, 38), 'pyredner.get_device', 'pyredner.get_device', ({}, {}), '()', False, 'import pyredner\n'), ((31, 59, 31, 80), 'pyredner.get_device', 'pyredner.get_device', ({}, {}), '()', False, 'import pyredner\n'), ((33, 38, 33, 59), 'pyredner.get_device', 'pyredner.get_device', ({}, {}), '()', False, 'import pyredner\n')]