import bpy

import os
import math
from mathutils import Vector
import random
import sys
import argparse

def render_objs(root_folder):
    full_folder = os.path.join(root_folder, "Obj_Files", "Full")
    render_folder = os.path.join(root_folder, "Img_Files")

    bpy.ops.wm.read_homefile(use_empty=True)

    light_data = bpy.data.lights.new(name="sun", type='SUN')
    light_data.energy = 7

    light_object = bpy.data.objects.new(name="sun_obj", object_data=light_data)
    light_object.rotation_euler = (math.radians(54), 0, math.radians(135))
    bpy.context.collection.objects.link(light_object)
    bpy.context.view_layer.objects.active = light_object
    light_object.location = (5, 5, 5)

    dg = bpy.context.evaluated_depsgraph_get() 
    dg.update()

    camera_data = bpy.data.cameras.new(name='Camera')
    camera_object = bpy.data.objects.new('Camera', camera_data)
    camera_object.location = (3, 3, 3)
    camera_object.rotation_euler = (math.radians(53), 0, math.radians(135))
    bpy.context.scene.collection.objects.link(camera_object)
    bpy.context.scene.camera = bpy.data.objects['Camera']

    collection = bpy.data.collections.new("import_group")
    bpy.context.scene.collection.children.link(collection)

    material = bpy.data.materials.new(name="my_material")
    material.use_nodes = True
    principled = material.node_tree.nodes.new(type='ShaderNodeBsdfPrincipled')
    inp = material.node_tree.nodes['Material Output'].inputs['Surface']
    material.node_tree.links.new(inp,principled.outputs['BSDF'])
    checker = material.node_tree.nodes.new(type='ShaderNodeTexChecker')
    material.node_tree.links.new(principled.inputs['Base Color'],checker.outputs['Color'])
    texcoord = material.node_tree.nodes.new(type='ShaderNodeTexCoord')
    material.node_tree.links.new(checker.inputs['Vector'],texcoord.outputs['UV'])
    
    checker.inputs['Scale'].default_value = 15
    
    obj_files = os.listdir(full_folder)
    obj_files.sort(key=lambda f: int(f.split("_")[1]))

    for onum, obj_file in enumerate(obj_files):
        print("Mesh", onum+1, "/", len(obj_files))
        filepath = os.path.join(full_folder, obj_file)
        name, ext = os.path.splitext(obj_file)

        imported_object = bpy.ops.import_scene.obj(filepath=filepath, axis_up="Z", axis_forward="Y")

        x_min = None
        x_max = None
        y_min = None
        y_max = None
        z_min = None
        z_max = None
        for i, o in enumerate(bpy.context.selected_objects):
            print(i+1, "/", len(bpy.context.selected_objects))
            bpy.context.view_layer.objects.active = o
            bpy.ops.object.collection_link(collection="import_group")
            bbox_corners = [o.matrix_world @ Vector(corner) for corner in o.bound_box]
            o_x_min = min(co[0] for co in bbox_corners)
            o_x_max = max(co[0] for co in bbox_corners)
            o_y_min = min(co[1] for co in bbox_corners)
            o_y_max = max(co[1] for co in bbox_corners)
            o_z_min = min(co[2] for co in bbox_corners)
            o_z_max = max(co[2] for co in bbox_corners)

            if x_min is None:
                x_min = o_x_min
                x_max = o_x_max
                y_min = o_y_min
                y_max = o_y_max
                z_min = o_z_min
                z_max = o_z_max
            else:
                x_min = min(x_min, o_x_min)
                x_max = max(x_max, o_x_max)
                y_min = min(y_min, o_y_min)
                y_max = max(y_max, o_y_max)
                z_min = min(z_min, o_z_min)
                z_max = max(z_max, o_z_max)

        max_bbox_length = max(x_max - x_min, y_max - y_min, z_max - z_min)
        goal_bbox_length = 2.6

        current_center = Vector(((x_max + x_min)/2, (y_max + y_min)/2, (z_max + z_min)/2))
        goal_center = Vector((0, 0, -0.2))

        scale_factor = goal_bbox_length / max_bbox_length
        translation = goal_center - (current_center * scale_factor)

        bpy.ops.transform.resize(value=(scale_factor, scale_factor, scale_factor))
        bpy.ops.transform.translate(value=translation)
        
        x_min = None
        x_max = None
        y_min = None
        y_max = None
        for i, o in enumerate(bpy.context.selected_objects):
            print(i+1, "/", len(bpy.context.selected_objects))
            bpy.context.view_layer.objects.active = o
            bpy.ops.object.collection_link(collection="import_group")
            
            if (o.data.uv_layers.active is None) or (len(o.data.uv_layers.active.data) == 0):
                continue
            
            for face in o.data.polygons:
                for loop_idx in face.loop_indices:
                    coord = o.data.uv_layers.active.data[loop_idx].uv
                    
                    if x_min is None:
                        x_min = coord.x
                        x_max = coord.x
                        y_min = coord.y
                        y_max = coord.y
                    else:
                        x_min = min(x_min, coord.x)
                        x_max = max(x_max, coord.x)
                        y_min = min(y_min, coord.y)
                        y_max = max(y_max, coord.y)
                        

        max_bbox_length = max(x_max - x_min, y_max - y_min)
        goal_bbox_length = 1
        
        current_center = Vector(((x_max + x_min)/2, (y_max + y_min)/2))
        goal_center = Vector((0, 0))
                                
        scale_factor = goal_bbox_length / max_bbox_length
        translation = goal_center - (current_center * scale_factor)

        for i, o in enumerate(bpy.context.selected_objects):
            print(i+1, "/", len(bpy.context.selected_objects))
            bpy.context.view_layer.objects.active = o
            bpy.ops.object.collection_link(collection="import_group")
            
            if (o.data.uv_layers.active is None) or (len(o.data.uv_layers.active.data) == 0):
                continue

            for face in o.data.polygons:
                for loop_idx in face.loop_indices:
                    coord = o.data.uv_layers.active.data[loop_idx].uv
                    
                    new_coord = coord * scale_factor + translation
                    o.data.uv_layers.active.data[loop_idx].uv = new_coord
        

