#~ Builtin Modules:       bpy, bpy.data, bpy.ops, bpy.props, bpy.types, bpy.context, bpy.utils, bgl, blf, mathutils
#~ Convenience Imports:   from mathutils import *; from math import *
#~ Convenience Variables: C = bpy.context, D = bpy.data
#~ 

import bpy
from bpy import data as D
from bpy import context as C
from mathutils import *
from math import *

import sys
import os.path
from os import path
import csv

"""Module for rendering images for the virtual data set. 
   Usage:
   > blender.exe --background --python Renders.py -- <cameraCalibration> <productModel> 
   <productReferencePose> <productRandomPoses> <packageModel> <packageReferencePose> 
   <packageRandomPoses> <outputDirectoryPath>
   
   Input arguments:
   * cameraCalibration - path to csv file containing camera calibration
   * productModel - path to the model of the product in the .obj file format
   * productReferencePose - csv containing the reference product pose
   * productRandomPoses - random poses of the product generated by RandomPoses.py
   * packageModel - path to the model of the package in the .obj file format
   * packageReferencePose - csv containing the reference package pose
   * packageRandomPoses - random poses of the package generated by RandomPoses.py
   * outputDirectoryPath - output directory for the rendered images

"""

__author__ = "Ondrej Klima"
__copyright__ = "Copyright 2020"
__credits__ = ["Ondrej Klima"]
__email__ = "iklima@fit.vutbr.cz"
__license__ = "BUT"
__version__ = "1.0"
__maintainer__ = "Ondrej Klima"

# Parsing input arguments
 
argv = sys.argv
try:
  argv = argv[argv.index("--") + 1:]
except ValueError:
  raise ValueError('Input arguments missing!')
  
try:
  cameraCalibrationFileName = argv[0]
except IndexError:
  raise IndexError('Camera pose file must be supplied on the command line')
if not path.isfile(cameraCalibrationFileName):
  raise ValueError('File "%s" does not exist!' % cameraCalibrationFileName)    

try:
  productFileName = argv[1]
except IndexError:
  raise IndexError('Product object must be supplied on the command line')
if not path.isfile(productFileName):
  raise ValueError('File "%s" does not exist!' % productFileName)

try:
  productReferencePoseFileName = argv[2]
except IndexError:
  raise IndexError('Product reference pose file must be supplied on the command line')
if not path.isfile(productReferencePoseFileName):
  raise ValueError('File "%s" does not exist!' % productReferencePoseFileName)
  
try:
  productRandomPosesFileName = argv[3]
except IndexError:
  raise IndexError('Product random poses file must be supplied on the command line')
if not path.isfile(productRandomPosesFileName):
  raise ValueError('File "%s" does not exist!' % productRandomPosesFileName)    

try:
  packagingFileName = argv[4]
except IndexError:
  raise IndexError('Packaging object must be supplied on the command line')
if not path.isfile(packagingFileName):
  raise ValueError('File "%s" does not exist!' % packagingFileName)

try:
  packagingReferencePoseFileName = argv[5]
except IndexError:
  raise IndexError('Package reference pose file object must be supplied on the command line')
if not path.isfile(packagingReferencePoseFileName):
  raise ValueError('File "%s" does not exist!' % packagingReferencePoseFileName)    
  
try:
  packagingRandomPosesFileName = argv[6]
except IndexError:
  raise IndexError('Package random poses file must be supplied on the command line')
if not path.isfile(packagingRandomPosesFileName):
  raise ValueError('File "%s" does not exist!' % packagingRandomPosesFileName)

try:
  outputDirectory = argv[7]
except IndexError:
  raise IndexError('Output directory for the rendered images must be supplied on the command line.')
    

# Delete the implicit Cube
try:
  bpy.ops.object.select_all(action='DESELECT')
  bpy.data.objects['Cube'].select_set(True)
  bpy.ops.object.delete() 
except KeyError:
  print('"Cube" object not found')

# Load the packaging object
packagingObject = bpy.ops.import_scene.obj(filepath=packagingFileName)
pack_object = bpy.context.selected_objects[0] 
print('Imported name: ', pack_object.name)

# Load the product object
productObject = bpy.ops.import_scene.obj(filepath=productFileName)
prod_object = bpy.context.selected_objects[0]
print('Imported name: ', prod_object.name)

# Set camera parameters
with open(cameraCalibrationFileName, newline='') as csvfile:
  cameraReader = csv.reader(csvfile, delimiter=',', quotechar='|')
  cameraData = list(cameraReader)

bpy.data.objects['Camera'].rotation_euler[0] = radians(float(cameraData[0][0]))
bpy.data.objects['Camera'].rotation_euler[1] = radians(float(cameraData[0][1]))
bpy.data.objects['Camera'].rotation_euler[2] = radians(float(cameraData[0][2]))

bpy.data.objects['Camera'].location[0] = float(cameraData[0][3])
bpy.data.objects['Camera'].location[1] = float(cameraData[0][4])
bpy.data.objects['Camera'].location[2] = float(cameraData[0][5])

