3D Modeling & Render API

import bpy
import math
import mathutils
import bmesh
import os

def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)
    for block in [bpy.data.meshes, bpy.data.cameras, bpy.data.lights, bpy.data.materials, bpy.data.curves]:
        for item in list(block):
            block.remove(item)

def look_at(obj_camera, target_position):
    loc_camera = obj_camera.location
    direction = mathutils.Vector(target_position) - loc_camera
    rot_quat = direction.to_track_quat('-Z', 'Y')
    obj_camera.rotation_euler = rot_quat.to_euler()

def look_at_light(obj_light, target_position):
    loc_light = obj_light.location
    direction = mathutils.Vector(target_position) - loc_light
    # Track target with -Z, align local Y axis of light (rectangle length) with world Z
    rot_quat = direction.to_track_quat('-Z', 'Z')
    obj_light.rotation_euler = rot_quat.to_euler()

def create_ceramic_material():
    mat = bpy.data.materials.new(name="CeramicGlossy")
    mat.use_nodes = True
    nodes = mat.node_tree.nodes
    links = mat.node_tree.links
    nodes.clear()
    
    # Create output node
    node_output = nodes.new(type='ShaderNodeOutputMaterial')
    
    # Create Mix Shader node
    mix_shader = nodes.new(type='ShaderNodeMixShader')
    links.new(mix_shader.outputs[0], node_output.inputs['Surface'])
    
    # 1. Mug Body Shader (Rich Deep Navy Blue)
    bsdf_mug = nodes.new(type='ShaderNodeBsdfPrincipled')
    # 2. Heart Shader (Vibrant Pink)
    bsdf_heart = nodes.new(type='ShaderNodeBsdfPrincipled')
    
    # Helper to set inputs
    def set_inputs(bsdf_node, color):
        def set_input(node, name, value):
            if name in node.inputs:
                node.inputs[name].default_value = value
        set_input(bsdf_node, 'Base Color', color)
        set_input(bsdf_node, 'Roughness', 0.06)
        set_input(bsdf_node, 'Metallic', 0.0)
        set_input(bsdf_node, 'Specular', 0.5)
        set_input(bsdf_node, 'Specular IOR Level', 0.5)
        set_input(bsdf_node, 'IOR', 1.55)
        set_input(bsdf_node, 'Clearcoat', 0.4)
        set_input(bsdf_node, 'Coat Weight', 0.4)
        set_input(bsdf_node, 'Clearcoat Roughness', 0.02)
        set_input(bsdf_node, 'Coat Roughness', 0.02)
        
    set_inputs(bsdf_mug, (0.02, 0.05, 0.18, 1.0)) # Rich Deep Navy Blue
    set_inputs(bsdf_heart, (0.95, 0.25, 0.5, 1.0)) # Vibrant Hot Pink
    
    # Link shaders to Mix Shader
    links.new(bsdf_mug.outputs['BSDF'], mix_shader.inputs[1])
    links.new(bsdf_heart.outputs['BSDF'], mix_shader.inputs[2])
    
    # 3. Procedural Heart Mask using Generated Coordinates (always normalized 0-1)
    tex_coord = nodes.new(type='ShaderNodeTexCoord')
    sep_xyz = nodes.new(type='ShaderNodeSeparateXYZ')
    # Use Generated output (index 0)
    links.new(tex_coord.outputs[0], sep_xyz.inputs['Vector'])
    
    # Center X and Y around 0.5
    sub_x = nodes.new(type='ShaderNodeMath')
    sub_x.operation = 'SUBTRACT'
    sub_x.inputs[1].default_value = 0.5
    links.new(sep_xyz.outputs['X'], sub_x.inputs[0])
    
    sub_y = nodes.new(type='ShaderNodeMath')
    sub_y.operation = 'SUBTRACT'
    sub_y.inputs[1].default_value = 0.5
    links.new(sep_xyz.outputs['Y'], sub_y.inputs[0])
    
    # theta = atan2(sub_y, sub_x)
    atan2_node = nodes.new(type='ShaderNodeMath')
    atan2_node.operation = 'ARCTAN2'
    links.new(sub_y.outputs['Value'], atan2_node.inputs[0])
    links.new(sub_x.outputs['Value'], atan2_node.inputs[1])
    
    # u = theta + pi/2 (centers the heart at local -Y side of the cylinder, opposite the handle)
    add_pi2 = nodes.new(type='ShaderNodeMath')
    add_pi2.operation = 'ADD'
    add_pi2.inputs[1].default_value = math.pi / 2.0
    links.new(atan2_node.outputs['Value'], add_pi2.inputs[0])
    
    # u' = u * scale_u (adjust width of the heart)
    scale_u = nodes.new(type='ShaderNodeMath')
    scale_u.operation = 'MULTIPLY'
    scale_u.inputs[1].default_value = 0.8 # perfect width scale
    links.new(add_pi2.outputs['Value'], scale_u.inputs[0])
    
    # abs_u = abs(u')
    abs_u = nodes.new(type='ShaderNodeMath')
    abs_u.operation = 'ABSOLUTE'
    links.new(scale_u.outputs['Value'], abs_u.inputs[0])
    
    # sq_u = u'^2
    sq_u = nodes.new(type='ShaderNodeMath')
    sq_u.operation = 'MULTIPLY'
    links.new(scale_u.outputs['Value'], sq_u.inputs[0])
    links.new(scale_u.outputs['Value'], sq_u.inputs[1])
    
    # sqrt_u = sqrt(abs_u)
    sqrt_u = nodes.new(type='ShaderNodeMath')
    sqrt_u.operation = 'SQRT'
    links.new(abs_u.outputs['Value'], sqrt_u.inputs[0])
    
    # offset = 0.5 * sqrt_u (adjust lobes curve of the heart)
    offset_node = nodes.new(type='ShaderNodeMath')
    offset_node.operation = 'MULTIPLY'
    offset_node.inputs[1].default_value = 0.5
    links.new(sqrt_u.outputs['Value'], offset_node.inputs[0])
    
    # v = z - 0.5 (center vertically)
    sub_z = nodes.new(type='ShaderNodeMath')
    sub_z.operation = 'SUBTRACT'
    sub_z.inputs[1].default_value = 0.5
    links.new(sep_xyz.outputs['Z'], sub_z.inputs[0])
    
    # v' = v * scale_v (adjust height of the heart)
    scale_v = nodes.new(type='ShaderNodeMath')
    scale_v.operation = 'MULTIPLY'
    scale_v.inputs[1].default_value = 1.1 # perfect height scale
    links.new(sub_z.outputs['Value'], scale_v.inputs[0])
    
    # diff = v' - offset
    diff_node = nodes.new(type='ShaderNodeMath')
    diff_node.operation = 'SUBTRACT'
    links.new(scale_v.outputs['Value'], diff_node.inputs[0])
    links.new(offset_node.outputs['Value'], diff_node.inputs[1])
    
    # sq_diff = diff^2
    sq_diff = nodes.new(type='ShaderNodeMath')
    sq_diff.operation = 'MULTIPLY'
    links.new(diff_node.outputs['Value'], sq_diff.inputs[0])
    links.new(diff_node.outputs['Value'], sq_diff.inputs[1])
    
    # D = sq_u + sq_diff
    dist_d = nodes.new(type='ShaderNodeMath')
    dist_d.operation = 'ADD'
    links.new(sq_u.outputs['Value'], dist_d.inputs[0])
    links.new(sq_diff.outputs['Value'], dist_d.inputs[1])
    
    # mask = D < R^2 (R=0.25 -> R^2=0.0625)
    compare_node = nodes.new(type='ShaderNodeMath')
    compare_node.operation = 'LESS_THAN'
    compare_node.inputs[1].default_value = 0.0625
    links.new(dist_d.outputs['Value'], compare_node.inputs[0])
    
    # Link mask to Mix Shader Factor
    links.new(compare_node.outputs['Value'], mix_shader.inputs[0])
    
    return mat

