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)
    
#****** SELECT RENDERING MODE (by default blender render) ***********
bpy.context.scene.render.engine = 'CYCLES'
print(bpy.context.scene.render.engine)
# Other parameters
#bpy.context.scene.cycles.progressive = 'PATH'
#bpy.context.scene.cycles.samples = 50
#bpy.context.scene.cycles.max_bounces = 1
#bpy.context.scene.cycles.min_bounces = 1
#bpy.context.scene.cycles.glossy_bounces = 1
#bpy.context.scene.cycles.transmission_bounces = 1
#bpy.context.scene.cycles.volume_bounces = 1
#bpy.context.scene.cycles.transparent_max_bounces = 1
#bpy.context.scene.cycles.transparent_min_bounces = 1
#bpy.context.scene.cycles.use_progressive_refine = True
#bpy.context.scene.render.tile_x = 64
#bpy.context.scene.render.tile_y = 64


#****** 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)
        

#************* 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


#****** 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 = 'SUN' );


#****** ADD MATERIALS *********
# Create material
plane_material = bpy.data.materials.new("plane_material") # new material

# Activate nodes for materials
plane_material.use_nodes = True
plane_nt = plane_material.node_tree
plane_nodes = plane_nt.nodes
plane_links = plane_nt.links

# clear previous materials nodes
while(plane_nodes): plane_nodes.remove(plane_nodes[0])

# Add nodes for materials
plane_output  = plane_nodes.new("ShaderNodeOutputMaterial")
plane_diffuse = plane_nodes.new("ShaderNodeBsdfDiffuse")
plane_texture = plane_nodes.new("ShaderNodeTexImage")
plane_texcoord   = plane_nodes.new("ShaderNodeTexCoord")

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

# Links materials nodes
plane_links.new( plane_output.inputs['Surface'], plane_diffuse.outputs['BSDF'])
plane_links.new(plane_diffuse.inputs['Color'],   plane_texture.outputs['Color'])
plane_links.new(plane_texture.inputs['Vector'],    plane_texcoord.outputs['Generated'])

# Distribute nodes along the x axis
for index, node in enumerate((plane_texcoord, plane_texture, plane_diffuse, plane_output)):
    node.location.x = 200.0 * index

# Link material to object
plane.data.materials.append(plane_material)


#************* 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)