import bpy


base_collection = bpy.data.scenes["Base"].view_layers["RenderLayer"].layer_collection
highlight_collection = bpy.data.scenes["Toon Highlight"].view_layers["RenderLayer"].layer_collection
shadows_collection = bpy.data.scenes["Toon Shadows"].view_layers["RenderLayer"].layer_collection

def set_exclude(model, value):
    base_collection.children[model].exclude = value
    highlight_collection.children[f"Highlight {model}"].exclude = value
    shadows_collection.children[f"Shadows {model}"].exclude = value

def set_color(model, color):
    for obj in base_collection.children[model].collection.all_objects.values():
        if len(obj.material_slots) < 3:
            continue
        obj.material_slots[1].material = bpy.data.materials[f"Liquid Surface {color}"]
        obj.material_slots[2].material = bpy.data.materials[f"Liquid Glass {color}"]

def render_bottle(color, model, output):
    set_exclude(model, False)
    set_color(model, color)
    bpy.data.scenes["Base"].render.filepath = output
    bpy.ops.render.render(write_still=True)
    set_exclude(model, True)

models = [
    "Cone",
    "Cone Slim",
    "Cone Inverted",
    "Cylinder",
    "Tube",
    "Tubes Duo",
    "Tubes Trio",
    "Sphere",
    "Sphere Tiny",
    "Sphere Double",
    "Sphere Tubed",
    "Cube",
    "Spiked",
    "Pyramid",
    "Triangle",
    "Triangle Alt",
    "Pentagon",
    "Hexagon",
]


bottles = [
    ["Red", "Cone", "//src/graphics/red.png"],
    ["Green", "Cube", "//src/graphics/green.png"],
    ["Black", "Spiked", "//src/graphics/black.png"],
    ["Cyan", "Sphere", "//src/graphics/cyan.png"],
    ["Purple", "Sphere Double", "//src/graphics/purple.png"],
    ["Yellow", "Triangle Alt", "//src/graphics/yellow.png"],
    ["White", "Cone Inverted", "//src/graphics/white.png"],
]

# Code to render a preview of all types
# bottles = [["Red", type, f"//types/{type}.png"] for type in models]

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

# Set active scene so the highlights/shadows isn't accidentally rendered
bpy.context.window.scene = bpy.data.scenes["Base"]

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

# Render the bottles
for color, model, output in bottles:
    render_bottle(color, model, output)

print("Done!")