def create_mug_body():
    # Cylinder radius=0.6, depth=1.4, location=(0, 0, 0.7)
    bpy.ops.mesh.primitive_cylinder_add(vertices=64, radius=0.6, depth=1.4, location=(0.0, 0.0, 0.7))
    cup = bpy.context.active_object
    cup.name = "MugBody"
    
    # Go to Edit Mode to delete top face and create foot ring
    bpy.ops.object.mode_set(mode='EDIT')
    bm = bmesh.from_edit_mesh(cup.data)
    bm.faces.ensure_lookup_table()
    
    # Find top face (face with highest z coordinate)
    top_face = max(bm.faces, key=lambda f: f.calc_center_bounds().z)
    bmesh.ops.delete(bm, geom=[top_face], context='FACES_ONLY')
    bm.faces.ensure_lookup_table()
    
    # Find bottom face (face with normal pointing down Z)
    bottom_face = None
    for f in bm.faces:
        if f.normal.dot(mathutils.Vector((0.0, 0.0, -1.0))) > 0.99:
            bottom_face = f
            break
            
    if bottom_face:
        # Inset bottom face to define foot ring width (0.08)
        inset_res = bmesh.ops.inset_individual(bm, faces=[bottom_face], thickness=0.08)
        new_bottom_face = inset_res['faces'][0]
        # Translate the inner face upward to create the recess (foot ring)
        bmesh.ops.translate(bm, vec=(0.0, 0.0, 0.04), verts=new_bottom_face.verts)
        
    bmesh.update_edit_mesh(cup.data)
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # Apply Scale
    bpy.ops.object.transform_apply(scale=True)
    
    # Add Solidify modifier
    solid = cup.modifiers.new(name="Solidify", type='SOLIDIFY')
    solid.thickness = 0.05
    solid.offset = -1 # Inward solidification
    
    # Apply Solidify modifier to bake geometry before boolean
    bpy.context.view_layer.objects.active = cup
    bpy.ops.object.modifier_apply(modifier="Solidify")
    
    return cup

