lighting.py

"""Optional "wow" lighting after an import: a physical sky with a matching sun, a
camera framing the building, EEVEE with shadows and ray tracing, AgX, and the
viewport switched to rendered shading. Everything is named so a re-import
updates the same world, sun and camera instead of adding more."""

from __future__ import annotations

import math
import os

import bpy
from mathutils import Vector

WORLD_NAME = "Pascal sky"
SUN_NAME = "Pascal sun"
CAMERA_NAME = "Pascal camera"
GROUND_NAME = "Pascal ground"
GROUND_COLOR = (0.22, 0.24, 0.23)
GROUND_RADIUS = 2000.0
EXPOSURE = 0.0
CAMERA_LENS = 32.0
CAMERA_AZIMUTH = math.radians(-35)
CAMERA_ELEVATION = math.radians(20)
ENGINE_CANDIDATES = ("BLENDER_EEVEE", "BLENDER_EEVEE_NEXT")

# Lighting comes from one of Blender's bundled world HDRIs (neutral, no colour
# cast) while the camera sees a physical sky backdrop that matches the sun; the
# sun lamp adds the crisp shadows a small HDRI cannot. "sky" is the pure
# physical-sky option.
PRESETS = {
    "daylight": {
        "label": "Daylight",
        "hdri": "courtyard.exr",
        "hdri_strength": 0.9,
        "hdri_rotation": math.radians(0),
        "sun_elevation": math.radians(42),
        "sun_rotation": math.radians(205),  # behind the camera: lit facades, shadows falling away
        "sun_energy": 2.5,
        "sun_color": (1.0, 0.98, 0.95),
        "sky_strength": 0.45,
    },
    "golden": {
        "label": "Golden hour",
        "hdri": "sunrise.exr",
        "hdri_strength": 1.2,
        "hdri_rotation": math.radians(120),
        "sun_elevation": math.radians(22),
        "sun_rotation": math.radians(205),
        "sun_energy": 3.5,
        "sun_color": (1.0, 0.86, 0.68),
        "sky_strength": 0.35,
    },
    "overcast": {
        "label": "Overcast",
        "hdri": "city.exr",
        "hdri_strength": 1.1,
        "hdri_rotation": math.radians(0),
        "sun_elevation": math.radians(50),
        "sun_rotation": math.radians(205),
        "sun_energy": 0.8,
        "sun_color": (1.0, 1.0, 1.0),
        "sky_strength": 0.5,
    },
    "night": {
        "label": "Night",
        "hdri": "night.exr",
        "hdri_strength": 0.2,
        "hdri_rotation": math.radians(0),
        "sun_elevation": math.radians(35),  # the moon
        "sun_rotation": math.radians(205),
        "sun_energy": 0.3,
        "sun_color": (0.55, 0.66, 1.0),
        "sky_strength": 0.03,
        "sky_sun_elevation": math.radians(-6),  # backdrop: last twilight
        "exposure": 1.0,
        "room_lights": True,
    },
    "sky": {
        "label": "Physical sky",
        "hdri": None,
        "hdri_strength": 0.0,
        "hdri_rotation": 0.0,
        "sun_elevation": math.radians(32),
        "sun_rotation": math.radians(205),
        "sun_energy": 4.5,
        "sun_color": (1.0, 0.93, 0.82),
        "sky_strength": 0.4,
    },
}
DEFAULT_PRESET = "daylight"
ROOM_LIGHT_NAME = "Pascal room light"
LIGHTS_COLLECTION = "Pascal lights"
ROOM_LIGHT_HEIGHT = 2.2
ROOM_LIGHT_POWER = 400.0
PORCH_LIGHT_POWER = 250.0
PORCH_LIGHT_HEIGHT = 2.6
ROOM_LIGHT_COLOR = (1.0, 0.76, 0.52)
ROOM_LIGHT_LIMIT = 64
PRESET_ITEMS = [(key, value["label"], "") for key, value in PRESETS.items()]


def _hdri_path(filename: str) -> str:
    # system_resource resolves directories only, never files.
    return os.path.join(bpy.utils.system_resource("DATAFILES", path="studiolights/world"), filename)


