# Copyright 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.

from __future__ import print_function
import tempfile
import os
import json
import argparse
import random
import math
import subprocess
import numpy as np
import restore_util
import get_b_box
import multiprocessing as mp
from collections import Counter
from datetime import datetime as dt
import sys
sys.path.insert(0, '/usr/local/lib/python')
sys.path.insert(0, '/usr/local/lib/python3.5/site-packages')
# import pycocotools.mask as maskUtils

"""
Renders random scenes using Blender, each with with a random number of objects;
each object has a random size, position, color, and shape. Objects will be
nonintersecting but may partially occlude each other. Output images will be
written to disk as PNGs, and we will also write a JSON file for each image with
ground-truth scene information.

This file expects to be run from Blender like this:

blender --background --python render_images.py -- [arguments to this script]
"""

INSIDE_BLENDER = True
try:
    import bpy
    import bpy_extras
    from mathutils import Vector
except ImportError as e:
    INSIDE_BLENDER = False
if INSIDE_BLENDER:
    try:
        import utils
        from utils import RetrySceneException
    except ImportError as e:
        print("\nERROR")
        print("Running render_images.py from Blender and cannot import utils.py.")
        print("You may need to add a .pth file to the site-packages of Blender's")
        print("bundled python with a command like this:\n")
        print("echo $PWD >> $BLENDER/$VERSION/python/lib/python3.5/site-packages/clevr.pth")
        print("\nWhere $BLENDER is the directory where Blender is installed, and")
        print("$VERSION is your Blender version (such as 2.78).")
        sys.exit(1)

parser = argparse.ArgumentParser()

# Input options
parser.add_argument('--base_scene_blendfile', default='data/base_scene.blend',
                    help="Base blender file on which all scenes are based; includes " +
                    "ground plane, lights, and camera.")
parser.add_argument('--properties_json', default='data/properties.json',
                    help="JSON file defining objects, materials, sizes, and colors. " +
                    "The \"colors\" field maps from CLEVR color names to RGB values; " +
                    "The \"sizes\" field maps from CLEVR size names to scalars used to " +
                    "rescale object models; the \"materials\" and \"shapes\" fields map " +
                    "from CLEVR material and shape names to .blend files in the " +
                    "--object_material_dir and --shape_dir directories respectively.")
parser.add_argument('--shape_dir', default='data/shapes',
                    help="Directory where .blend files for object models are stored")
parser.add_argument('--material_dir', default='data/materials',
                    help="Directory where .blend files for materials are stored")
parser.add_argument('--conf_class_combos_json', default=None, required=True,
                    help="Path to a JSON file mapping obj attributes to a list of " +
                    "dicts. This specifies the attribute combinations that are required to be in the image.")
parser.add_argument('--gt_class_combos_json', default=None, required=True,
                    help="Path to a JSON file mapping obj attributes to a list of " +
                    "dicts. This specifies the ground-truth attribute combinations that are required to be in the image.")
parser.add_argument('--img_class_id', type=int, required=True,
                    help="Class label of the generated images")

# Settings for objects
parser.add_argument('--min_objects', default=3, type=int,
                    help="The minimum number of objects to place in each scene")
parser.add_argument('--max_objects', default=10, type=int,
                    help="The maximum number of objects to place in each scene")
parser.add_argument('--min_dist', default=0.25, type=float,
                    help="The minimum allowed distance between object centers")
parser.add_argument('--margin', default=0.4, type=float,
                    help="Along all cardinal directions (left, right, front, back), all " +
                    "objects will be at least this distance apart. This makes resolving " +
                    "spatial relationships slightly less ambiguous.")
parser.add_argument('--min_pixels_per_object', default=200, type=int,
                    help="All objects will have at least this many visible pixels in the " +
                    "final rendered images; this ensures that no objects are fully " +
                    "occluded by other objects.")
parser.add_argument('--max_retries', default=50, type=int,
                    help="The number of times to try placing an object before giving up and " +
                    "re-placing all objects in the scene.")
parser.add_argument('--max_img_retries', default=3, type=int,
                    help="The number of times to try with a new scene, before reducing the number of objects.")