def create_mug_handle():
    # Create curve
    curve_data = bpy.data.curves.new('HandleCurve', type='CURVE')
    curve_data.dimensions = '3D'
    curve_data.resolution_u = 16
    
    polyline = curve_data.splines.new('BEZIER')
    polyline.bezier_points.add(3) # total of 4 points
    
    # P0: top attach point (slightly inside the cup cylinder wall at Y = 0.58, Z = 1.05)
    p0 = polyline.bezier_points[0]
    p0.co = (0.0, 0.58, 1.05)
    p0.handle_left = (0.0, 0.58, 1.15)
    p0.handle_right = (0.0, 0.78, 1.05)
    
    # P1: top outer curve
    p1 = polyline.bezier_points[1]
    p1.co = (0.0, 1.15, 0.9)
    p1.handle_left = (0.0, 0.98, 1.0)
    p1.handle_right = (0.0, 1.25, 0.8)
    
    # P2: bottom outer curve
    p2 = polyline.bezier_points[2]
    p2.co = (0.0, 1.15, 0.5)
    p2.handle_left = (0.0, 1.25, 0.6)
    p2.handle_right = (0.0, 0.98, 0.4)
    
    # P3: bottom attach point
    p3 = polyline.bezier_points[3]
    p3.co = (0.0, 0.58, 0.35)
    p3.handle_left = (0.0, 0.78, 0.35)
    p3.handle_right = (0.0, 0.58, 0.25)
    
    for p in polyline.bezier_points:
        p.handle_left_type = 'FREE'
        p.handle_right_type = 'FREE'
        
    curve_data.bevel_depth = 0.08
    curve_data.bevel_resolution = 6
    
    handle_obj = bpy.data.objects.new('Handle', curve_data)
    bpy.context.collection.objects.link(handle_obj)
    
    # Convert to mesh
    bpy.context.view_layer.objects.active = handle_obj
    handle_obj.select_set(True)
    bpy.ops.object.convert(target='MESH')
    
    # Flatten the handle on X axis to make it oval shaped
    handle_obj.scale = (0.7, 1.0, 1.0)
    bpy.ops.object.transform_apply(scale=True)
    
    return handle_obj

def merge_cup_and_handle(cup, handle):
    bpy.ops.object.select_all(action='DESELECT')
    cup.select_set(True)
    bpy.context.view_layer.objects.active = cup
    
    # Attempt Boolean Union
    union_success = False
    try:
        bool_mod = cup.modifiers.new(name="UnionHandle", type='BOOLEAN')
        bool_mod.operation = 'UNION'
        bool_mod.object = handle
        bool_mod.solver = 'EXACT'
        
        # Apply modifier
        bpy.ops.object.modifier_apply(modifier="UnionHandle")
        # Delete handle
        bpy.data.objects.remove(handle, do_unlink=True)
        union_success = True
        print("Boolean Union of cup and handle succeeded.")
    except Exception as e:
        print(f"Boolean Union failed: {e}. Falling back to simple Join.")
        # Rollback: remove boolean modifier if it exists
        if "UnionHandle" in cup.modifiers:
            cup.modifiers.remove(cup.modifiers["UnionHandle"])
        # Fallback to Join
        bpy.ops.object.select_all(action='DESELECT')
        cup.select_set(True)
        handle.select_set(True)
        bpy.context.view_layer.objects.active = cup
        bpy.ops.object.join()
        
    # Apply bevel modifier first to round corners and fillet joints
    bev = cup.modifiers.new(name="Bevel", type='BEVEL')
    bev.width = 0.015
    bev.segments = 3
    bev.limit_method = 'ANGLE'
    bev.angle_limit = math.radians(25)
    
    # Apply subdivision modifier to make everything beautifully smooth
    sub = cup.modifiers.new(name="Subdivision", type='SUBSURF')
    sub.levels = 3
    sub.render_levels = 3
    
    # Shade smooth
    bpy.ops.object.shade_smooth()
    
    return cup

