from __future__ import print_function
import math, sys, random, argparse, json, os, tempfile
from datetime import datetime as dt
from collections import Counter
try:
    from PIL import Image, ImageFilter
except ImportError:
    Image = None
    ImageFilter = None

INSIDE_BLENDER = True
try:
  import bpy, bpy_extras
  from mathutils import Vector
except ImportError as e:
  INSIDE_BLENDER = False

def extract_args(input_argv=None):
  if input_argv is None:
    input_argv = sys.argv
  output_argv = []
  if '--' in input_argv:
    idx = input_argv.index('--')
    output_argv = input_argv[(idx + 1):]
  return output_argv

def parse_args(parser, argv=None):
  return parser.parse_args(extract_args(argv))

def delete_object(obj):
  if not INSIDE_BLENDER:
    return
  bpy.ops.object.select_all(action='DESELECT')
  obj.select_set(True)
  bpy.context.view_layer.objects.active = obj
  bpy.ops.object.delete()

def get_camera_coords(cam, pos):
  if not INSIDE_BLENDER:
    return (0, 0, 0)
  scene = bpy.context.scene
  x, y, z = bpy_extras.object_utils.world_to_camera_view(scene, cam, pos)
  scale = scene.render.resolution_percentage / 100.0
  w = int(scale * scene.render.resolution_x)
  h = int(scale * scene.render.resolution_y)
  px = int(round(x * w))
  py = int(round(h - y * h))
  return (px, py, z)

def set_layer(obj, layer_idx):
  if not INSIDE_BLENDER:
    return
  obj.layers[layer_idx] = True
  for i in range(len(obj.layers)):
    obj.layers[i] = (i == layer_idx)

def add_object(object_dir, name, scale, loc, theta=0):
  if not INSIDE_BLENDER:
    return
  object_dir = os.path.abspath(os.path.normpath(object_dir))
  if not os.path.exists(object_dir):
    print(f"ERROR: Object directory does not exist: {object_dir}")
    return
  count = 0
  for obj in bpy.data.objects:
    if obj.name.startswith(name):
      count += 1

  blend_file = os.path.join(object_dir, '%s.blend' % name)
  blend_file = os.path.abspath(blend_file).replace('\\', '/')
  if not os.path.exists(blend_file):
    print(f"ERROR: Blend file does not exist: {blend_file}")
    return
  directory = blend_file + '/Object/'
  try:
    bpy.ops.wm.append(
      directory=directory,
      filename=name,
      filter_blender=True
    )
  except Exception as e:
    error_msg = str(e)
    print(f"ERROR: Failed to load object {name} from {directory}: {error_msg}")
    print(f"  Error type: {type(e).__name__}")
    raise

  new_name = '%s_%d' % (name, count)
  bpy.data.objects[name].name = new_name

  x, y = loc
  bpy.context.view_layer.objects.active = bpy.data.objects[new_name]
  bpy.context.object.rotation_euler[2] = theta
  bpy.ops.transform.resize(value=(scale, scale, scale))
  bpy.ops.transform.translate(value=(x, y, scale))

def load_materials(material_dir):
  if not INSIDE_BLENDER:
    return
  material_dir = os.path.abspath(os.path.normpath(material_dir))
  if not os.path.exists(material_dir):
    print(f"ERROR: Material directory does not exist: {material_dir}")
    return
  for fn in os.listdir(material_dir):
    if not fn.endswith('.blend'): continue
    name = os.path.splitext(fn)[0]
    blend_file = os.path.join(material_dir, fn)
    blend_file = os.path.abspath(blend_file).replace('\\', '/')
    if not os.path.exists(blend_file):
      print(f"ERROR: Blend file does not exist: {blend_file}")
      continue
    directory = blend_file + '/NodeTree/'
    try:
      bpy.ops.wm.append(
        directory=directory,
        filename=name,
        filter_blender=True
      )
    except Exception as e:
      error_msg = str(e)
      print(f"ERROR: Failed to load material {name} from {directory}: {error_msg}")
      print(f"  Error type: {type(e).__name__}")
      raise