#         colors = [(1, 0.25, 1, 1), (1, 0.25, 0.25, 1), (1, 0.75, 0.25, 1), (1, 1, 0.25, 1), (0.25, 1, 0.25, 1), \
#                   (0.25, 1, 1, 1), (0.25, 0.25, 1, 1)]
        
        
        colors_rgb = [(228,26,28), (55,126,184), (77,175,74), (152,78,163), (255,127,0), (255,255,51)]
        colors = [tuple([c/255 for c in color]) + (1,) for color in colors_rgb]
        
        checker.inputs['Color1'].default_value = random.choice(colors)
        checker.inputs['Color2'].default_value = (0.95, 0.95, 0.95, 1.)

        for i, o in enumerate(bpy.context.selected_objects):
            o.active_material_index = 0

            for i in range(len(o.material_slots)):
                bpy.ops.object.material_slot_remove({'object': o})

            o.data.materials.append(material)

        bpy.context.scene.render.filepath = os.path.join(render_folder, name + ".png")
        bpy.context.scene.render.film_transparent = True
        bpy.context.scene.render.engine = "CYCLES"
        
        bpy.context.scene.render.resolution_x = bpy.context.scene.render.resolution_y

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

        for o in bpy.context.selected_objects:
            bpy.data.objects.remove(o)
            
            
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Render all objs")
    parser.add_argument("-r", "--root-folder", type=str, required=True, dest="root_folder")
    
    argv = sys.argv
    
    if "--" not in argv:
        argv = []
    else:
        argv = argv[argv.index("--") + 1:]
    
    args = parser.parse_args(argv)
        
    root_folder = os.path.abspath(args.root_folder)
    
    render_objs(root_folder)