def create_backdrop():
    # Define a smooth profile in the Y-Z plane
    profile = [
        (-10.0, 0.0),
        (2.0, 0.0),
        (3.0, 0.02),
        (3.8, 0.1),
        (4.5, 0.3),
        (5.1, 0.7),
        (5.5, 1.3),
        (5.8, 2.1),
        (6.0, 3.1),
        (6.1, 4.3),
        (6.15, 6.0),
        (6.15, 12.0)
    ]
    
    # Generate vertices and faces extruded along X
    verts = []
    faces = []
    
    # Vertices at X = -15
    for y, z in profile:
        verts.append((-15.0, y, z))
    # Vertices at X = 15
    for y, z in profile:
        verts.append((15.0, y, z))
        
    n = len(profile)
    for i in range(n - 1):
        faces.append((i, i + 1, n + i + 1, n + i))
        
    mesh = bpy.data.meshes.new("Backdrop")
    mesh.from_pydata(verts, [], faces)
    mesh.update()
    
    obj = bpy.data.objects.new("Backdrop", mesh)
    bpy.context.collection.objects.link(obj)
    
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    bpy.ops.object.shade_smooth()
    
    # Backdrop material
    mat = bpy.data.materials.new("BackdropMat")
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get("Principled BSDF")
    if bsdf:
        # Clean white studio background color
        if 'Base Color' in bsdf.inputs:
            bsdf.inputs['Base Color'].default_value = (0.95, 0.94, 0.91, 1.0) # Warm off-white/cream
        if 'Roughness' in bsdf.inputs:
            bsdf.inputs['Roughness'].default_value = 0.8
        if 'Specular' in bsdf.inputs:
            bsdf.inputs['Specular'].default_value = 0.2
    obj.data.materials.append(mat)
    return obj