def _world(scene: bpy.types.Scene, preset: dict) -> bpy.types.World:
    world = bpy.data.worlds.get(WORLD_NAME) or bpy.data.worlds.new(WORLD_NAME)
    scene.world = world
    world.use_nodes = True
    nodes, links = world.node_tree.nodes, world.node_tree.links
    nodes.clear()

    sky = nodes.new("ShaderNodeTexSky")
    sky.location = (-600, -250)
    for sky_type in ("MULTIPLE_SCATTERING", "NISHITA"):
        try:
            sky.sky_type = sky_type
            break
        except TypeError:
            continue
    for name, value in (
        ("sun_elevation", preset.get("sky_sun_elevation", preset["sun_elevation"])),
        ("sun_rotation", preset["sun_rotation"]),
        ("sun_intensity", 1.0),
        ("sun_disc", False),  # the sun lamp provides the direct light
        ("altitude", 50.0),
        ("air_density", 1.0),
        ("dust_density", 2.0),
        ("ozone_density", 1.0),
        ("turbidity", 3.5),
    ):
        if hasattr(sky, name):
            setattr(sky, name, value)
    sky_strength = nodes.new("ShaderNodeMath")
    sky_strength.operation = "MULTIPLY"
    sky_strength.location = (-350, -250)
    sky_strength.inputs[1].default_value = preset["sky_strength"]
    # Math takes a float; scale each channel via a Mix in multiply mode instead.
    sky_scale = nodes.new("ShaderNodeMix")
    sky_scale.data_type = "RGBA"
    sky_scale.blend_type = "MULTIPLY"
    sky_scale.location = (-350, -250)
    sky_scale.inputs[0].default_value = 1.0
    sky_scale.inputs[7].default_value = (preset["sky_strength"],) * 3 + (1.0,)
    nodes.remove(sky_strength)
    links.new(sky.outputs["Color"], sky_scale.inputs[6])

    background = nodes.new("ShaderNodeBackground")
    background.location = (100, 0)
    background.inputs["Strength"].default_value = 1.0
    output = nodes.new("ShaderNodeOutputWorld")
    output.location = (350, 0)

    hdri = preset["hdri"]
    if hdri:
        image = bpy.data.images.load(_hdri_path(hdri), check_existing=True)
        coords = nodes.new("ShaderNodeTexCoord")
        coords.location = (-1100, 150)
        mapping = nodes.new("ShaderNodeMapping")
        mapping.location = (-900, 150)
        mapping.inputs["Rotation"].default_value = (0.0, 0.0, preset["hdri_rotation"])
        environment = nodes.new("ShaderNodeTexEnvironment")
        environment.location = (-650, 150)
        environment.image = image
        env_scale = nodes.new("ShaderNodeMix")
        env_scale.data_type = "RGBA"
        env_scale.blend_type = "MULTIPLY"
        env_scale.location = (-350, 150)
        env_scale.inputs[0].default_value = 1.0
        env_scale.inputs[7].default_value = (preset["hdri_strength"],) * 3 + (1.0,)
        light_path = nodes.new("ShaderNodeLightPath")
        light_path.location = (-350, 400)
        # Lighting from the HDRI, the sky as what the camera actually sees.
        pick = nodes.new("ShaderNodeMix")
        pick.data_type = "RGBA"
        pick.location = (-100, 0)
        links.new(coords.outputs["Generated"], mapping.inputs["Vector"])
        links.new(mapping.outputs["Vector"], environment.inputs["Vector"])
        links.new(environment.outputs["Color"], env_scale.inputs[6])
        links.new(light_path.outputs["Is Camera Ray"], pick.inputs[0])
        links.new(env_scale.outputs[2], pick.inputs[6])
        links.new(sky_scale.outputs[2], pick.inputs[7])
        links.new(pick.outputs[2], background.inputs["Color"])
    else:
        links.new(sky_scale.outputs[2], background.inputs["Color"])
    links.new(background.outputs["Background"], output.inputs["Surface"])
    return world


def _sun(root: bpy.types.Collection, preset: dict) -> bpy.types.Object:
    sun = bpy.data.objects.get(SUN_NAME)
    if sun is None or sun.type != "LIGHT":
        light = bpy.data.lights.new(SUN_NAME, "SUN")
        sun = bpy.data.objects.new(SUN_NAME, light)
        root.objects.link(sun)
    light = sun.data
    light.energy = preset["sun_energy"]
    light.color = preset["sun_color"]
    light.angle = math.radians(0.6)
    light.use_shadow = True
    # A sun lamp shines down its local -Z: tilt it from the zenith by the
    # complement of the elevation, then spin it to the sky's sun rotation.
    sun.rotation_euler = (math.pi / 2 - preset["sun_elevation"], 0.0, preset["sun_rotation"])
    sun.location = (0.0, 0.0, 10.0)
    return sun


