import math, sys, random, argparse, json, os, tempfile
from datetime import datetime as dt
from collections import Counter
import itertools
import numpy as np


INSIDE_BLENDER = True
try:
  import bpy, bpy_extras
  from mathutils import Vector
except ImportError as e:
  INSIDE_BLENDER = False
if INSIDE_BLENDER:
  try:
    import utils
  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 >> /Applications/Blender.app/Contents/Resources/3.5/python/lib/python3.10/site-packages/clevr.pth")
    sys.exit(1)

parser = argparse.ArgumentParser()

# Input options
parser.add_argument('--shape_dir', default='input/shapes',
    help="Directory where .blend files for object models are stored")
parser.add_argument('--material_dir', default='input/materials',
    help="Directory where .blend files for materials are stored")

# Settings for objects
parser.add_argument('--dataset', default='single-body_2d_3classes')
#parser.add_argument('--dataset', default='single-body_3d_3classes')
#parser.add_argument('--dataset', default='single-body_2d_3classes_cont')
#parser.add_argument('--dataset', default='single-body_2d_3classes_fail')
#parser.add_argument('--dataset', default='single-body_2d_4classes')
#parser.add_argument('--dataset', default='two-body_2d_3classes')
#parser.add_argument('--dataset', default='two-body_2d_3classes_colored')
parser.add_argument('--noise', default=0.05, type=float)

# Output settings
parser.add_argument('--filename_prefix', default='CLEVR',
    help="This prefix will be prepended to the rendered images and JSON scenes")

# Rendering options
parser.add_argument('--use_gpu', default=0, type=int,
    help="Setting --use_gpu 1 enables GPU-accelerated rendering using CUDA. ")
parser.add_argument('--width', default=56, type=int,
    help="The width (in pixels) for the rendered images")
parser.add_argument('--height', default=56, type=int,
    help="The height (in pixels) for the rendered images")

def main(args):
  num_digits = 6
  properties_json = "input/properties_" + args.dataset + ".json"

  with open(properties_json, 'r') as f:
      properties = json.load(f)

  keys, values = zip(*properties.items())
  permutations0 = [dict(zip(keys, v)) for v in itertools.product(*values)]
  noise = args.noise
  

  for mode, niter in zip(["template_1.3"],[1]):
  #for mode, niter in zip(["test", "train"],[200, 1500]):

      all_scene_paths = []
      output_dir = "output/" + args.dataset #+ "_" + str(noise)
      if not os.path.isdir(output_dir):
          os.makedirs(output_dir)

      output_dir = output_dir + "/"+mode+"/"
      if not os.path.isdir(output_dir):
          os.makedirs(output_dir)

      if "train_" in mode: permutations = [{'shapes': '0', 'colors': '0', 'sizes': '0'}]
      elif "panda" in mode: permutations = [{'shapes': '1', 'colors': '0', 'sizes': '0'}]
      else: permutations = permutations0

      i = 1
      for permutation in permutations:
          if "panda" in mode: 
              prefix = args.filename_prefix + "_" + "".join(permutation.values()) + "panda_"
          else:
              prefix = args.filename_prefix + "_" + "".join(permutation.values()) + "_"

          img_template =  '%s%%0%dd.png' % (prefix, num_digits)
          scene_template = '%s%%0%dd.json' % (prefix, num_digits)

          img_template = os.path.join(output_dir, img_template)
          scene_template = os.path.join(output_dir, scene_template)

          for _ in range(niter): 
              img_path = img_template % i
              scene_path = scene_template % i
              all_scene_paths.append(scene_path)
              if "two" in args.dataset: 
                  num_objects = 2
              else:
                  num_objects = 1
              render_scene(args,
                num_objects=num_objects,
                output_index=i,
                output_image=img_path,
                output_scene=scene_path,
                permutation=permutation,
                noise=noise, 
                mode=mode
              )
              if not "template" in mode: i += 1