# Output settings
parser.add_argument('--start_idx', default=0, type=int,
                    help="The index at which to start for numbering rendered images. Setting " +
                    "this to non-zero values allows you to distribute rendering across " +
                    "multiple machines and recombine the results later.")
parser.add_argument('--num_images', default=5, type=int,
                    help="The number of images to render")
parser.add_argument('--filename_prefix', default='CLEVR',
                    help="This prefix will be prepended to the rendered images and JSON scenes")
parser.add_argument('--split', default='new',
                    help="Name of the split for which we are rendering. This will be added to " +
                    "the names of rendered images, and will also be stored in the JSON " +
                    "scene structure for each image.")
parser.add_argument('--output_image_dir', default='../output/images/',
                    help="The directory where output images will be stored. It will be " +
                    "created if it does not exist.")
parser.add_argument('--output_scene_dir', default='../output/scenes/',
                    help="The directory where output JSON scene structures will be stored. " +
                    "It will be created if it does not exist.")
parser.add_argument('--output_scene_file', default='../output/CLEVR_scenes.json',
                    help="Path to write a single JSON file containing all scene information")
parser.add_argument('--output_blend_dir', default='output/blendfiles',
                    help="The directory where blender scene files will be stored, if the " +
                    "user requested that these files be saved using the " +
                    "--save_blendfiles flag; in this case it will be created if it does " +
                    "not already exist.")
parser.add_argument('--save_blendfiles', type=int, default=0,
                    help="Setting --save_blendfiles 1 will cause the blender scene file for " +
                    "each generated image to be stored in the directory specified by " +
                    "the --output_blend_dir flag. These files are not saved by default " +
                    "because they take up ~5-10MB each.")
parser.add_argument('--version', default='1.0',
                    help="String to store in the \"version\" field of the generated JSON file")
parser.add_argument('--license',
                    default="Creative Commons Attribution (CC-BY 4.0)",
                    help="String to store in the \"license\" field of the generated JSON file")
parser.add_argument('--date', default=dt.today().strftime("%m/%d/%Y"),
                    help="String to store in the \"date\" field of the generated JSON file; " +
                    "defaults to today's date")

# Rendering options
parser.add_argument('--use_gpu', default=0, type=int,
                    help="Setting --use_gpu 1 enables GPU-accelerated rendering using CUDA. " +
                    "You must have an NVIDIA GPU with the CUDA toolkit installed for " +
                    "to work.")
parser.add_argument('--width', default=320, type=int,
                    help="The width (in pixels) for the rendered images")
parser.add_argument('--height', default=240, type=int,
                    help="The height (in pixels) for the rendered images")
parser.add_argument('--key_light_jitter', default=1.0, type=float,
                    help="The magnitude of random jitter to add to the key light position.")
parser.add_argument('--fill_light_jitter', default=1.0, type=float,
                    help="The magnitude of random jitter to add to the fill light position.")
parser.add_argument('--back_light_jitter', default=1.0, type=float,
                    help="The magnitude of random jitter to add to the back light position.")
parser.add_argument('--camera_jitter', default=0.5, type=float,
                    help="The magnitude of random jitter to add to the camera position")
parser.add_argument('--render_num_samples', default=512, type=int,
                    help="The number of samples to use when rendering. Larger values will " +
                    "result in nicer images but will cause rendering to take longer.")
parser.add_argument('--render_min_bounces', default=8, type=int,
                    help="The minimum number of bounces to use for rendering.")
parser.add_argument('--render_max_bounces', default=8, type=int,
                    help="The maximum number of bounces to use for rendering.")
parser.add_argument('--render_tile_size', default=256, type=int,
                    help="The tile size to use for rendering. This should not affect the " +
                    "quality of the rendered image but may affect the speed; CPU-based " +
                    "rendering may achieve better performance using smaller tile sizes " +
                    "while larger tile sizes may be optimal for GPU-based rendering.")
parser.add_argument('--num_parallel_threads', default=1, type=int,
                    help="The number of threads to be run in parallel to generate the images")