def apply_filter_to_image(image_path, filter_type, filter_strength):
  image_path = os.path.abspath(image_path)
  if Image is None:
    print(f"ERROR: PIL/Image not available, cannot apply filter {filter_type}")
    return
  if not os.path.exists(image_path):
    print(f"ERROR: Image file does not exist: {image_path}")
    return
  
  try:
    img = Image.open(image_path)
    
    if filter_type == 'blur':
      radius = max(1, int(filter_strength))
      img = img.filter(ImageFilter.GaussianBlur(radius=radius))
    
    elif filter_type == 'vignette':
      width, height = img.size
      center_x, center_y = width // 2, height // 2
      max_dist = math.sqrt(center_x**2 + center_y**2)
      
      img = img.convert('RGB')
      pixels = img.load()
      for y in range(height):
        for x in range(width):
          dist = math.sqrt((x - center_x)**2 + (y - center_y)**2)
          factor = 1.0 - (dist / max_dist) * (filter_strength / 5.0)
          factor = max(0.0, min(1.0, factor))
          
          r, g, b = pixels[x, y]
          pixels[x, y] = (int(r * factor), int(g * factor), int(b * factor))
    
    elif filter_type == 'fisheye':
      width, height = img.size
      center_x, center_y = width / 2.0, height / 2.0
      max_radius = min(center_x, center_y)
      
      img = img.convert('RGB')
      output = Image.new('RGB', (width, height))
      out_pixels = output.load()
      in_pixels = img.load()
      
      for y in range(height):
        for x in range(width):
          dx = (x - center_x) / max_radius
          dy = (y - center_y) / max_radius
          distance = math.sqrt(dx*dx + dy*dy)
          
          if distance > 1.0:
            out_pixels[x, y] = (0, 0, 0)
          else:
            theta = math.atan2(dy, dx)
            r_normalized = distance
            r_distorted = r_normalized * (1.0 + filter_strength * (1.0 - r_normalized))
            r_distorted = min(1.0, r_distorted)
            
            src_x = int(center_x + r_distorted * max_radius * math.cos(theta))
            src_y = int(center_y + r_distorted * max_radius * math.sin(theta))
            
            if 0 <= src_x < width and 0 <= src_y < height:
              out_pixels[x, y] = in_pixels[src_x, src_y]
            else:
              out_pixels[x, y] = (0, 0, 0)
      
      img = output
    
    img.save(image_path)
    print(f"[OK] Applied {filter_type} filter (strength: {filter_strength:.2f})")
  except Exception as e:
    import traceback
    print(f"ERROR applying filter {filter_type}: {e}")
    traceback.print_exc()
    raise

def add_material(name, **properties):
  if not INSIDE_BLENDER:
    return
  mat_count = len(bpy.data.materials)
  bpy.ops.material.new()
  mat = bpy.data.materials['Material']
  mat.name = 'Material_%d' % mat_count
  obj = bpy.context.active_object
  assert len(obj.data.materials) == 0
  obj.data.materials.append(mat)

  output_node = None
  for n in mat.node_tree.nodes:
    if n.name == 'Material Output':
      output_node = n
      break

  group_node = mat.node_tree.nodes.new('ShaderNodeGroup')
  group_node.node_tree = bpy.data.node_groups[name]

  for inp in group_node.inputs:
    if inp.name in properties:
      inp.default_value = properties[inp.name]

  mat.node_tree.links.new(
      group_node.outputs['Shader'],
      output_node.inputs['Surface'],
  )

parser = argparse.ArgumentParser()
parser.add_argument('--scene_file', default=None,
    help="Optional JSON file to load scene from. If provided, renders from JSON instead of generating random scenes.")
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('--shape_color_combos_json', default=None,
    help="Optional path to a JSON file mapping shape names to a list of " +
         "allowed color names for that shape. This allows rendering images " +
         "for CLEVR-CoGenT.")

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.15, type=float,
    help="The minimum allowed distance between object centers")