def render_scene(args,
    num_objects=5,
    output_index=0,
    output_image='render.png',
    output_scene='render_json',
    permutation=None, 
    noise=0.05, 
    mode=None 
  ):

  if "2d" in args.dataset: 
      base_scene_blendfile = "input/my_base_scene_2d.blend"
  elif "3d" in args.dataset:
      base_scene_blendfile = "input/my_base_scene_3d.blend"
  bpy.ops.wm.open_mainfile(filepath=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 = 'BLENDER_EEVEE'
  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

  # Some CYCLES-specific stuff
  bpy.data.worlds['World'].cycles.sample_as_light = True
  bpy.context.scene.cycles.blur_glossy = 1.0
  bpy.context.scene.cycles.samples = 1 
  bpy.context.scene.cycles.transparent_min_bounces = 8 
  bpy.context.scene.cycles.transparent_max_bounces = 8 
  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 = {
      '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
  #bpy.ops.mesh.primitive_plane_add(location=(0,0,0),enter_editmode=False,rotation=(90,0,90))
  #plane = bpy.context.selected_objects[0]
  #plane.name = "ground"
  #plane.select_set(False)


  # 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)

  # Now make some random objects
  objects, blender_objects = add_random_objects(scene_struct, num_objects, args, camera, permutation, noise, mode)
  #bpy.ops.transform.rotate(value=0.1*np.pi, orient_axis='Y')
  #bpy.ops.transform.rotate(value=0.1*np.pi, orient_axis='X')
  #bpy.ops.transform.rotate(value=45, orient_axis='Z')


  # 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)


  if "two-body_2d_3classes" in args.dataset: 
      my_scene_struct = scene_struct['objects'][0]["distance"]
  elif "single-body_2d_4classes"==args.dataset:
      my_scene_struct = [scene_struct['objects'][0].get(key) for key in ["size","rgba","y"]]
  elif "single-body_2d_3classes_position"==args.dataset:
      my_scene_struct = [scene_struct['objects'][0].get(key) for key in ["y","rgba","theta"]]
  else:
      #my_scene_struct = [scene_struct['objects'][0].get(key) for key in ["theta","size","rgba"]]
      my_scene_struct = [scene_struct['objects'][0].get(key) for key in ["size","rgba"]]
  with open(output_scene, 'w') as f:
    json.dump(my_scene_struct, f, indent=2)


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

  # Load the property file
  properties_json = "input/properties_" + args.dataset + ".json"

  with open(properties_json, 'r') as f:
    properties = json.load(f)
  """
    color_name_to_rgba = {}
    for name, rgb in properties['colors'].items():
      #rgba = [c / 255.0 for c in rgb] + [1.0]
      rgba = rgb + [1.0]
      color_name_to_rgba[name] = rgba
  """

  #y_top = 2.0 + random.uniform(-0.7, 0.7)
  #y_bottom = -2.0 + random.uniform(-0.7, 0.7)
  y_top = 1.5 + random.uniform(-1.2, 1.2)
  y_bottom = -1.5 + random.uniform(-1.2, 1.2)

  positions = []
  objects = []
  blender_objects = []
  for i in range(num_objects):

    if "single-body_2d_3classes"==args.dataset:
        size_name = permutation["sizes"]
        if "template_" in mode or mode=="train_001":
            if size_name=="0": r = 2.6
            if size_name=="1": r = float(mode.split("_")[-1]) 
        else:
            r = properties["sizes"][size_name] + random.uniform(-0.5,0.5)
    elif "two-body_2d_3classes" in args.dataset: 
        r = 1.0 
    #elif "3d" in args.dataset:
    #    r = 1.3 
    elif "single-body_2d_4classes" in args.dataset:
        size_name = permutation["sizes"]
        r = properties["sizes"][size_name] + random.uniform(-0.3,0.3)
    else:
        size_name = permutation["sizes"]
        r = properties["sizes"][size_name] + random.uniform(-noise,noise)


    if "two-body_2d_3classes" in args.dataset: 
       x = random.uniform(-noise, noise)
       distance_name = permutation["distances"]
       if distance_name=="0": 
           if i==0: y = y_top 
           if i==1: y = y_bottom 
       if distance_name=="1": 
           if i==0: y = y_bottom 
           if i==1: y = y_top 
    elif "single-body_2d_4classes" in args.dataset:
       x = 0.
       position_name = permutation["positions"]
       y = properties["positions"][position_name] + random.uniform(-0.7, 0.7)
    elif "single-body_2d_3classes_position" in args.dataset:
       x = 0.
       position_name = permutation["positions"]
       y = properties["positions"][position_name] + random.uniform(-1.2, 1.2)
    else: 
       x = 0.0
       y = 0.0

    if "two-body_2d_3classes"==args.dataset: 
        rgba = [0.0, 0.0, 0.0, 1.0]
        obj_name_out = permutation["shapes"+str(i)]
        obj_name = properties["shapes"+str(i)][obj_name_out] 
        color_name = "0"
    elif "two-body_2d_3classes_colored"==args.dataset: 
        obj_name_out = permutation["shapes"+str(i)]
        obj_name = properties["shapes"+str(i)][obj_name_out] 
        if obj_name_out=="0": rgba = [1.0, 0.0, 0.0, 1.0]; color_name = "0"
        if obj_name_out=="1": rgba = [0.0, 0.0, 1.0, 1.0]; color_name = "1"
        if obj_name_out=="2": rgba = [0.0, 1.0, 0.0, 1.0]; color_name = "2"
    else:
        if "single-body_2d_3classes"==args.dataset or "single-body_2d_4classes"==args.dataset or "single-body_2d_3classes_fail"==args.dataset or "single-body_3d_3classes"==args.dataset:  
            obj_name_out = permutation["shapes"]
            obj_name = properties["shapes"][obj_name_out] 
        else:
            obj_name = "Cone" 
        if "template_" in mode or mode=="train_100" or mode=="panda100":
            color_name = permutation["colors"]
            if color_name=="0": rgba = [ 0.9, 0.1, 0.1, 1.0 ]
            if color_name=="1": rgba = [ 0.1, 0.1, 0.9, 1.0 ]
            if color_name=="2": rgba = [ 0.1, 0.9, 0.1, 1.0 ]
        else:
            color_name = permutation["colors"]
            rgba = properties["colors"][color_name] + [1.0]
            rgba[0] += random.uniform(-noise,noise)
            rgba[1] += random.uniform(-noise,noise)
            rgba[2] += random.uniform(-noise,noise)


    #if "single-body_2d_4classes" in args.dataset:
    #   position_name = permutation["rotations"]
    #   theta = properties["rotations"][position_name] + random.uniform(-0.25, 0.25)
    if "single-body_2d_3classes_cont"==args.dataset:
       position_name = permutation["rotations"]
       theta = properties["rotations"][position_name] + random.uniform(-0.125, 0.125)
    elif "single-body_2d_3classes_disc"==args.dataset:  
       position_name = permutation["rotations"]
       theta = properties["rotations"][position_name] 
    else:
       if "_1" in obj_name: 
           theta = 0.5
       if "_2" in obj_name: 
           theta = 1.0
       else:
           theta = 0.

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

    # Attach a random material
    #mat_name, mat_name_out = random.choice(material_mapping)
    #utils.add_material(mat_name, Color=rgba)
    utils.add_material("Rubber", Color=rgba)

    # Record data about the object in the scene data structure
    pixel_coords = utils.get_camera_coords(camera, obj.location)
    if "two-body_2d_3classes" in args.dataset:
        if distance_name=="0": distance = y_top - y_bottom
        if distance_name=="1": distance = y_bottom - y_top
    else:
        distance = 0.
    objects.append({
      #'shape': obj_name_out,
      'size': r,
      'distance': distance,
      'x': x, 
      'y': y, 
      'theta': theta, 
      'color': color_name, 
      'rgba': rgba, 
      '3d_coords': tuple(obj.location)
    })


  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


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 my_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')