def render_single_scene(i, img_template, scene_template, blend_template):
    img_path = img_template % (i + args.start_idx)
    scene_path = scene_template % (i + args.start_idx)
    blend_path = None
    if args.save_blendfiles == 1:
        blend_path = blend_template % (i + args.start_idx)
    num_objects = random.randint(args.min_objects, args.max_objects)
    render_scene(args,
                 num_objects=num_objects,
                 output_index=(i + args.start_idx),
                 output_split=args.split,
                 output_image=img_path,
                 output_scene=scene_path,
                 output_blendfile=blend_path,
                 )
    return scene_path


def main(args):
    num_digits = 6
    args.split = 'classid_'+str(args.img_class_id)
    prefix = '%s_%s_' % (args.filename_prefix, args.split)
    img_template = '%s%%0%dd.png' % (prefix, num_digits)
    scene_template = '%s%%0%dd.json' % (prefix, num_digits)
    blend_template = '%s%%0%dd.blend' % (prefix, num_digits)
    img_template = os.path.join(args.output_image_dir, img_template)
    scene_template = os.path.join(args.output_scene_dir, scene_template)
    blend_template = os.path.join(args.output_blend_dir, blend_template)

    if not os.path.isdir(args.output_image_dir):
        os.makedirs(args.output_image_dir)
    if not os.path.isdir(args.output_scene_dir):
        os.makedirs(args.output_scene_dir)
    if args.save_blendfiles == 1 and not os.path.isdir(args.output_blend_dir):
        os.makedirs(args.output_blend_dir)

    pool = mp.Pool(args.num_parallel_threads)
    # launching multiple evaluations asynchronously *may* use more processes
    multiple_results = [pool.apply_async(render_single_scene,
                                         args=(i, img_template, scene_template, blend_template))
                        for i in range(args.num_images)]

    all_scene_paths = [res.get() for res in multiple_results]
    all_scene_paths.sort()

    # After rendering all images, combine the JSON files for each scene into a
    # single JSON file.
    all_scenes = []
    for scene_path in all_scene_paths:
        with open(scene_path, 'r') as f:
            all_scenes.append(json.load(f))
    output = {
        'info': {
            'date': args.date,
            'version': args.version,
            'split': args.split,
            'license': args.license,
        },
        'scenes': all_scenes
    }

    json_pth = args.output_scene_file.split(
        '.json')[0]+'_classid_'+str(args.img_class_id)+'.json'
    with open(json_pth, 'w') as f:
        json.dump(output, f)


