# Run as: blender -b <filename> -P <this_script> -- <image_path>
import argparse
import math
import os 
from pathlib import Path
import sys
home = str(Path.home())
sys.path.insert(0, os.path.join(home, "code/blender"))
sys.path.insert(0, "/app")
import numpy as np
import util

import bpy 
import mathutils
from mathutils import Vector

# Assume the last argument is image path
print("sys.argv: {}".format(sys.argv[0]))
print("sys.argv: {}".format(sys.argv[1]))
print("sys.argv: {}".format(sys.argv[2]))

p = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
p.add_argument('--num_observations', type=int, required=True, help='Number of observations')
p.add_argument('--resolution', type=int, required=True, help='Resolution size.')
p.add_argument('--sphere_radius', type=float, required=True, help='Sphere radius.')
p.add_argument('--sampling_mode', type=str, required=True, help='Options: spherical and archimedean_spiral and z_height_constant')
p.add_argument('--render_source', type=str, required=True, default='geon_blender_files', help='Options: spherical and archimedean_spiral and z_height_constant')
p.add_argument('--background_color', type=str, required=True, help='Options: white and black')
p.add_argument('--save_dir', type=str, required=True, help='The path the output will be dumped to.')
p.add_argument('--write_z_buffer', action='store_true', default=False, help='If you want to save depth map as well')
p.add_argument('--scale_x', type=float, required=True, help='Scale the object in x-axis.')
p.add_argument('--scale_y', type=float, required=True, help='Scale the object in y-axis.')
p.add_argument('--scale_z', type=float, required=True, help='Scale the object in z-axis.')


argv = sys.argv
argv = sys.argv[sys.argv.index("--") + 1:]


args = p.parse_args(argv)

dirname = os.path.join(args.save_dir, "geon_"+str(args.num_observations)+"cam_" \
            +str(args.sphere_radius)+"radius" \
            +args.sampling_mode+str(args.resolution)+"res_"+args.background_color+"_bg")
if args.write_z_buffer:
    dirname += "_with_depthmap"
filename = os.path.splitext(os.path.basename(sys.argv[2]))[0] # filename == geon_type
bg_color = args.background_color


def start_rendering(scene, foldername, filename, resolution=100):
    # render settings
    scene.render.resolution_x = resolution
    scene.render.resolution_y = resolution 
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = 'PNG'
    scene.render.filepath = os.path.join(foldername, filename)
    bpy.ops.render.render(write_still = 1)


# e.g. output_dir = data/Geon3D/geon_30cam/geon_axis_c_cross_c_sweep_ce_termination_c
def new_render(scene, output_dir, blender_cam2world_matrices, \
    write_cam_params=False, resolution=100, write_z_buffer=False):
    # camera settings 
    scene.camera.data.sensor_height = scene.camera.data.sensor_width
    old_lens = scene.camera.data.lens
    print(old_lens)
    now_lens = scene.camera.data.lens
    print(now_lens)
    # render settings
    scene.render.resolution_x = resolution
    scene.render.resolution_y = resolution 
    scene.render.resolution_percentage = 100
    scene.render.image_settings.file_format = 'PNG'

    if write_cam_params:
        img_dir = os.path.join(output_dir, 'rgb')
        if write_z_buffer:
            depth_dir = os.path.join(output_dir, 'depth')
        pose_dir = os.path.join(output_dir, 'pose')

        util.cond_mkdir(img_dir)
        util.cond_mkdir(pose_dir)
    else:
        img_dir = output_dir
        util.cond_mkdir(img_dir)

    if write_cam_params:
        K = util.get_calibration_matrix_K_from_blender(scene.camera.data)
        with open(os.path.join(output_dir, 'intrinsics.txt'),'w') as intrinsics_file:
            intrinsics_file.write('%f %f %f 0.\n'%(K[0][0], K[0][2], K[1][2]))
            intrinsics_file.write('0. 0. 0.\n')
            intrinsics_file.write('1.\n')
            intrinsics_file.write('%d %d\n'%(resolution, resolution))

    for i in range(len(blender_cam2world_matrices)):
        scene.camera.matrix_world = blender_cam2world_matrices[i]

        # Render the object
        if os.path.exists(os.path.join(img_dir, '%06d.png' % i)):
            continue

        # Render the color image
        scene.render.filepath = os.path.join(img_dir, '%06d.png'%i)
        bpy.context.scene.render.engine = 'CYCLES'
        bpy.ops.render.render(write_still=True)

        if write_cam_params:
            # Write out camera pose
            RT = util.get_world2cam_from_blender_cam(scene.camera)
            cam2world = RT.inverted()
            with open(os.path.join(pose_dir, '%06d.txt'%i),'w') as pose_file:
                matrix_flat = []
                for j in range(4):
                    for k in range(4):
                        matrix_flat.append(cam2world[j][k])
                pose_file.write(' '.join(map(str, matrix_flat)) + '\n')
        
        if write_z_buffer:
            tmp_nodetree = bpy.context.scene.node_tree
            nodetree = bpy.context.scene.node_tree
            bpy.context.area.ui_type = 'CompositorNodeTree'
            node1 = nodetree.nodes.new("ShaderNodeCameraData")
            node2 = nodetree.nodes.new("ShaderNodeEmission")
            nodetree.links.new(node1.outputs["Image"],node2.inputs[1])


            bpy.context.area.ui_type = 'CONSOLE'
            