def _ground(root: bpy.types.Collection, bounds: tuple[Vector, Vector] | None) -> bpy.types.Object:
    """A wide neutral disc under the site: the physical sky is dark below the
    horizon and would otherwise show past the edge of the exported ground."""
    ground = bpy.data.objects.get(GROUND_NAME)
    if ground is None or ground.type != "MESH":
        mesh = bpy.data.meshes.new(GROUND_NAME)
        segments = 64
        points = [
            (GROUND_RADIUS * math.cos(2 * math.pi * i / segments), GROUND_RADIUS * math.sin(2 * math.pi * i / segments), 0.0)
            for i in range(segments)
        ]
        mesh.from_pydata(points, [], [list(range(segments))])
        mesh.update()
        ground = bpy.data.objects.new(GROUND_NAME, mesh)
        root.objects.link(ground)
    material = bpy.data.materials.get(GROUND_NAME)
    if material is None:
        material = bpy.data.materials.new(GROUND_NAME)
        material.use_nodes = True
        material.diffuse_color = (*GROUND_COLOR, 1.0)
        principled = next((n for n in material.node_tree.nodes if n.type == "BSDF_PRINCIPLED"), None)
        if principled is not None:
            principled.inputs["Base Color"].default_value = (*GROUND_COLOR, 1.0)
            principled.inputs["Roughness"].default_value = 1.0
    if not ground.data.materials:
        ground.data.materials.append(material)
    floor = bounds[0].z if bounds else 0.0
    ground.location = (0.0, 0.0, floor - 0.01)
    ground.hide_select = True
    return ground


def _zone_anchor(zone: bpy.types.Object) -> Vector | None:
    polygon = zone.get("polygon")
    if polygon is None or len(polygon) < 3:
        return None
    points = [(float(p[0]), -float(p[1])) for p in polygon]
    cx = sum(x for x, _ in points) / len(points)
    cy = sum(y for _, y in points) / len(points)
    return zone.matrix_world @ Vector((cx, cy, ROOM_LIGHT_HEIGHT))


def _room_lights(root: bpy.types.Collection, objects, enabled: bool) -> int:
    """Warm bulbs in every room (one per zone, at its centroid) so a night
    scene glows from the inside. Rebuilt on every re-light; removed when the
    preset does not ask for them."""
    for obj in list(bpy.data.objects):
        if obj.name.split(".")[0] == ROOM_LIGHT_NAME:
            bpy.data.objects.remove(obj)
    collection = bpy.data.collections.get(LIGHTS_COLLECTION)
    if not enabled:
        if collection is not None:
            bpy.data.collections.remove(collection)
        return 0
    if collection is None:
        collection = bpy.data.collections.new(LIGHTS_COLLECTION)
    if collection.name not in root.children:
        root.children.link(collection)
    anchors = [
        anchor
        for zone in objects
        if zone.get("kind") == "zone" and (anchor := _zone_anchor(zone)) is not None
    ]
    if not anchors:
        bounds = _bounds(objects)
        if bounds is not None:
            low, high = bounds
            anchors = [Vector(((low.x + high.x) / 2, (low.y + high.y) / 2, low.z + ROOM_LIGHT_HEIGHT))]
    bulbs = [(anchor, ROOM_LIGHT_POWER) for anchor in anchors[:ROOM_LIGHT_LIMIT]]
    # The spawn point is where visitors arrive: a porch light there reads as the entrance.
    for spawn in objects:
        if spawn.get("kind") == "spawn":
            bulbs.append((spawn.matrix_world @ Vector((0.0, 0.0, PORCH_LIGHT_HEIGHT)), PORCH_LIGHT_POWER))
            break
    for anchor, power in bulbs:
        data = bpy.data.lights.new(ROOM_LIGHT_NAME, "POINT")
        data.energy = power
        data.color = ROOM_LIGHT_COLOR
        data.shadow_soft_size = 0.25
        light = bpy.data.objects.new(ROOM_LIGHT_NAME, data)
        light.location = anchor
        collection.objects.link(light)
    return len(bulbs)


def _bounds(objects) -> tuple[Vector, Vector] | None:
    low = Vector((math.inf,) * 3)
    high = Vector((-math.inf,) * 3)
    found = False
    for obj in objects:
        if obj.type != "MESH" or obj.hide_render or obj.get("kind") == "zone-floor":
            continue
        if obj.name.split(".")[0] == GROUND_NAME:
            continue
        for corner in obj.bound_box:
            point = obj.matrix_world @ Vector(corner)
            low = Vector(map(min, low, point))
            high = Vector(map(max, high, point))
            found = True
    return (low, high) if found else None