def render_scene(args,
                 num_objects=5,
                 output_index=0,
                 output_split='none',
                 output_image='render.png',
                 output_scene='render_json',
                 output_blendfile=None,
                 ):
    # Load the main blendfile
    bpy.ops.wm.open_mainfile(filepath=args.base_scene_blendfile)

    # Load materials
    utils.load_materials(args.material_dir)

    # Set render arguments so we can get pixel coordinates later.
    # We use functionality specific to the CYCLES renderer so BLENDER_RENDER
    # cannot be used.
    render_args = bpy.context.scene.render
    render_args.engine = "CYCLES"
    render_args.filepath = output_image
    render_args.resolution_x = args.width
    render_args.resolution_y = args.height
    render_args.resolution_percentage = 100
    render_args.tile_x = args.render_tile_size
    render_args.tile_y = args.render_tile_size
    if args.use_gpu == 1:
        # Blender changed the API for enabling CUDA at some point
        if bpy.app.version < (2, 78, 0):
            bpy.context.user_preferences.system.compute_device_type = 'CUDA'
            bpy.context.user_preferences.system.compute_device = 'CUDA_0'
        else:
            cycles_prefs = bpy.context.user_preferences.addons['cycles'].preferences
            cycles_prefs.compute_device_type = 'CUDA'

    # Some CYCLES-specific stuff
    bpy.data.worlds['World'].cycles.sample_as_light = True
    bpy.context.scene.cycles.blur_glossy = 2.0
    bpy.context.scene.cycles.samples = args.render_num_samples
    bpy.context.scene.cycles.transparent_min_bounces = args.render_min_bounces
    bpy.context.scene.cycles.transparent_max_bounces = args.render_max_bounces
    if args.use_gpu == 1:
        bpy.context.scene.cycles.device = 'GPU'

    # This will give ground-truth information about the scene and its objects
    scene_struct = {
        'split': output_split,
        'image_index': output_index,
        'image_filename': os.path.basename(output_image),
        'objects': [],
        'directions': {},
    }

    # Put a plane on the ground so we can compute cardinal directions
    bpy.ops.mesh.primitive_plane_add(radius=5)
    plane = bpy.context.object

    def rand(L):
        return 2.0 * L * (random.random() - 0.5)

    # Add random jitter to camera position
    if args.camera_jitter > 0:
        for i in range(3):
            bpy.data.objects['Camera'].location[i] += rand(args.camera_jitter)

    # Figure out the left, up, and behind directions along the plane and record
    # them in the scene structure
    camera = bpy.data.objects['Camera']
    plane_normal = plane.data.vertices[0].normal
    cam_behind = camera.matrix_world.to_quaternion() * Vector((0, 0, -1))
    cam_left = camera.matrix_world.to_quaternion() * Vector((-1, 0, 0))
    cam_up = camera.matrix_world.to_quaternion() * Vector((0, 1, 0))
    plane_behind = (cam_behind - cam_behind.project(plane_normal)).normalized()
    plane_left = (cam_left - cam_left.project(plane_normal)).normalized()
    plane_up = cam_up.project(plane_normal).normalized()

    # Delete the plane; we only used it for normals anyway. The base scene file
    # contains the actual ground plane.
    utils.delete_object(plane)

    # Save all six axis-aligned directions in the scene struct
    scene_struct['directions']['behind'] = tuple(plane_behind)
    scene_struct['directions']['front'] = tuple(-plane_behind)
    scene_struct['directions']['left'] = tuple(plane_left)
    scene_struct['directions']['right'] = tuple(-plane_left)
    scene_struct['directions']['above'] = tuple(plane_up)
    scene_struct['directions']['below'] = tuple(-plane_up)

    # Add random jitter to lamp positions
    if args.key_light_jitter > 0:
        for i in range(3):
            bpy.data.objects['Lamp_Key'].location[i] += rand(
                args.key_light_jitter)
    if args.back_light_jitter > 0:
        for i in range(3):
            bpy.data.objects['Lamp_Back'].location[i] += rand(
                args.back_light_jitter)
    if args.fill_light_jitter > 0:
        for i in range(3):
            bpy.data.objects['Lamp_Fill'].location[i] += rand(
                args.fill_light_jitter)

    # Now make some random objects, catch if the scene must be recreated from scratch, because too many placement
    # tries have been done
    img_retries = 0
    while True:
        img_retries += 1
        try:
            objects, blender_objects = add_random_objects(
                scene_struct, num_objects, args, camera)
        except RetrySceneException:
            if img_retries > args.max_img_retries:
                num_objects = random.randint(args.min_objects, args.max_objects)
                print("\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\nMaking new scene with {} objects!\n"
                  "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n".format(num_objects))
            # img_retries += 1
            # if img_retries > args.max_img_retries and num_objects > args.min_objects:
            #     num_objects -= 1
            # print("\nNumber_objects {}\n".format(num_objects))
            continue
        break

    for i in range(len(objects)):
        objects[i]['idx'] = i + 1

    print("\nScene for class {} created with {} objects\n".format(args.img_class_id, len(objects)))

    # get b_box
    box_dict = get_b_box.main(bpy.context, blender_objects)

    def build_rendermask_graph(num_obj):
        # 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

        scene = bpy.context.scene
        nodes = scene.node_tree.nodes

        render_layers = nodes['Render Layers']

        ofile_nodes = [nodes.new("CompositorNodeOutputFile")
                       for _ in range(num_obj)]
        for _i, of_node in enumerate(ofile_nodes):
            of_node.base_path = "./tmp_graph_output/indexob{}_{}".format(
                _i, output_index)

        idmask_list = [nodes.new("CompositorNodeIDMask")
                       for _ in range(num_obj)]
        for _i, o_node in enumerate(idmask_list):
            o_node.index = _i + 1

        bpy.data.scenes['Scene'].render.layers['RenderLayer'].use_pass_object_index = True

        for _i in range(num_obj):
            scene.node_tree.links.new(
                render_layers.outputs['IndexOB'],
                idmask_list[_i].inputs[0]
            )
            scene.node_tree.links.new(
                idmask_list[_i].outputs[0],
                ofile_nodes[_i].inputs['Image']
            )

    def get_diff_obj_points():
        obj_index = dict()
        index = 0
        for obj in blender_objects:
            index += 1
            obj.pass_index = index
            obj_index[obj.name] = index
        return index

    index = get_diff_obj_points()
    build_rendermask_graph(index)

    # Render the scene and dump the scene data structure
    scene_struct['objects'] = objects
    scene_struct['relationships'] = compute_all_relationships(scene_struct)
    while True:
        try:
            bpy.ops.render.render(write_still=True)
            break
        except Exception as e:
            print(e)

    # save_as_json
    cmd = ['python', './restore_img2json.py', str(index), str(output_index)]
    res = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    res.wait()
    if res.returncode != 0:
        print("  os.wait:exit status != 0\n")
        result = res.stdout.read()
        print("after read: {}".format(result))
        raise Exception('error in img2json')

    obj_mask = json.load(open('/tmp/obj_mask_{}.json'.format(output_index)))
    _path = '/tmp/obj_mask_{}.json'.format(output_index)
    os.system('rm ' + _path)

    for i in range(len(objects)):
        objects[i]['bbox'] = box_dict[i+1]
        # get uncompressed rle
        objects[i]['mask'] = obj_mask[str(i+1)]

    # add class id
    scene_struct['class_id'] = args.img_class_id

    with open(output_scene, 'w') as f:
        json.dump(scene_struct, f, indent=2)

    if output_blendfile is not None:
        bpy.ops.wm.save_as_mainfile(filepath=output_blendfile)