def create_light(location):
    light_data = bpy.data.lights.new('light', type='POINT')
    light_data.energy = 1000
    light = bpy.data.objects.new('light', light_data)
    light.data.use_shadow = False
    bpy.context.collection.objects.link(light)
    light.location = mathutils.Vector(location)

def calc_camera_location(default_cam_loc, cam_id, num_cam):
    # defualt_cam_loc's data type -> Vector
    x, y, z = default_cam_loc[0], default_cam_loc[1], default_cam_loc[2]
    radius = Vector((x, y, 0)).length
    angle = 2 * math.pi * (cam_id-1) / num_cam #random.random()
    
    # Randomly place the camera on a circle around the object
    # at the same height as the main camera
    new_camera_pos = Vector((radius * math.cos(angle), radius * math.sin(angle), z))
    return new_camera_pos
    

def create_camera(scene, default_cam, target_obj, cam_id, num_cam):
    if cam_id == 1:
        return scene, default_cam
    else:
        default_cam_loc = default_cam.location
        new_camera_pos = calc_camera_location(default_cam_loc, cam_id, num_cam)

        cam_data = bpy.data.cameras.new('camera')
        cam = bpy.data.objects.new('camera', cam_data)
        #scene.objects.link(cam)
        bpy.context.collection.objects.link(cam)
        cam.location = new_camera_pos
        scene.camera = cam

        track_to = cam.constraints.new('TRACK_TO')
        track_to.name = "OB1_Track"
        track_to.target = target_obj
        track_to.up_axis = 'UP_Y'
        track_to.track_axis = 'TRACK_NEGATIVE_Z'

        return scene, cam


def create_camera_old():
    # Create the camera
    cam_data = bpy.data.cameras.new('camera')
    cam = bpy.data.objects.new('camera', cam_data)
    #scene.objects.link(cam)
    bpy.context.collection.objects.link(cam)
    scene.camera = cam

    cam.location = mathutils.Vector((10, -5, 5))
    #cam.rotation_euler = mathutils.Euler((0.9, 0.0, 1.1))


    track_to = cam.constraints.new('TRACK_TO')
    track_to.name = "OB1_Track"
    track_to.target = target_obj
    track_to.up_axis = 'UP_Y'
    track_to.track_axis = 'TRACK_NEGATIVE_Z'

    #filepath = os.path.join(dirname, "image_test1")
    filename_1 = filename + '_1'
    start_rendering(scene, filename_1, resolution=100)

    z = cam.location[2]
    radius = Vector((cam.location[0], cam.location[1], 0)).length
    angle = 2 * math.pi * random.random()
    
    # Randomly place the camera on a circle around the object
    # at the same height as the main camera
    new_camera_pos = Vector((radius * math.cos(angle), radius * math.sin(angle), z))
    # Create the camera
    cam_data = bpy.data.cameras.new('camera2')
    cam = bpy.data.objects.new('camera2', cam_data)
    bpy.context.collection.objects.link(cam)
    cam.location = new_camera_pos
    scene.camera = cam
   
    # Add a new track to constraint and set it to track your object
    track_to = cam.constraints.new('TRACK_TO')
    track_to.name = "OB1_Track"
    track_to.target = target_obj
    track_to.up_axis = 'UP_Y'
    track_to.track_axis = 'TRACK_NEGATIVE_Z'