def build_scene(blend_path, img_path):
    print("Setting up coffee mug scene...")
    clear_scene()
    
    # 1. Create backdrop
    create_backdrop()
    
    # 2. Create the template mug
    cup_body = create_mug_body()
    handle = create_mug_handle()
    template_mug = merge_cup_and_handle(cup_body, handle)
    template_mug.name = "Mug_Template"
    
    # Assign Ceramic Material
    ceramic_mat = create_ceramic_material()
    # Explicitly clear slots to prevent any None slots from boolean union
    template_mug.data.materials.clear()
    template_mug.data.materials.append(ceramic_mat)
    
    # 3. Create Left Mug (rotated slightly back, handle visible to the right-back)
    mug1 = template_mug.copy()
    mug1.data = template_mug.data.copy()
    mug1.data.materials.clear()
    mug1.data.materials.append(ceramic_mat)
    mug1.name = "Mug_Left"
    mug1.location = (-0.95, -0.2, 0.0) # Moved left and forward to create separation
    mug1.rotation_euler.z = math.radians(-105)
    bpy.context.collection.objects.link(mug1)
    
    # 4. Create Right Mug (rotated slightly forward, handle visible to the right-front)
    mug2 = template_mug.copy()
    mug2.data = template_mug.data.copy()
    mug2.data.materials.clear()
    mug2.data.materials.append(ceramic_mat)
    mug2.name = "Mug_Right"
    mug2.location = (0.9, 0.2, 0.0) # Moved right and backward to prevent collision
    mug2.rotation_euler.z = math.radians(-65)
    bpy.context.collection.objects.link(mug2)
    
    # 5. Delete the template mug
    bpy.data.objects.remove(template_mug, do_unlink=True)
    
    # 6. Set up Camera (Studio portrait composition)
    camera_data = bpy.data.cameras.new('Camera')
    camera = bpy.data.objects.new('Camera', camera_data)
    # Position back and up looking at the center of the two mugs
    camera.location = (0.0, -7.5, 3.8)
    camera_data.lens = 85 # portrait focal length to reduce perspective distortion
    bpy.context.collection.objects.link(camera)
    look_at(camera, (0.0, 0.0, 0.65))
    bpy.context.scene.camera = camera
    
    # 7. Set up Lighting (3-Point soft studio lighting using vertical strip lights for glossy reflections)
    # Key Light (Front-Left, warm strip)
    key_light_data = bpy.data.lights.new(name="KeyLight", type='AREA')
    key_light_data.shape = 'RECTANGLE'
    key_light_data.size = 0.4
    key_light_data.size_y = 2.5
    key_light_data.energy = 800
    key_light_data.color = (1.0, 0.98, 0.96)
    key_light = bpy.data.objects.new(name="KeyLight", object_data=key_light_data)
    key_light.location = (-3.0, -3.0, 4.0)
    bpy.context.collection.objects.link(key_light)
    look_at_light(key_light, (0.0, 0.0, 0.6))
    
    # Fill Light (Front-Right, cool strip)
    fill_light_data = bpy.data.lights.new(name="FillLight", type='AREA')
    fill_light_data.shape = 'RECTANGLE'
    fill_light_data.size = 0.6
    fill_light_data.size_y = 3.0
    fill_light_data.energy = 350
    fill_light_data.color = (0.96, 0.98, 1.0)
    fill_light = bpy.data.objects.new(name="FillLight", object_data=fill_light_data)
    fill_light.location = (3.0, -2.5, 3.0)
    bpy.context.collection.objects.link(fill_light)
    look_at_light(fill_light, (0.0, 0.0, 0.6))
    
    # Rim Light / Backlight (Back-Center-Right, highlight)
    rim_light_data = bpy.data.lights.new(name="RimLight", type='AREA')
    rim_light_data.shape = 'RECTANGLE'
    rim_light_data.size = 0.5
    rim_light_data.size_y = 2.0
    rim_light_data.energy = 500
    rim_light_data.color = (1.0, 1.0, 1.0)
    rim_light = bpy.data.objects.new(name="RimLight", object_data=rim_light_data)
    rim_light.location = (1.0, 3.0, 4.5)
    bpy.context.collection.objects.link(rim_light)
    look_at_light(rim_light, (0.0, 0.0, 0.6))
    
    # Soft environment world color
    bpy.context.scene.world.use_nodes = True
    world_nodes = bpy.context.scene.world.node_tree.nodes
    if 'Background' in world_nodes:
        world_nodes['Background'].inputs['Color'].default_value = (0.05, 0.05, 0.05, 1.0)
        world_nodes['Background'].inputs['Strength'].default_value = 1.0
        
    # 8. Render Engine Configuration (Cycles)
    scene = bpy.context.scene
    scene.render.engine = 'CYCLES'
    scene.cycles.samples = 128
    scene.cycles.use_denoising = True
    scene.render.resolution_x = 1200
    scene.render.resolution_y = 1200
    scene.render.filepath = img_path
    
    # Attempt to enable GPU rendering (Metal on macOS)
    try:
        preferences = bpy.context.preferences
        cycles_preferences = preferences.addons['cycles'].preferences
        cycles_preferences.compute_device_type = 'METAL'
        cycles_preferences.get_devices()
        for device in cycles_preferences.devices:
            if device.type == 'METAL':
                device.use = True
                scene.cycles.device = 'GPU'
                print(f"Enabled Metal GPU rendering device: {device.name}")
    except Exception as e:
        print(f"GPU configuration failed, using CPU: {e}")
        scene.cycles.device = 'CPU'
        
    # Save the .blend file
    os.makedirs(os.path.dirname(blend_path), exist_ok=True)
    bpy.ops.wm.save_as_mainfile(filepath=blend_path)
    print(f"Saved .blend file to: {blend_path}")
    
    # Render the scene
    print("Rendering high quality coffee mug scene...")
    bpy.ops.render.render(write_still=True)
    print(f"Render complete and saved to: {img_path}")

if __name__ == "__main__":
    blend = "coffee-mug.blend"
    img = "coffee-mug_render.png"
    build_scene(blend, img)