import bpy
import numpy as np
import os
import sys


#****** CREATE DIRECTORY WHERE TO SAVE IMAGES ***********
img_dir = '/tmp/images/'
if not os.path.exists(img_dir):
    os.makedirs(img_dir)
    
#****** SET GPU AS DEVICE ***********
# Before check your gpu compute device type
ver = bpy.app.version[0]*1000 + bpy.app.version[1]*10 + bpy.app.version[2]
if ver < 2782 :
# Before blender 2.78b
	bpy.context.scene.cycles.device = 'GPU'
	bpy.context.user_preferences.system.compute_device_type = 'CUDA'
else :
# Since blender 2.78b
	sysp = bpy.context.user_preferences.addons['cycles'].preferences
	sysp.compute_device_type = 'CUDA'
	bpy.context.scene.cycles.device = 'GPU'


#****** INITIALIZE THE BLENDER 3D WORLD ********
# gather list of items of interest.
candidate_list = [item.name for item in bpy.data.objects if item.type == "MESH" or item.type == "LAMP" or item.type == "CAMERA"]

# select them only.
for object_name in candidate_list:
    bpy.data.objects[object_name].select = True

# remove all selected.
bpy.ops.object.delete()

# remove the meshes, they have no users anymore.
for item in bpy.data.meshes:
    bpy.data.meshes.remove(item)
    
# delete all materials
for i in bpy.data.materials.values():
    bpy.data.materials.remove(i)

# delete all textures
for i in bpy.data.textures.values():
    bpy.data.textures.remove(i)

# delete all images 
for i in bpy.data.images.values():
    # delete image path, this is only possible without a user
    i.user_clear()
    # delete all, except »Render Result«
    if i.name != "Render Result":
        bpy.data.images.remove(i)

#****** ADD OBJECTS ********
# Add a mesh
bpy.ops.mesh.primitive_plane_add(location=(0,0,0), radius=1.75)
plane = bpy.data.objects['Plane']

# Add a camera
bpy.ops.object.camera_add( location = ( 1.5, 0, 5), rotation = ( 0,0.3,0 ) )

# Add a lamp
bpy.ops.object.lamp_add( location = ( 1.0, -1.0, 6.0 ), type = 'POINT' );

#****** ADD MATERIALS *********
# create material
plane_material = bpy.data.materials.new("plane_material") # new material
plane_material_texture = plane_material.texture_slots.add() # add texture slot to material
plane_material.specular_intensity = 0.0

# edit texture
plane_image_path = os.path.expanduser('ton_image.jpg') # get image (needs the lib os)
plane_image = bpy.data.images.load(plane_image_path) # load image

# create texture
plane_texture = bpy.data.textures.new("plane_texture", type = 'IMAGE') # create texture type "image"
plane_texture.image = plane_image # link picture to image

# link texture and material
plane_material_texture.texture = plane_texture
plane_material_texture.texture_coords = 'UV'
plane_material_texture.mapping = 'FLAT'

# prepare object
bpy.context.scene.objects.active = plane # select pyramid object
bpy.ops.object.mode_set(mode='EDIT') # switch to edit mode
bpy.ops.uv.smart_project() # automatically unwrap object
bpy.ops.object.mode_set(mode='OBJECT') # switch back to object mode
plane.data.materials.append(plane_material) # link material to object

#************* INIT NODE COMPOSITOR FOR DEPTH RENDERING ***********
# switch on nodes
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
  
# clear default nodes
for n in tree.nodes:
    tree.nodes.remove(n)
  
# create input render layer node
rl = tree.nodes.new('CompositorNodeRLayers')      
rl.location = 185,285

# create output node for color image
out_c_node = tree.nodes.new('CompositorNodeOutputFile')
out_c_node.location = 600, 200
out_c_node.format.file_format = 'PNG'

# create output node for z map
out_z_node = tree.nodes.new('CompositorNodeOutputFile')
out_z_node.location = 600, 400
out_z_node.format.file_format = 'OPEN_EXR'

#links
links.new(rl.outputs[2], out_z_node.inputs[0])  # render node to output file for Z image
links.new(rl.outputs[0], out_c_node.inputs[0])  # render node to output file for color image

#************* RENDER AND SAVE IMAGES ***********
# Define image resolution for rendering
bpy.data.scenes["Scene"].render.resolution_x = 250
bpy.data.scenes["Scene"].render.resolution_y = 250
# resolution percentage : have to be 100% to have the whole image resolution defined earlier
bpy.context.scene.render.resolution_percentage = 100

# Activate which camera will be used for rendering
# (if more than one camera are defined, the rendering have to be repeated for each camera)
bpy.context.scene.camera = bpy.data.objects["Camera"]

# Define path where to save images
str4 = img_dir + "image"
out_c_node.base_path = str4 + "000"
out_z_node.base_path = str4 +"_Z_"+ "000"

# Render
bpy.ops.render.render()#write_still=True)

#************* EXIT BLENDER  ***********
sys.exit(0)