def rescale_geon(x, y, z):
    '''
    x = 1.1 will scale the object x1.1 in the x-axis.
    '''
    bpy.ops.transform.resize(value=(x, y, z), orient_type='GLOBAL', 
        orient_matrix=((1, 0, 0), (0, 1, 0), (0, 0, 1)), orient_matrix_type='GLOBAL', 
        constraint_axis=(True, False, False), mirror=True, use_proportional_edit=False, 
        proportional_edit_falloff='SMOOTH', proportional_size=1, 
        use_proportional_connected=False, use_proportional_projected=False)


def main():
    # After importing
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.normals_make_consistent(inside=False)
    bpy.ops.object.mode_set(mode='OBJECT')

    target_obj = bpy.data.objects['Geon']
    maxDim = max(target_obj.dimensions)
    target_obj.dimensions = target_obj.dimensions / maxDim
    override = bpy.context.copy()
    override["selected_editable_objects"] = [target_obj]
    bpy.ops.object.origin_set(override, type='ORIGIN_GEOMETRY', center='BOUNDS')
    target_obj.location = (0,0,0)

    rescale_geon(x=args.scale_x, y=args.scale_y, z=args.scale_z)    
    override = bpy.context.copy()
    override["selected_editable_objects"] = [target_obj]
    bpy.ops.object.origin_set(override, type='ORIGIN_GEOMETRY', center='BOUNDS')
    target_obj.location = (0,0,0)


    scene = bpy.context.scene
    if bg_color == 'black':
        bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0, 0, 0, 1)
    elif bg_color == 'gray':
        pass
    else:
        raise NotImplementedError()

    create_light((3, -4.2, 5))
    create_light((-3, 4.2, 5))
    create_light((3, -4.2, -5))
    create_light((-3, 4.2, -5))


    if args.sampling_mode == 'spherical':
        cam_locations = util.sample_spherical(args.num_observations, args.sphere_radius)
    elif args.sampling_mode == 'archimedian_spiral':
        cam_locations = util.get_archimedean_spiral(args.sphere_radius, args.num_observations)
    else:
        raise NotImplementedError("sampling_mode={} is not implemented.".format(args.sampling_mode))

    # Run a for loop to render an object from
    # multiple angles
    default_cam = bpy.data.objects['Camera'] #.location 
    subfolder = '%03.1f'%args.scale_x + '_%03.1f'%args.scale_y + '_%03.1f'%args.scale_z
    foldername = os.path.join(os.path.join(dirname, filename), subfolder) # filename=geon_type

    obj_location = np.zeros((1,3))

    cv_poses = util.look_at(cam_locations, obj_location)
    blender_poses = [util.cv_cam2world_to_bcam2world(m) for m in cv_poses]

    new_render(scene, foldername, blender_poses, write_cam_params=True, \
            resolution=args.resolution, write_z_buffer=args.write_z_buffer)


def main_z_height_constant():
    # Assume object, material and texture name (and settings) are valid
    target_obj = bpy.data.objects['Geon']
    scene = bpy.context.scene
    rescale_geon(x=args.scale_x, y=args.scale_y, z=args.scale_z) 

    # Render to separate file, identified by texture file
    # Change the background color to black
    if bg_color == 'black':
        bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0, 0, 0, 1)
    elif bg_color == 'gray':
        pass
    else:
        raise NotImplementedError()

    # Render still image, automatically write to output path
    resolution=args.resolution #224
    create_light((3, -4.2, 5))
    create_light((-3, 4.2, 5))
    create_light((3, -4.2, -5))
    create_light((-3, 4.2, -5))
    # default light location can be found by bpy.data.objects['Light'].location 
    # Vector((4.076245307922363, 1.0054539442062378, 5.903861999511719)) 

    # Run a for loop to render an object from
    # multiple angles
    default_cam = bpy.data.objects['Camera'] #.location 
    foldername = os.path.join(dirname, filename)
    for i in range(1, args.num_observations+1):
        if not os.path.exists(foldername):
            os.makedirs(foldername)
        new_scene, new_cam = create_camera(scene, default_cam, target_obj, i, args.num_observations)
        start_rendering(scene=new_scene, foldername=foldername, 
            filename="cam_" + str(i), resolution=resolution)

if __name__=='__main__':
    if  args.render_source == 'geon_blender_files':
        if args.sampling_mode == 'z_height_constant':
            main_z_height_constant()
        else:
            main()
    if args.render_source == 'ply':
        '''
        Moved main_z_height_constant_new(fpath) to geon_analysis folder.
        '''
        raise NotImplementedError("args.render_source {} is not defined.".format(args.render_source))
    else:
        raise NotImplementedError("args.render_source {} is not defined.".format(args.render_source))