bpy.data.cameras['Camera'].lens = float(cameraData[0][6])

# Product reference pose 
with open(productReferencePoseFileName, newline='') as csvfile:
  productReferenceReader = csv.reader(csvfile, delimiter=',', quotechar='|')
  productReferenceData = list(productReferenceReader)

refProductRotation = [.0, .0, .0]
refProductLocation = [.0, .0, .0]

refProductRotation[0] = float(productReferenceData[0][0])
refProductRotation[1] = float(productReferenceData[0][1])
refProductRotation[2] = float(productReferenceData[0][2])

refProductLocation[0] = float(productReferenceData[0][3])
refProductLocation[1] = float(productReferenceData[0][4])
refProductLocation[2] = float(productReferenceData[0][5])

refProductRotationRad = [0, 0, 0]
for i in range(len(refProductRotation)):
  refProductRotationRad[i] = radians(refProductRotation[i])
  
# Package reference pose
with open(packagingReferencePoseFileName, newline='') as csvfile:
  packageReferenceReader = csv.reader(csvfile, delimiter=',', quotechar='|')
  packageReferenceData = list(packageReferenceReader)

refPackageRotation = [.0, .0, .0]
refPackageLocation = [.0, .0, .0]

refPackageRotation[0] = float(packageReferenceData[0][0])
refPackageRotation[1] = float(packageReferenceData[0][1])
refPackageRotation[2] = float(packageReferenceData[0][2])

refPackageLocation[0] = float(packageReferenceData[0][3])
refPackageLocation[1] = float(packageReferenceData[0][4])
refPackageLocation[2] = float(packageReferenceData[0][5])

# Set product pose
bpy.data.objects[prod_object.name].location = refProductLocation
bpy.data.objects[prod_object.name].rotation_euler = refProductRotation

# Rendering engine set up
bpy.data.scenes['Scene'].render.engine = 'CYCLES'
bpy.data.scenes['Scene'].cycles.feature_set = 'SUPPORTED'
bpy.data.scenes['Scene'].cycles.device = 'GPU'
bpy.data.scenes['Scene'].cycles.use_denoising = True
bpy.data.scenes['Scene'].cycles.denoiser = 'OPENIMAGEDENOISE'
bpy.data.scenes['Scene'].cycles.samples = 256

bpy.data.scenes['Scene'].render.resolution_y = 2048;

# Set up light
bpy.data.lights['Light'].type = 'SUN'
bpy.data.lights['Light'].energy = 5

# Rendering border
bpy.context.scene.render.border_min_x = 0.3
bpy.context.scene.render.border_min_y = 0.2
bpy.context.scene.render.border_max_x = 0.7
bpy.context.scene.render.border_max_y = 0.8

bpy.context.scene.render.use_border = True
bpy.context.scene.render.use_crop_to_border = True

bpy.data.objects[pack_object.name].hide_render = True

with open(productRandomPosesFileName, newline='') as csvfile:
  posesReader = csv.reader(csvfile, delimiter=',', quotechar='|')
  productData = list(posesReader)
  
with open(packagingRandomPosesFileName, newline='') as csvfile:
  posesReader = csv.reader(csvfile, delimiter=',', quotechar='|')
  packageData = list(posesReader)
  
bpy.data.objects[prod_object.name].hide_render = False     
bpy.data.objects[pack_object.name].hide_render = False
bpy.context.scene.render.film_transparent = False  
  
i = 0
for row in productData:
  packRow = packageData[i]
      
  # Product poses
  rotation = row[0:3] # rotation is in degree units 
  #rotation.reverse()
  rotationRad = [0, 0, 0]

  for j in range(len(refProductRotation)):
    rotationRad[j] = radians(refProductRotation[j] + float(rotation[j]))
   
  bpy.data.objects[prod_object.name].rotation_euler = rotationRad  
    
  translation = [0, 0, 0]
  for j in range(len(refProductRotation)):
    bpy.data.objects[prod_object.name].location[j] = float(row[3 + j]) + refProductLocation[j] 
  
  bpy.data.objects[pack_object.name].rotation_euler[0] = radians(float(packRow[0]) + refPackageRotation[0])  
  bpy.data.objects[pack_object.name].rotation_euler[1] = radians(float(packRow[1]) + refPackageRotation[1])  
  bpy.data.objects[pack_object.name].rotation_euler[2] = radians(float(packRow[2]) + refPackageRotation[2])    
    
  translation = [0, 0, 0]
  for j in range(len(translation)):
    bpy.data.objects[pack_object.name].location[j] = float(packRow[3 + j]) + refPackageLocation[j]    
    
  bpy.data.objects[prod_object.name].hide_render = False     
  bpy.data.objects[pack_object.name].hide_render = False
  bpy.context.scene.render.film_transparent = False    
  bpy.context.scene.render.filepath = outputDirectory  + '\\%d.png' % i
  bpy.ops.render.render(write_still = True)

  i += 1