import bpy
import os

base_path = "src/graphics/technology"

base_collection = bpy.data.scenes["Scene"].view_layers["Base"].layer_collection
shadows_collection = bpy.data.scenes["Scene"].view_layers["Shadows"].layer_collection

def set_exclude(model, value):
	base_collection.children[model].exclude = value
	shadows_collection.children[model].exclude = value

def set_color(model, variant, values):
	for obj in base_collection.children[model].collection.all_objects.values():
		if obj.data.name == "Floor":
			obj.hide_render = False
		if "liquid_color" not in obj:
			continue
		obj["liquid_color"][0] = values[0]
		obj["liquid_color"][1] = values[1]
		obj["liquid_color"][2] = values[2]
		obj.update_tag()
		glass = "glass_frozen" if variant == "frozen" else "glass_no_back_grime"
		obj.material_slots[0].material = bpy.data.materials[glass]
		obj.material_slots[1].material = bpy.data.materials[f"liquid_clear_surface"]
		obj.material_slots[2].material = bpy.data.materials[f"liquid_clear_glass"]

def set_empty(model, variant):
	for obj in base_collection.children[model].collection.all_objects.values():
		if obj.data.name == "Floor":
			obj.hide_render = True
		if "liquid_color" not in obj:
			continue
		glass = "glass_frozen" if variant == "frozen" else "glass_no_back_grime"
		obj.material_slots[0].material = bpy.data.materials[glass]
		obj.material_slots[1].material = bpy.data.materials[f"transparent"]
		obj.material_slots[2].material = bpy.data.materials[f"glass_inside"]
			

def render_bottle(color, variant, model, output):
	set_exclude(model, False)
	if color != "empty":
		set_color(model, variant, colors[color])
	else:
		set_empty(model, variant)
	bpy.data.scenes["Scene"].render.filepath = output
	bpy.ops.render.render(write_still=True)
	set_exclude(model, True)

models = [
	"cone_normal",
	"cone_slim",
	"cone_inverted",
	"cylinder",
	"tube_one",
	"tube_two",
	"tube_three",
	"sphere_normal",
	"sphere_tiny",
	"sphere_double",
	"sphere_tubed",
	"sphere_hemi",
	"sphere_spiked",
	"hourglass",
	"torus",
	"klein",
	"pyramid",
	"cube",
	"triangle",
	"triangle_alt",
	"pentagon",
	"hexagon",
]

colors = {
	"red": [0.554, 0.044, 0.046],
	"green": [0.075, 0.573, 0.075],
	"black": [0.050, 0.050, 0.050], 
	"cyan": [0.066, 0.430, 0.592],
	"purple": [0.233, 0.039, 0.372],
	"yellow": [0.536, 0.402, 0.000],
	"white": [0.515, 0.534, 0.555],
	"orange": [0.644, 0.277, 0.012],
	"pink": [0.557, 0.043, 0.433],
	"blue": [0.038, 0.050, 0.606],
	"lime": [0.523, 0.570, 0.000],
	"empty": None,
}


# Code to render all types
bottles = []
for model in models:
	bottles.append(["empty", "clear", model, f"//{base_path}/{model}_empty.png"])
	variants = ["clear"]
	if model == "cone_normal":
		variants.append("frozen")
	for variant in variants:
		for color in colors:
			if color == "empty":
				continue
			bottles.append([color, variant, model, f"//{base_path}/{model}_{variant}_{color}.png"])

# Code for prototyping speedup
#bpy.data.scenes["Scene"].cycles.samples = 1024

# Reset visibility of all bottles
for model in models:
	set_exclude(model, True)

# Render the bottles
for color, variant, model, output in bottles:
	# Code to only render files that don't exist
	#if os.path.exists(bpy.path.abspath(output)): continue

	print("Rendering", output)
	render_bottle(color, variant, model, output)

print("Done!")
