import bpy
import os.path
import math
import sys
import os
import csv
from random import sample
#file_path = 'D:/3D-FUTURE-model/final_choosed_FUTURE_chairs_text_caption_clear.csv'
#file_list = []
#with open(file_path, 'r') as r:
#    reader = csv.reader(r)
#    for row in reader:
#        file_name = row[0]
#        file_name = file_name.split('_')[0]
#        if file_name not in file_list:
#            file_list.append(file_name)
 
C = bpy.context
D = bpy.data
scene = D.scenes['Scene']


#inters =[i for i in range(48,72)] # [i for i in range(0,24)]#+
#fixed_views = [i for i in range(40,70)]
 
# cameras: a list of camera positions
# a camera position is defined by two parameters: (theta, phi),
# where we fix the "r" of (r, theta, phi) in spherical coordinate system.
 
# 5个固定视角：前 右 后 左 上
# cameras = [(60, 0), (60, 90), (60, 180), (60, 270),(0, 0)]
 
# 连续视角：固定某一个角度，然后环视3D物体，环视角度间隔inter=30度
fixed_view = 60
inter = 12
cameras = [(fixed_view, i) for i in range(0, 360, inter)] # 这会生成360/12=30个视角图片
render_setting = scene.render
 
# 输出图像大小 (W, H)
w = 224
h = 224
render_setting.resolution_x = w*2
render_setting.resolution_y = h*2

haved_folder = os.listdir('E:\\objaverse\\imgs\\')
count = 1
'''****************************************************************'''


def main(count):
    
    root_dir = 'E:\\objaverse\\objaverse\\'
    #for i,dir in enumerate(os.listdir(root_dir)[3736:4736]):
    for i,dir in enumerate(os.listdir(root_dir)):
    #for i,dir in enumerate(file_list):
#        inter = sample(inters, 20)
#        fixed_view = sample(fixed_views, 20)
        
        # cameras = [camera for camera in zip(inter,fixed_view)]
        #model_path = 'D:/tmp_data/767764f5-3ba5-457a-b040-4f52f11663ce/normalized_model.obj'        # 输入: 单个 .off 或 dataset.txt
        if dir in haved_folder:
            continue

        
        model_path = os.path.join(root_dir+dir, 'model.obj')
        
        image_dir = 'E:\\objaverse\\imgs\\'+dir         # 输出: 保存多视角图像的路径
        
        init_camera()
        fix_camera_to_origin()
        print('model path is ********', model_path) # model_path:'./airplane.off'
        # extension = '.'+str(count).zfill(3)
        count+=1
        do_model(model_path, image_dir, cameras)
#        except:
#            #delete_model('solid')
#            delete_model('normalized')
#            count-=1
#            continue
        

        
   
 
def init_camera():
    cam = D.objects['Camera']
    # D.objects['Plane'].location[2] = -0.56 + jiange  # calcuate the z of plane.(By Gang Li.)
    # select the camera object
    bpy.context.view_layer.objects.active = cam
      
    cam.select_set(True)
 
    # set the rendering mode to orthogonal and scale
    C.object.data.type = 'ORTHO'
    C.object.data.ortho_scale = 1.5  # 2.
 
 
def fix_camera_to_origin():
    origin_name = 'Origin'
    # create origin
    try:
        origin = D.objects[origin_name]
    except KeyError:
        bpy.ops.object.empty_add(type='SPHERE')
        D.objects['Empty'].name = origin_name
        origin = D.objects[origin_name]
 
    origin.location = (0, 0, 0)
 
    cam = D.objects['Camera']
    # scene.objects.active = cam
    bpy.context.view_layer.objects.active = cam
    cam.select_set(True)
 
    if 'Track To' not in cam.constraints:
        bpy.ops.object.constraint_add(type='TRACK_TO')
 
    cam.constraints['Track To'].target = origin
    cam.constraints['Track To'].track_axis = 'TRACK_NEGATIVE_Z'
    cam.constraints['Track To'].up_axis = 'UP_Y'
 
'''
def do_model(model_path, image_dir, extension, cameras):
    # model_path= 'F:\\DATA3D\ModelNet10_MV\\bathtub\\train\\bathtub_0003.off'
    # image_dir = 'F:\\DATA3D\\ModelNet10_MV_32_train\\'
    
    name = load_model(model_path) # -> name = 'bathtub_0003'

    name = name # +extension
    center_model(name)
    normalize_model(name)
 
    image_subdir = os.path.join(image_dir, name.split('_')[0], name) # path: image_dir\\bathtub\\bathtub_0003
    
    for i, c in enumerate(cameras):
        move_camera(c)
        render()
        save(image_dir, '%s_%d' % (name, i))
 
    delete_model(name)
'''