def _camera(scene: bpy.types.Scene, root: bpy.types.Collection, objects) -> bpy.types.Object | None:
    bounds = _bounds(objects)
    if bounds is None:
        return None
    low, high = bounds
    center = (low + high) / 2
    size = max(high - low)
    camera = bpy.data.objects.get(CAMERA_NAME)
    if camera is None or camera.type != "CAMERA":
        data = bpy.data.cameras.new(CAMERA_NAME)
        camera = bpy.data.objects.new(CAMERA_NAME, data)
        root.objects.link(camera)
    camera.data.lens = CAMERA_LENS
    camera.data.clip_end = max(1000.0, size * 20)
    fov = 2 * math.atan(18.0 / CAMERA_LENS)  # 36 mm sensor width
    distance = (size * 0.62) / math.tan(fov / 2)
    direction = Vector(
        (
            math.cos(CAMERA_ELEVATION) * math.sin(CAMERA_AZIMUTH),
            -math.cos(CAMERA_ELEVATION) * math.cos(CAMERA_AZIMUTH),
            math.sin(CAMERA_ELEVATION),
        )
    )
    target = Vector((center.x, center.y, low.z + (high.z - low.z) * 0.4))
    camera.location = target + direction * distance
    camera.rotation_euler = (-direction).to_track_quat("-Z", "Y").to_euler()
    scene.camera = camera
    return camera


ENGINE_CHOICES = [
    ("KEEP", "Keep current", "Leave the render engine as it is"),
    ("EEVEE", "EEVEE", "Switch to EEVEE for an instant real-time preview"),
    ("CYCLES", "Cycles", "Switch to Cycles for path-traced renders"),
]


def engine_preference() -> str:
    try:
        return str(bpy.context.preferences.addons[__package__].preferences.render_engine)
    except (KeyError, AttributeError):
        return "KEEP"


def _apply_engine(scene: bpy.types.Scene, choice: str) -> None:
    """The engine is the user's call; only switch when they asked for it."""
    candidates = {"EEVEE": ENGINE_CANDIDATES, "CYCLES": ("CYCLES",)}.get(choice, ())
    for engine in candidates:
        try:
            scene.render.engine = engine
            return
        except TypeError:
            continue


def _render_settings(scene: bpy.types.Scene, engine_choice: str = "KEEP", exposure: float = EXPOSURE) -> None:
    _apply_engine(scene, engine_choice)
    if scene.render.engine in ENGINE_CANDIDATES:
        eevee = scene.eevee
        for name, value in (
            ("use_shadows", True),
            ("use_raytracing", True),
            ("use_fast_gi", True),
            ("taa_render_samples", 64),
            ("shadow_ray_count", 2),
        ):
            if hasattr(eevee, name):
                setattr(eevee, name, value)
    elif scene.render.engine == "CYCLES" and hasattr(scene, "cycles"):
        cycles = scene.cycles
        for name, value in (("use_preview_denoising", True), ("use_denoising", True)):
            if hasattr(cycles, name):
                setattr(cycles, name, value)
    view = scene.view_settings
    for name, value in (("view_transform", "AgX"), ("look", "AgX - Medium High Contrast")):
        try:
            setattr(view, name, value)
        except TypeError:
            pass
    view.exposure = exposure
    if scene.world is not None and hasattr(scene.world, "use_sun_shadow"):
        scene.world.use_sun_shadow = False


def _rendered_viewports(hdri: str | None) -> None:
    if bpy.app.background:
        return
    for window in bpy.context.window_manager.windows:
        for area in window.screen.areas:
            if area.type != "VIEW_3D":
                continue
            for space in area.spaces:
                if space.type == "VIEW_3D":
                    shading = space.shading
                    shading.type = "RENDERED"
                    shading.use_scene_lights_render = True
                    shading.use_scene_world_render = True
                    if hdri:
                        try:
                            shading.studio_light = hdri
                            shading.studiolight_background_alpha = 1.0
                            shading.studiolight_background_blur = 0.5
                        except TypeError:
                            pass
            area.tag_redraw()


def preset_preference() -> str:
    try:
        return str(bpy.context.preferences.addons[__package__].preferences.lighting_preset)
    except (KeyError, AttributeError):
        return DEFAULT_PRESET


def setup_lighting(context, root: bpy.types.Collection, objects, preset_name: str | None = None) -> None:
    preset = PRESETS.get(preset_name or preset_preference(), PRESETS[DEFAULT_PRESET])
    scene = context.scene
    _world(scene, preset)
    _sun(root, preset)
    _ground(root, _bounds(objects))
    _camera(scene, root, objects)
    _room_lights(root, objects, bool(preset.get("room_lights")))
    _render_settings(scene, engine_preference(), float(preset.get("exposure", EXPOSURE)))
    _rendered_viewports(preset["hdri"])