def place_obj(scene_struct, args,
              cand_obj, cur_objs_attr_dict, blender_objects, positions,
              relationship_constraints, general_position, r):
    """
    Place a candidate object, making sure it is far enough from the borders and other objects.
    :param scene_struct:
    :param blender_objects:
    :param positions:
    :param general_position:
    :param r:
    :param args:
    :return:
    """
    # Try to place the object, ensuring that we don't intersect any existing
    # objects and that we are more than the desired margin away from all existing
    # objects along all cardinal directions.
    num_tries = 0
    while True:
        # first make sure higher level constraints hold
        relationship_constraint_hold = False
        while not relationship_constraint_hold:
            # If we try and fail to place an object too many times, then delete all
            # the objects in the scene and send message to wrapper.
            num_tries += 1
            if num_tries > args.max_retries:
                for obj in blender_objects:
                    utils.delete_object(obj)
                raise RetrySceneException()
                # print("\nRetrying on new image!!\n")
                # return add_random_objects(scene_struct, num_objects, args, camera)

            # first catch general placement constraints
            if general_position:
                x, y = utils.sample_obj_place_with_general_constraint(general_position, r, scene_struct)
            else:
                x, y = utils.sample_obj_place_with_general_constraint(None, r, scene_struct)
            # next check if position holds for any given relationship constraints
            relationship_constraint_hold = utils.check_relationship_constraints(x, y, r, scene_struct, cand_obj,
                                                                          cur_objs_attr_dict, relationship_constraints,
                                                                          cur_class_id=args.img_class_id)
        print("Constraints hold {}!".format(num_tries))

        # Check to make sure the new object is further than min_dist from all
        # other objects, and further than margin along the four cardinal directions
        dists_good = True
        margins_good = True
        for (xx, yy, rr) in positions:
            dx, dy = x - xx, y - yy
            dist = math.sqrt(dx * dx + dy * dy)
            if dist - r - rr < args.min_dist:
                dists_good = False
                break
            for direction_name in ['left', 'right', 'front', 'behind']:
                direction_vec = scene_struct['directions'][direction_name]
                assert direction_vec[2] == 0
                margin = dx * direction_vec[0] + dy * direction_vec[1]
                if 0 < margin < args.margin:
                    print(margin, args.margin, direction_name)
                    print('BROKEN MARGIN!')
                    margins_good = False
                    break
            if not margins_good:
                break

        if dists_good and margins_good:
            break

    general_position_name = utils.get_general_obj_placement_name(x, y, r, scene_struct)
    x_corrected, y_corrected = utils.compute_corrected_x_y(x, y, r, scene_struct)

    return x, y, general_position_name, x_corrected, y_corrected