def do_model(model_path, image_dir, cameras):
    # model_path= 'F:\\DATA3D\ModelNet10_MV\\bathtub\\train\\bathtub_0003.off'
    # image_dir = 'F:\\DATA3D\\ModelNet10_MV_32_train\\'
    
    # name = load_model(model_path) # -> name = 'bathtub_0003'
    
    # bpy.context.space_data.params.directory = model_path.split('chair_')[0] # "D:\\3D-FUTURE-chairs_clear2\\test1"
    # bpy.ops.wm.obj_import(filepath=model_path, directory=model_path.split('chair_')[0], files=[{"name":'chair_'+model_path.split('chair_')[1], "name":'chair_'+model_path.split('chair_')[1]}])
    file_name = load_model(model_path)
    
    model_name = bpy.context.selected_objects[0].name
    
    #center_model(model_name)
    #normalize_model(model_name)
    
    image_subdir = os.path.join(image_dir, file_name.split('.obj')[0]) # path: image_dir\\bathtub\\bathtub_0003
    name = file_name.split('.obj')[0]
    
    for i, c in enumerate(cameras):
        move_camera(c)
        render()
        save(image_dir, '%s_%d' % (name, i))
 
    delete_model(name)

 
 
def load_model(model_path):
    # single .off: model_path='./airplane.off'
    # dataset.txt: model_path= 'F:\\DATA3D\ModelNet10_MV\\bathtub\\train\\bathtub_0003.off'
    d = os.path.dirname(model_path) # invalide for .off file
    ext = model_path.split('.')[-1] # ext: 'off'
 
    # Attention!  win10: ..path.split('\\')  linux: ..path.split('/')
    _model_path_tmp = model_path.split('\\')[-1] # _model_path_tmp: 'bathtub_0003.off'
    name = os.path.basename(_model_path_tmp).split('.')[0] # bathtub_0003
    # handle weird object naming by Blender for stl files
    if ext == 'stl':
        name = name.title().replace('_', ' ')
 
    if name not in D.objects:
        print('loading :' + name)
        if ext == 'stl':
            bpy.ops.import_mesh.stl(filepath=model_path, directory=d,
                                    filter_glob='*.stl')
        elif ext == 'off':
            bpy.ops.import_mesh.off(filepath=model_path, filter_glob='*.off')
        elif ext == 'obj':
            bpy.ops.import_scene.obj(filepath=model_path, filter_glob='*.obj')
        else:
            print('Currently .{} file type is not supported.'.format(ext))
            exit(-1)
    return name # name='airplane' -> 'bathtub_0003'
 
 
def delete_model(name):
    for ob in scene.objects:
        if ob.name.startswith('Camera') or ob.name.startswith('Origin') or ob.name.startswith('Plane') or ob.name.startswith('Sun'):
            continue
        elif ob.type == 'MESH': #and ob.name.startswith(name):
            print('I am choosed.')
            ob.select_set(True)
        else:
            ob.select_set(False)
            
    20/0
    bpy.ops.object.delete()
 
 
def center_model(name):
    bpy.ops.object.origin_set(type='GEOMETRY_ORIGIN')
    D.objects[name].location = (0, 0, 0)
 
 
def normalize_model(name):
    obj = D.objects[name]
    dim = obj.dimensions
    print('original dim:' + str(dim))
    if max(dim) > 0:
        dim = dim / max(dim)
    obj.dimensions = dim
 
    print('new dim:' + str(dim))
 
 
def move_camera(coord):
    def deg2rad(deg):
        return deg * math.pi / 180.
 
    r = 3.
    theta, phi = deg2rad(coord[0]), deg2rad(coord[1])
    loc_x = r * math.sin(theta) * math.cos(phi)
    loc_y = r * math.sin(theta) * math.sin(phi)
    loc_z = r * math.cos(theta)
 
    D.objects['Camera'].location = (loc_x, loc_y, loc_z)
 
 
def render():
    #bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0.557195, 0.0176854, 0.00812319, 1)
    bpy.ops.render.render()
 
 
def save(image_dir, name):
    path = os.path.join(image_dir, name + '.png')
    D.images['Render Result'].save_render(filepath=path)
    print('save to ' + path)
 
 
if __name__ == '__main__':
    main(count)