parser.add_argument('--margin', default=0.2, 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=50, 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=100, 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('--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")

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('--output_image', default=None,
    help="Output image path (used when rendering from JSON)")

MIN_VISIBLE_FRACTION = 0.001
MIN_VISIBLE_FRACTION_PARTIAL_OCCLUSION = 0.0005
MIN_PIXELS_FLOOR = 50

BASE_MIN_VISIBILITY_FRACTION = 0.9
CF_OCCLUSION_MIN_VISIBILITY_FRACTION = 0.3
CF_OCCLUSION_MAX_VISIBILITY_FRACTION = 0.5
CF_OCCLUSION_HARD_MIN_FRACTION = 0.2


def min_visible_pixels(width, height, fraction=MIN_VISIBLE_FRACTION, floor=MIN_PIXELS_FLOOR):
  return max(floor, int(width * height * fraction))


BACKGROUND_COLORS = {
    'default': None,
    'gray': (0.5, 0.5, 0.5),
    'blue': (0.2, 0.4, 0.8),
    'green': (0.2, 0.6, 0.3),
    'brown': (0.4, 0.3, 0.2),
    'purple': (0.5, 0.3, 0.6),
    'orange': (0.8, 0.5, 0.2),
    'white': (0.9, 0.9, 0.9),
    'dark_gray': (0.2, 0.2, 0.2),
    'red': (0.7, 0.2, 0.2),
    'yellow': (0.8, 0.8, 0.3),
    'cyan': (0.3, 0.7, 0.8),
}

LIGHTING_PRESETS = {
    'default': {'key': 1.0, 'fill': 0.5, 'back': 0.3},
    'bright': {'key': 12.0, 'fill': 6.0, 'back': 4.0},
    'dim': {'key': 0.008, 'fill': 0.004, 'back': 0.002},
    'warm': {'key': 5.0, 'fill': 0.8, 'back': 0.3, 'color': (1.0, 0.5, 0.2)},
    'cool': {'key': 4.0, 'fill': 2.0, 'back': 1.5, 'color': (0.2, 0.5, 1.0)},
    'dramatic': {'key': 15.0, 'fill': 0.005, 'back': 0.002},
}

def set_background_color(color_name):
  """Set the world background color"""
  if not INSIDE_BLENDER:
    return
  if color_name not in BACKGROUND_COLORS or BACKGROUND_COLORS[color_name] is None:
    return
  
  rgb = BACKGROUND_COLORS[color_name]
  
  world = bpy.context.scene.world
  if world is None:
    world = bpy.data.worlds.new("World")
    bpy.context.scene.world = world
  
  world.use_nodes = True
  nodes = world.node_tree.nodes
  
  bg_node = None
  for node in nodes:
    if node.type == 'BACKGROUND':
      bg_node = node
      break
  
  if bg_node is None:
    bg_node = nodes.new(type='ShaderNodeBackground')
  
  bg_node.inputs['Color'].default_value = (rgb[0], rgb[1], rgb[2], 1.0)
  bg_node.inputs['Strength'].default_value = 1.0
  
  print(f"Set background color to {color_name}: RGB{rgb}")

def set_ground_color(color_name):
  if not INSIDE_BLENDER:
    return
  if color_name not in BACKGROUND_COLORS or BACKGROUND_COLORS[color_name] is None:
    return
  
  rgb = BACKGROUND_COLORS[color_name]
  
  ground = None
  for obj in bpy.data.objects:
    if 'ground' in obj.name.lower() or 'plane' in obj.name.lower():
      ground = obj
      break
  
  if ground is None:
    return
  
  if len(ground.data.materials) == 0:
    mat = bpy.data.materials.new(name="Ground_Material")
    ground.data.materials.append(mat)
  else:
    mat = ground.data.materials[0]
  
  mat.use_nodes = True
  nodes = mat.node_tree.nodes
  
  bsdf = None
  for node in nodes:
    if node.type == 'BSDF_PRINCIPLED':
      bsdf = node
      break
  
  if bsdf:
    bsdf.inputs['Base Color'].default_value = (rgb[0], rgb[1], rgb[2], 1.0)
    print(f"Set ground color to {color_name}")

def set_lighting(lighting_name):
  """Set lighting conditions"""
  if not INSIDE_BLENDER:
    return
  if lighting_name not in LIGHTING_PRESETS:
    return
  
  preset = LIGHTING_PRESETS[lighting_name]
  
  lamp_names = ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back']
  intensity_keys = ['key', 'fill', 'back']
  
  for lamp_name, int_key in zip(lamp_names, intensity_keys):
    if lamp_name in bpy.data.objects:
      lamp_obj = bpy.data.objects[lamp_name]
      if lamp_obj.data and hasattr(lamp_obj.data, 'energy'):
        base_energy = lamp_obj.data.energy
        lamp_obj.data.energy = base_energy * preset.get(int_key, 1.0)
        
        if 'color' in preset and hasattr(lamp_obj.data, 'color'):
          lamp_obj.data.color = preset['color']
  
  print(f"Set lighting to {lighting_name}")

def render_from_json(args):
  if not INSIDE_BLENDER:
    print("ERROR: render_from_json must be run inside Blender")
    return
  
  output_dir = os.path.dirname(args.output_image) if args.output_image else '.'
  if output_dir and not os.path.exists(output_dir):
    os.makedirs(output_dir)
  
  with open(args.scene_file, 'r') as f:
    scene_struct = json.load(f)
  
  num_objects = len(scene_struct.get('objects', []))
  print(f"Scene has {num_objects} objects")
  
  base_scene_path = os.path.abspath(args.base_scene_blendfile)
  bpy.ops.wm.open_mainfile(filepath=base_scene_path)
  
  try:
    load_materials(args.material_dir)
  except Exception as e:
    print(f"Warning: Could not load materials: {e}")

  background_color = scene_struct.get('background_color', None)
  if background_color:
    set_background_color(background_color)
    set_ground_color(background_color)
  
  lighting = scene_struct.get('lighting', None)
  if lighting:
    set_lighting(lighting)
  
  render_args = bpy.context.scene.render
  render_args.engine = "CYCLES"
  render_args.filepath = args.output_image
  render_args.resolution_x = args.width
  render_args.resolution_y = args.height
  render_args.resolution_percentage = 100
  
  if args.use_gpu == 1:
    try:
      bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
      bpy.context.scene.cycles.device = 'GPU'
      print("[OK] GPU rendering enabled")
    except Exception as e:
      print(f"Warning: Could not enable GPU: {e}")
  
  bpy.context.scene.cycles.samples = args.render_num_samples
  
  filter_type = scene_struct.get('filter_type')
  filter_strength = scene_struct.get('filter_strength', 1.0)
  
  if filter_type == 'fisheye':
    camera = bpy.data.objects.get('Camera')
    if camera and camera.data:
      cam_data = camera.data
      if cam_data.type == 'PERSP':
        cam_data.lens = cam_data.lens * 0.7
        print(f"[OK] Zoomed out camera for fisheye: lens={cam_data.lens:.1f}mm")
  
  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
    size_mapping = properties['sizes']
  
  shape_semantic_to_file = properties['shapes']
  material_semantic_to_file = properties['materials']
  
  blender_objects = []
  print("Adding objects to scene...")
  for i, obj_info in enumerate(scene_struct.get('objects', [])):
    x, y, z = obj_info['3d_coords']
    r = size_mapping[obj_info['size']]
    semantic_shape = obj_info['shape']
    
    if semantic_shape == 'cube':
      r /= math.sqrt(2)
    
    if semantic_shape not in shape_semantic_to_file:
      print(f"ERROR: Shape '{semantic_shape}' not found")
      continue
    
    shape_file_name = shape_semantic_to_file[semantic_shape]
    
    try:
      add_object(args.shape_dir, shape_file_name, r, (x, y), theta=obj_info['rotation'])
    except Exception as e:
      print(f"Error adding object {i}: {e}")
      continue
    if INSIDE_BLENDER and bpy.context.object:
      blender_objects.append(bpy.context.object)
    
    rgba = color_name_to_rgba[obj_info['color']]
    semantic_material = obj_info['material']
    
    if semantic_material not in material_semantic_to_file:
      print(f"ERROR: Material '{semantic_material}' not found")
      continue
    
    mat_file_name = material_semantic_to_file[semantic_material]
    
    try:
      add_material(mat_file_name, Color=rgba)
    except Exception as e:
      print(f"Warning: Could not add material: {e}")
  
  if blender_objects:
    cf_meta = scene_struct.get('cf_metadata') or {}
    cf_type = cf_meta.get('cf_type', '')

    visibility_info = None
    if INSIDE_BLENDER and Image is not None:
      try:
        visibility_info = compute_visibility_fractions(blender_objects)
      except Exception as e:
        print(f"Warning: compute_visibility_fractions failed during render: {e}")
        visibility_info = None

    all_visible = True
    fail_reason = 'unknown visibility failure'

    if visibility_info is not None:
      ratios, scene_counts, full_counts = visibility_info

      if cf_type == 'occlusion_change':
        too_hidden = [
          (i, r) for i, r in enumerate(ratios)
          if full_counts[i] > 0 and r < CF_OCCLUSION_HARD_MIN_FRACTION
        ]
        band_objects = [
          (i, r) for i, r in enumerate(ratios)
          if full_counts[i] > 0 and CF_OCCLUSION_MIN_VISIBILITY_FRACTION <= r <= CF_OCCLUSION_MAX_VISIBILITY_FRACTION
        ]

        if too_hidden:
          all_visible = False
          min_r = min(r for (_, r) in too_hidden)
          fail_reason = (f'at least one object is too occluded in occlusion_change CF; '
                         f'min visibility fraction={min_r:.3f} '
                         f'(required >= {CF_OCCLUSION_HARD_MIN_FRACTION})')
        elif not band_objects:
          all_visible = False
          fail_reason = (f'no object falls into required occlusion band '
                         f'[{CF_OCCLUSION_MIN_VISIBILITY_FRACTION}, '
                         f'{CF_OCCLUSION_MAX_VISIBILITY_FRACTION}]')
        else:
          all_visible = True

      else:
        too_occluded = [
          (i, r) for i, r in enumerate(ratios)
          if full_counts[i] > 0 and r < BASE_MIN_VISIBILITY_FRACTION
        ]
        if too_occluded:
          all_visible = False
          min_r = min(r for (_, r) in too_occluded)
          fail_reason = (f'at least one object is too occluded in base scene; '
                         f'min visibility fraction={min_r:.3f} '
                         f'(required >= {BASE_MIN_VISIBILITY_FRACTION})')
        else:
          all_visible = True

    else:
      # Fallback to legacy absolute pixel-based visibility when we cannot
      # compute per-object relative visibility (e.g., PIL not available).
      w = getattr(args, 'width', 320)
      h = getattr(args, 'height', 240)
      if cf_type == 'occlusion_change':
        min_pixels = min_visible_pixels(w, h, MIN_VISIBLE_FRACTION_PARTIAL_OCCLUSION, MIN_PIXELS_FLOOR)
      else:
        base = min_visible_pixels(w, h, MIN_VISIBLE_FRACTION, MIN_PIXELS_FLOOR)
        min_pixels = max(getattr(args, 'min_pixels_per_object', MIN_PIXELS_FLOOR), base)
      all_visible = check_visibility(blender_objects, min_pixels)
      if not all_visible:
        fail_reason = 'at least one object has too few visible pixels'

    if not all_visible:
      print(f'Visibility check failed: {fail_reason}')
      for obj in blender_objects:
        try:
          delete_object(obj)
        except Exception:
          pass
      sys.exit(1)
  
  filter_type = scene_struct.get('filter_type')
  filter_strength = scene_struct.get('filter_strength', 1.0)
  
  print(f"Rendering to {args.output_image}...")
  
  try:
    bpy.ops.render.render(write_still=True)
    print("[OK] Rendering complete!")
  except Exception as e:
    print(f"Error during rendering: {e}")
    sys.exit(1)
  
  post_filter_type = scene_struct.get('filter_type')
  if post_filter_type and post_filter_type != 'fisheye':
    if Image is None:
      print(f"Warning: PIL not available, cannot apply post-filter {post_filter_type}")
    elif not os.path.exists(args.output_image):
      print(f"Warning: Output image does not exist: {args.output_image}")
    else:
      try:
        post_filter_strength = scene_struct.get('filter_strength', 1.0)
        apply_filter_to_image(args.output_image, post_filter_type, post_filter_strength)
      except Exception as e:
        import traceback
        print(f"Warning: Failed to apply post-filter {post_filter_type}: {e}")
        traceback.print_exc()

def main(args):
  if args.scene_file:
    render_from_json(args)
    return
  
  num_digits = 6
  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)
  
  all_scene_paths = []
  for i in range(args.num_images):
    img_path = img_template % (i + args.start_idx)
    scene_path = scene_template % (i + args.start_idx)
    all_scene_paths.append(scene_path)
    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,
    )

  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
  }
  if args.output_scene_file:
    output_dir = os.path.dirname(os.path.abspath(args.output_scene_file))
    if output_dir:
      os.makedirs(output_dir, exist_ok=True)
    with open(args.output_scene_file, '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,
  ):

  base_scene_path = os.path.abspath(args.base_scene_blendfile)
  bpy.ops.wm.open_mainfile(filepath=base_scene_path)
  load_materials(args.material_dir)

  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
  if args.use_gpu == 1:
    bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
    bpy.context.preferences.addons['cycles'].preferences.get_devices()
    for device in bpy.context.preferences.addons['cycles'].preferences.devices:
      device.use = True

  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'

  scene_struct = {
      'split': output_split,
      'image_index': output_index,
      'image_filename': os.path.basename(output_image),
      'objects': [],
      'directions': {},
  }

  bpy.ops.mesh.primitive_plane_add(size=10, location=(0, 0, 0))
  plane = bpy.context.object

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

  if args.camera_jitter > 0:
    for i in range(3):
      bpy.data.objects['Camera'].location[i] += rand(args.camera_jitter)

  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_object(plane)

  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)

  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)

  objects, blender_objects = add_random_objects(scene_struct, num_objects, args, camera)
  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)

  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 add_random_objects(scene_struct, num_objects, args, camera, max_scene_attempts=10):
  scene_attempt = 0
  while scene_attempt < max_scene_attempts:
    scene_attempt += 1

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

    shape_color_combos = None
    if args.shape_color_combos_json is not None:
      with open(args.shape_color_combos_json, 'r') as f:
        shape_color_combos = list(json.load(f).items())

    positions = []
    objects = []
    blender_objects = []
    for i in range(num_objects):
      size_name, r = random.choice(size_mapping)

      num_tries = 0
      while True:
        num_tries += 1
        if num_tries > args.max_retries:
          for obj in blender_objects:
            delete_object(obj)
          break
        x = random.uniform(-3, 3)
        y = random.uniform(-3, 3)
        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
      
      if num_tries > args.max_retries:
        break

      if shape_color_combos is None:
        obj_name, obj_name_out = random.choice(object_mapping)
        color_name, rgba = random.choice(list(color_name_to_rgba.items()))
      else:
        obj_name_out, color_choices = random.choice(shape_color_combos)
        color_name = random.choice(color_choices)
        obj_name = [k for k, v in object_mapping if v == obj_name_out][0]
        rgba = color_name_to_rgba[color_name]

      if obj_name == 'Cube':
        r /= math.sqrt(2)

      theta = 360.0 * random.random()
      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))

      mat_name, mat_name_out = random.choice(material_mapping)
      add_material(mat_name, Color=rgba)

      pixel_coords = 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,
      })

    if len(objects) < num_objects:
      continue

    visibility_info = None
    if INSIDE_BLENDER and Image is not None:
      try:
        visibility_info = compute_visibility_fractions(blender_objects)
      except Exception as e:
        print(f"Warning: compute_visibility_fractions failed during scene generation: {e}")
        visibility_info = None

    all_visible = True
    if visibility_info is not None:
      ratios, scene_counts, full_counts = visibility_info
      min_ratio = min((r for r in ratios if full_counts[ratios.index(r)] > 0), default=1.0)
      all_visible = all(
        (full_counts[i] == 0) or (ratios[i] >= BASE_MIN_VISIBILITY_FRACTION)
        for i in range(len(ratios))
      )
      if not all_visible:
        print(f'Some objects are too occluded in generated scene; '
              f'min visibility fraction={min_ratio:.3f} (required >= {BASE_MIN_VISIBILITY_FRACTION})')
    else:
      # Fallback to legacy absolute pixel-based visibility when PIL or Blender context is unavailable.
      min_pixels = max(args.min_pixels_per_object, min_visible_pixels(args.width, args.height))
      all_visible = check_visibility(blender_objects, min_pixels)

    if not all_visible:
      print('Some objects are occluded; replacing objects')
      for obj in blender_objects:
        delete_object(obj)
      continue

    return objects, blender_objects

  raise RuntimeError(f"Failed to generate a valid scene after {max_scene_attempts} attempts")


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 j.
  """
  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 compute_visibility_fractions(blender_objects):
  if not INSIDE_BLENDER or not blender_objects:
    return None
  if Image is None:
    return None

  # First pass: all objects together (occluded counts).
  fd, path = tempfile.mkstemp(suffix='.png')
  os.close(fd)
  try:
    colors_list = render_shadeless(blender_objects, path, use_distinct_colors=True)
    img = Image.open(path).convert('RGB')
    w, h = img.size
    pix = img.load()
    color_to_idx = {}
    for i, (r, g, b) in enumerate(colors_list):
      key = (round(r * 255), round(g * 255), round(b * 255))
      color_to_idx[key] = i
    scene_counts = [0] * len(blender_objects)
    for y in range(h):
      for x in range(w):
        key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
        if key in color_to_idx:
          scene_counts[color_to_idx[key]] += 1
  finally:
    try:
      os.remove(path)
    except Exception:
      pass

  # Second pass: per-object "full area" with other objects hidden.
  full_counts = []
  original_hide_render = [obj.hide_render for obj in blender_objects]
  try:
    for idx, obj in enumerate(blender_objects):
      # Hide all other objects, ensure this one is visible.
      for j, other in enumerate(blender_objects):
        if j == idx:
          other.hide_render = False
        else:
          other.hide_render = True

      fd_i, path_i = tempfile.mkstemp(suffix='.png')
      os.close(fd_i)
      try:
        colors_list = render_shadeless([obj], path_i, use_distinct_colors=True)
        img = Image.open(path_i).convert('RGB')
        w, h = img.size
        pix = img.load()
        color_to_idx = {}
        for i, (r, g, b) in enumerate(colors_list):
          key = (round(r * 255), round(g * 255), round(b * 255))
          color_to_idx[key] = i
        count = 0
        for y in range(h):
          for x in range(w):
            key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
            if key in color_to_idx:
              count += 1
        full_counts.append(count)
      finally:
        try:
          os.remove(path_i)
        except Exception:
          pass
  finally:
    # Restore previous hide_render flags.
    for obj, prev in zip(blender_objects, original_hide_render):
      obj.hide_render = prev

  visibility = []
  for scene_c, full_c in zip(scene_counts, full_counts):
    if full_c <= 0:
      visibility.append(0.0)
    else:
      visibility.append(float(scene_c) / float(full_c))

  return visibility, scene_counts, full_counts


def check_visibility(blender_objects, min_pixels_per_object):
  """
  Legacy absolute pixel-count visibility check, kept as a fallback when
  relative per-object visibility cannot be computed.
  """
  if not INSIDE_BLENDER or not blender_objects:
    return True
  if Image is None:
    return True
  fd, path = tempfile.mkstemp(suffix='.png')
  os.close(fd)
  try:
    colors_list = render_shadeless(blender_objects, path, use_distinct_colors=True)
    img = Image.open(path).convert('RGB')
    w, h = img.size
    pix = img.load()
    color_to_idx = {}
    for i, (r, g, b) in enumerate(colors_list):
      key = (round(r * 255), round(g * 255), round(b * 255))
      color_to_idx[key] = i
    counts = [0] * len(blender_objects)
    for y in range(h):
      for x in range(w):
        key = (pix[x, y][0], pix[x, y][1], pix[x, y][2])
        if key in color_to_idx:
          counts[color_to_idx[key]] += 1
    all_visible = all(c >= min_pixels_per_object for c in counts)
    return all_visible
  finally:
    try:
      os.remove(path)
    except Exception:
      pass


def render_shadeless(blender_objects, path='flat.png', use_distinct_colors=False):
  """
  Render a version of the scene with shading disabled and unique materials
  assigned to all objects. The image itself is written to path. This is used to ensure
  that all objects will be visible in the final rendered scene (when check_visibility is enabled).
  Returns a list of (r,g,b) colors in object order (for visibility counting when use_distinct_colors=True).
  """
  render_args = bpy.context.scene.render

  old_filepath = render_args.filepath
  old_engine = render_args.engine

  render_args.filepath = path
  render_args.engine = 'BLENDER_EEVEE_NEXT'
  
  view_layer = bpy.context.scene.view_layers[0]
  old_use_pass_combined = view_layer.use_pass_combined
  
  for obj_name in ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back', 'Ground']:
    if obj_name in bpy.data.objects:
      obj = bpy.data.objects[obj_name]
      obj.hide_render = True

  n = len(blender_objects)
  object_colors = [] if use_distinct_colors else set()
  old_materials = []
  for i, obj in enumerate(blender_objects):
    if len(obj.data.materials) > 0:
      old_materials.append(obj.data.materials[0])
    else:
      old_materials.append(None)
    
    mat = bpy.data.materials.new(name='Material_%d' % i)
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    nodes.clear()
    
    node_emission = nodes.new(type='ShaderNodeEmission')
    node_output = nodes.new(type='ShaderNodeOutputMaterial')
    
    if use_distinct_colors:
      r = (i + 1) / (n + 1)
      g, b = 0.5, 0.5
      object_colors.append((r, g, b))
    else:
      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))
    
    node_emission.inputs['Color'].default_value = (r, g, b, 1.0)
    mat.node_tree.links.new(node_emission.outputs['Emission'], node_output.inputs['Surface'])
    
    if len(obj.data.materials) > 0:
      obj.data.materials[0] = mat
    else:
      obj.data.materials.append(mat)

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

  for mat, obj in zip(old_materials, blender_objects):
    if mat is not None:
      obj.data.materials[0] = mat
    elif len(obj.data.materials) > 0:
      obj.data.materials.clear()

  for obj_name in ['Lamp_Key', 'Lamp_Fill', 'Lamp_Back', 'Ground']:
    if obj_name in bpy.data.objects:
      obj = bpy.data.objects[obj_name]
      obj.hide_render = False

  render_args.filepath = old_filepath
  render_args.engine = old_engine

  return object_colors


if __name__ == '__main__':
  if INSIDE_BLENDER:
    argv = 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')