def add_random_objects(scene_struct, num_objects, args, camera):
    """
    Add random objects to the current blender scene
    """

    # Load the property file
    with open(args.properties_json, 'r') as f:
        properties = json.load(f)
        color_name_to_rgba = {}
        for name, rgb in properties['colors'].items():
            rgba = [float(c) / 255.0 for c in rgb] + [1.0]
            color_name_to_rgba[name] = rgba
        material_mapping = [(v, k) for k, v in properties['materials'].items()]
        object_mapping = [(v, k) for k, v in properties['shapes'].items()]
        size_mapping = list(properties['sizes'].items())

    # Load attribute combination dictionaries for gt class examples
    with open(args.gt_class_combos_json, 'r') as f:
        gt_combos_all_classes, relationship_constraints = utils.convert_json_dict(f)

    # Load attribute combination dictionaries for confounding class examples
    with open(args.conf_class_combos_json, 'r') as f:
        conf_combos_all_classes, _ = utils.convert_json_dict(f)

    # catch for classes with variations of one rule
    conf_combos_all_classes, gt_combos_all_classes= utils.check_for_classes_with_probabilities(conf_combos_all_classes,
                                                                                               gt_combos_all_classes,
                                                                                               args.img_class_id)

    # attributes that have to be in the image based on current class id
    must_attr_combos = conf_combos_all_classes[args.img_class_id]
    num_must_attr_combos = len(must_attr_combos)
    # if more objects required for class rule
    if num_must_attr_combos > args.min_objects:
        num_objects = random.randint(num_must_attr_combos, args.max_objects)

    positions = []
    objects = []
    blender_objects = []
    cur_objs_attr_dict = {}  # list of lists of the attributes of the current objects
    # cur_objs_attr_dict = {2: {"color": "red", "size": "large", "shape": "cube", "material":"metal", "general_position": "top right", "corr_position": (-2, 1)},
    #                       3: {"color": "red", "size": "large", "shape": "cube", "material":"metal", "general_position": "top right", "corr_position": (-1, 0)}}  # list of lists of the attributes of the current objects
    for obj_id in range(num_objects):
        # enforce that first objects to be placed are the must be placed objects
        if obj_id < num_must_attr_combos:
            must_attr_combo = must_attr_combos[obj_id]

            # if any attr type is an empty string, choose randomly
            if must_attr_combo["shape"]:
                obj_name_out = must_attr_combo["shape"]
                obj_name = [
                    k for k, v in object_mapping if v == obj_name_out][0]
            else:
                obj_name, obj_name_out = random.choice(object_mapping)

            if must_attr_combo["size"]:
                size_name = must_attr_combo["size"]
                r = [r for k, r in size_mapping if k == size_name][0]
            else:
                size_name, r = random.choice(size_mapping)

            if must_attr_combo["material"]:
                mat_name_out = must_attr_combo["material"]
                mat_name = [
                    k for k, v in material_mapping if v == mat_name_out][0]
            else:
                mat_name, mat_name_out = random.choice(material_mapping)

            if must_attr_combo["color"]:
                color_name = must_attr_combo["color"]
                rgba = color_name_to_rgba[color_name]
            else:
                color_name, rgba = random.choice(
                    list(color_name_to_rgba.items()))

            cand_obj = {"shape": obj_name_out, "size": size_name, "material": mat_name_out, "color": color_name}

            x, y, general_position_name, x_corrected, y_corrected = place_obj(scene_struct,
                                                                                    args, cand_obj,
                                                                                    cur_objs_attr_dict, blender_objects,
                                                                                    positions, relationship_constraints,
                                                                                    general_position=
                                                                                    must_attr_combo[
                                                                                        "general_position"], r=r)

        else:
            conflict = True

            # randomly choose an object, s.t. it is not in conflict with must attr combos of the
            # other confounding classes
            num_tries = 0
            while(conflict):
                num_tries += 1
                if num_tries > args.max_retries:
                    print("Too often conflict with obj attributes")
                    for obj in blender_objects:
                        utils.delete_object(obj)
                    # return add_random_objects(scene_struct, num_objects, args, camera)
                    raise RetrySceneException

                # Choose a random size and material
                size_name, r = random.choice(size_mapping)
                mat_name, mat_name_out = random.choice(material_mapping)

                # Choose random color and shape
                obj_name, obj_name_out = random.choice(object_mapping)
                color_name, rgba = random.choice(
                    list(color_name_to_rgba.items()))

                cand_obj = {"shape": obj_name_out, "size": size_name, "material": mat_name_out, "color": color_name}

                x, y, general_position_name, x_corrected, y_corrected = place_obj(scene_struct,
                                                                                        args, cand_obj,
                                                                                        cur_objs_attr_dict,
                                                                                        blender_objects,
                                                                                        positions,
                                                                                        relationship_constraints,
                                                                                        general_position="", r=r)

                # add candidate obj to temporary list
                tmp_obj_attr_dict = cur_objs_attr_dict.copy()
                tmp_obj_attr_dict[obj_id] = {"shape": obj_name_out, "size": size_name,
                                             "material": mat_name_out, "color": color_name,
                                             "general_position": general_position_name}

                conflict = utils.check_conflicts_all_classes(gt_combos_all_classes,
                                                             tmp_obj_attr_dict,
                                                             args.img_class_id)
                print("Conflict concerning attributes? {}".format(conflict))

        print("\nObject {} of CLass {} has been placed\n".format(obj_id, args.img_class_id))

        cur_objs_attr_dict[obj_id] = {"shape": obj_name_out, "size": size_name,
                                     "material": mat_name_out, "color": color_name,
                                     "general_position": general_position_name,
                                     "corr_position": (x_corrected, y_corrected)}

        # For cube, adjust the size a bit
        if obj_name == 'Cube':
            r /= math.sqrt(2)

        # Choose random orientation for the object.
        theta = 360.0 * random.random()

        # Actually add the object to the scene
        utils.add_object(args.shape_dir, obj_name, r, (x, y), theta=theta)
        obj = bpy.context.object
        blender_objects.append(obj)
        positions.append((x, y, r))

        # Attach material
        utils.add_material(mat_name, Color=rgba)

        # Record data about the object in the scene data structure
        pixel_coords = utils.get_camera_coords(camera, obj.location)
        objects.append({
            'shape': obj_name_out,
            'size': size_name,
            'material': mat_name_out,
            '3d_coords': tuple(obj.location),
            'rotation': theta,
            'pixel_coords': pixel_coords,
            'color': color_name,
        })

    # Check that all objects are at least partially visible in the rendered image
    all_visible = check_visibility(blender_objects, args.min_pixels_per_object)
    if not all_visible:
        # If any of the objects are fully occluded then start over; delete all
        # objects from the scene and place them all again.
        print('Some objects are occluded; replacing objects')
        for obj in blender_objects:
            utils.delete_object(obj)
        # return add_random_objects(scene_struct, num_objects, args, camera)
        raise RetrySceneException

    return objects, blender_objects


def compute_all_relationships(scene_struct, eps=0.2):
    """
    Computes relationships between all pairs of objects in the scene.

    Returns a dictionary mapping string relationship names to lists of lists of
    integers, where output[rel][i] gives a list of object indices that have the
    relationship rel with object i. For example if j is in output['left'][i] then
    object j is left of object i.
    """
    all_relationships = {}
    for name, direction_vec in scene_struct['directions'].items():
        if name == 'above' or name == 'below':
            continue
        all_relationships[name] = []
        for i, obj1 in enumerate(scene_struct['objects']):
            coords1 = obj1['3d_coords']
            related = set()
            for j, obj2 in enumerate(scene_struct['objects']):
                if obj1 == obj2:
                    continue
                coords2 = obj2['3d_coords']
                diff = [coords2[k] - coords1[k] for k in [0, 1, 2]]
                dot = sum(diff[k] * direction_vec[k] for k in [0, 1, 2])
                if dot > eps:
                    related.add(j)
            all_relationships[name].append(sorted(list(related)))
    return all_relationships


def check_visibility(blender_objects, min_pixels_per_object):
    """
    Check whether all objects in the scene have some minimum number of visible
    pixels; to accomplish this we assign random (but distinct) colors to all
    objects, and render using no lighting or shading or antialiasing; this
    ensures that each object is just a solid uniform color. We can then count
    the number of pixels of each color in the output image to check the visibility
    of each object.

    Returns True if all objects are visible and False otherwise.
    """
    f, path = tempfile.mkstemp(suffix='.png')
    object_colors = render_shadeless(blender_objects, path=path)
    img = bpy.data.images.load(path)
    p = list(img.pixels)
    color_count = Counter((p[i], p[i + 1], p[i + 2], p[i + 3])
                          for i in range(0, len(p), 4))
    os.remove(path)
    if len(color_count) != len(blender_objects) + 1:
        return False
    for _, count in color_count.most_common():
        if count < min_pixels_per_object:
            return False
    return True


def render_shadeless(blender_objects, path='flat.png'):
    """
    Render a version of the scene with shading disabled and unique materials
    assigned to all objects, and return a set of all colors that should be in the
    rendered image. The image itself is written to path. This is used to ensure
    that all objects will be visible in the final rendered scene.
    """
    render_args = bpy.context.scene.render

    # Cache the render args we are about to clobber
    old_filepath = render_args.filepath
    old_engine = render_args.engine
    old_use_antialiasing = render_args.use_antialiasing

    # Override some render settings to have flat shading
    render_args.filepath = path
    render_args.engine = 'BLENDER_RENDER'
    render_args.use_antialiasing = False

    # Move the lights and ground to layer 2 so they don't render
    utils.set_layer(bpy.data.objects['Lamp_Key'], 2)
    utils.set_layer(bpy.data.objects['Lamp_Fill'], 2)
    utils.set_layer(bpy.data.objects['Lamp_Back'], 2)
    utils.set_layer(bpy.data.objects['Ground'], 2)

    # Add random shadeless materials to all objects
    object_colors = set()
    old_materials = []
    for i, obj in enumerate(blender_objects):
        old_materials.append(obj.data.materials[0])
        bpy.ops.material.new()
        mat = bpy.data.materials['Material']
        mat.name = 'Material_%d' % i
        while True:
            r, g, b = [random.random() for _ in range(3)]
            if (r, g, b) not in object_colors:
                break
        object_colors.add((r, g, b))
        mat.diffuse_color = [r, g, b]
        mat.use_shadeless = True
        obj.data.materials[0] = mat

    # Render the scene
    bpy.ops.render.render(write_still=True)

    # Undo the above; first restore the materials to objects
    for mat, obj in zip(old_materials, blender_objects):
        obj.data.materials[0] = mat

    # Move the lights and ground back to layer 0
    utils.set_layer(bpy.data.objects['Lamp_Key'], 0)
    utils.set_layer(bpy.data.objects['Lamp_Fill'], 0)
    utils.set_layer(bpy.data.objects['Lamp_Back'], 0)
    utils.set_layer(bpy.data.objects['Ground'], 0)

    # Set the render settings back to what they were
    render_args.filepath = old_filepath
    render_args.engine = old_engine
    render_args.use_antialiasing = old_use_antialiasing

    return object_colors


if __name__ == '__main__':
    if INSIDE_BLENDER:
        # Run normally
        argv = utils.extract_args()
        args = parser.parse_args(argv)
        main(args)
    elif '--help' in sys.argv or '-h' in sys.argv:
        parser.print_help()
    else:
        print('This script is intended to be called from blender like this:')
        print()
        print('blender --background --python render_images.py -- [args]')
        print()
        print('You can also run as a standalone python script to view all')
        print('arguments like this:')
        print()
        print('python render_images.py --help')
