Post-processing

"""Color grading: exposure, contrast, saturation and temperature, and a 3D lookup table.

Each segment grades the same test card through one full-frame pass, set with
``scene.segment(..., post=...)``: the tiles show hue and saturation, the ramp
below them shows the tone curve. The first pass leaves every color as it is.
Each label is the post-process of its segment.
"""

import os

from gaanim import Brush, PostProcess, Scene, Zones

HOLD = 0.8
LUT = "assets/teal_orange.cube"

scene = Scene()
background = scene.canvas.color("background")
foreground = scene.canvas.color("foreground")
page = scene.layout.zones(Zones.rows(["6fr", "1fr"], gap="24px", padding="page_padding"))


def test_card(zone):
    """Fill ``zone`` with four colored tiles, each holding a dark shape, above a gray ramp."""
    tiles, ramp = zone.split(Zones.rows(["4fr", "1fr"], gap="24px"))
    cells = tiles.split(Zones.columns(4, gap="24px"))
    roles = ["accent", "success", "danger", "chart"]
    shapes = [
        lambda: scene.geometry.circle(1.0),
        lambda: scene.geometry.regular_polygon(3, 1.0),
        lambda: scene.geometry.star(5, 1.0, 0.45),
        lambda: scene.geometry.annulus(1.0, 0.55),
    ]
    for cell, role, shape in zip(cells, roles, shapes):
        scene.geometry.rounded_rect(cell.width, cell.height, 0.2).fill(scene.canvas.color(role)).place(cell)
        shape().fill(background).place(cell.inset("48px"), fit="contain")
    (left, y), (right, _) = ramp.point("left"), ramp.point("right")
    gray = Brush.linear([background, foreground], start=(left, y), end=(right, y))
    scene.geometry.rounded_rect(ramp.width, ramp.height, 0.2).fill(gray).place(ramp)


def grade(code, post):
    """One segment of the test card through ``post``, labelled with ``code``."""
    scene.segment(code, post=post)
    test_card(page[0])
    scene.text(code, role="code").place(page[1], anchor="left")
    scene.wait(HOLD)


# color_grade works in linear light: exposure is in stops and temperature
# runs from -1 (cool) to 1 (warm); then contrast pivots about mid-gray and
# saturation about luma. 1 keeps contrast and saturation, 0 the rest.
GRADES = [
    ("PostProcess.color_grade()", PostProcess.color_grade()),
    ("PostProcess.color_grade(exposure=1.0)", PostProcess.color_grade(exposure=1.0)),
    ("PostProcess.color_grade(exposure=-1.0)", PostProcess.color_grade(exposure=-1.0)),
    ("PostProcess.color_grade(contrast=1.8)", PostProcess.color_grade(contrast=1.8)),
    ("PostProcess.color_grade(contrast=0.5)", PostProcess.color_grade(contrast=0.5)),
    ("PostProcess.color_grade(saturation=0.0)", PostProcess.color_grade(saturation=0.0)),
    ("PostProcess.color_grade(saturation=2.0)", PostProcess.color_grade(saturation=2.0)),
    ("PostProcess.color_grade(temperature=1.0)", PostProcess.color_grade(temperature=1.0)),
    ("PostProcess.color_grade(temperature=-1.0)", PostProcess.color_grade(temperature=-1.0)),
    # A .cube lookup table maps every color to another; strength blends it
    # with the original.
    ("PostProcess.lut(LUT)", PostProcess.lut(LUT)),
    ("PostProcess.lut(LUT, strength=0.5)", PostProcess.lut(LUT, strength=0.5)),
]
for code, post in GRADES:
    grade(code, post)

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # The middle of every segment.
    scene.snapshots(snapshots, [HOLD * (index + 0.5) for index in range(len(GRADES))])
else:
    scene.render()

"""Finishing passes: vignette, grain, chromatic aberration, halftone, dither, CRT, pixelate and glitch.

Each segment draws the same test card through one PostProcess preset, set
for the whole frame with ``scene.segment(..., post=...)``. The preset's main
argument rises from a value that leaves the frame untouched to the one in
the label. The label stays out of the post-process with
``post_process(False)``, so it reads under every look.
"""

import os

from gaanim import Brush, PostProcess, Scene, Zones

LEAD = 0.3
RAMP = 0.6
HOLD = 0.6

scene = Scene()
background = scene.canvas.color("background")
foreground = scene.canvas.color("foreground")
page = scene.layout.zones(Zones.rows(["6fr", "1fr"], gap="24px", padding="page_padding"))


def test_card(zone):
    """Fill ``zone`` with four colored tiles, each holding a dark shape, above a gray ramp."""
    tiles, ramp = zone.split(Zones.rows(["4fr", "1fr"], gap="24px"))
    cells = tiles.split(Zones.columns(4, gap="24px"))
    roles = ["accent", "success", "danger", "chart"]
    shapes = [
        lambda: scene.geometry.circle(1.0),
        lambda: scene.geometry.regular_polygon(3, 1.0),
        lambda: scene.geometry.star(5, 1.0, 0.45),
        lambda: scene.geometry.annulus(1.0, 0.55),
    ]
    for cell, role, shape in zip(cells, roles, shapes):
        scene.geometry.rounded_rect(cell.width, cell.height, 0.2).fill(scene.canvas.color(role)).place(cell)
        shape().fill(background).place(cell.inset("48px"), fit="contain")
    (left, y), (right, _) = ramp.point("left"), ramp.point("right")
    gray = Brush.linear([background, foreground], start=(left, y), end=(right, y))
    scene.geometry.rounded_rect(ramp.width, ramp.height, 0.2).fill(gray).place(ramp)


def finish(template, build, start, end):
    """One segment of the test card through ``build(amount)``, as ``amount`` goes from ``start`` to ``end``.

    Every preset argument takes a number, a Parameter or a Computed; the
    label shows the preset with its final number. Returns the segment start.
    """
    amount = scene.viz.parameter(start)
    code = template.format(end)
    begin = scene.cursor
    scene.segment(code, post=build(amount))
    test_card(page[0])
    scene.text(code, role="code").place(page[1], anchor="left").post_process(False)
    scene.wait(LEAD)
    scene.play([amount.animate.set(end).duration(RAMP)])
    scene.wait(HOLD)
    return begin


# Film finishes: vignette darkens toward the corners (softness sets where the
# falloff starts), grain adds brightness noise in cells of `size` pixels,
# renewed 24 times a second unless animated=False, and chromatic aberration
# splits red and blue toward the corners.
FINISHES = [
    ("PostProcess.vignette({})", lambda amount: PostProcess.vignette(amount), 0.0, 0.8),
    ("PostProcess.vignette({}, softness=0.2)", lambda amount: PostProcess.vignette(amount, softness=0.2), 0.0, 1.0),
    ("PostProcess.grain({})", lambda amount: PostProcess.grain(amount), 0.0, 0.15),
    ("PostProcess.grain({}, size=4)", lambda amount: PostProcess.grain(amount, size=4), 0.0, 0.3),
    ("PostProcess.chromatic_aberration({})", lambda amount: PostProcess.chromatic_aberration(amount), 0.0, 0.02),
    # Styles: print dots, ordered dithering to a few levels per channel, a
    # curved screen with scanlines, blocks, and bands that jump 12 times a
    # second (seed picks another pattern). Sizes are pixels of a 1080p frame.
    ("PostProcess.halftone({})", lambda amount: PostProcess.halftone(amount), 1, 12),
    ("PostProcess.dither({})", lambda amount: PostProcess.dither(amount), 256, 4),
    ("PostProcess.dither({})", lambda amount: PostProcess.dither(amount), 256, 2),
    ("PostProcess.crt({})", lambda amount: PostProcess.crt(amount), 0.0, 1.0),
    ("PostProcess.pixelate({})", lambda amount: PostProcess.pixelate(amount), 1, 16),
    ("PostProcess.glitch({})", lambda amount: PostProcess.glitch(amount), 0.0, 0.5),
    ("PostProcess.glitch({}, seed=3)", lambda amount: PostProcess.glitch(amount, seed=3), 0.0, 1.0),
]
starts = {template.format(end): finish(template, build, start, end) for template, build, start, end in FINISHES}

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # Every look at full strength; halftone and pixelate also early in their
    # ramp.
    scene.snapshots(snapshots, sorted(
        [begin + LEAD + RAMP + HOLD - 0.2 for begin in starts.values()]
        + [starts[code] + LEAD + 0.3 * RAMP for code in ("PostProcess.halftone(12)", "PostProcess.pixelate(16)")]
    ))
else:
    scene.render()

"""Bloom: threshold, intensity and radius, and a glow that pulses with a Parameter.

Each segment draws the same outlines and bars through one bloom pass over the
whole frame. Everything brighter than the threshold glows: the bright outlines
by default, the gray bars only when the threshold drops below their
brightness. Each label is the post-process of its segment.
"""

import os

from gaanim import PostProcess, Scene, Zones

HOLD = 1.0
PULSE = 0.8

scene = Scene()
page = scene.layout.zones(Zones.rows(["6fr", "1fr"], gap="24px", padding="page_padding"))


def neon(zone):
    """Bright outlines over three filled bars of decreasing brightness."""
    shapes, bars = zone.split(Zones.rows(["4fr", "1fr"], gap="48px"))
    outlines = [
        (lambda: scene.geometry.circle(1.0), "accent"),
        (lambda: scene.geometry.star(5, 1.0, 0.45), "success"),
        (lambda: scene.geometry.square(1.6), "danger"),
        (lambda: scene.geometry.regular_polygon(3, 1.0), "foreground"),
    ]
    for cell, (shape, role) in zip(shapes.split(Zones.columns(4, gap="48px")), outlines):
        shape().no_fill().stroke(scene.canvas.color(role), 0.16).place(cell, fit="contain")
    for cell, role in zip(bars.split(Zones.columns(3, gap="48px")), ["chart", "muted", "rule"]):
        scene.geometry.rounded_rect(cell.width, cell.height, cell.height / 2).fill(scene.canvas.color(role)).place(cell)


def glow(code, post):
    """One segment of the outlines through ``post``, labelled with ``code``."""
    scene.segment(code, post=post)
    neon(page[0])
    scene.text(code, role="code").place(page[1], anchor="left")


# threshold (0..1) is the brightness where the glow starts, intensity scales
# the added light and radius (0..1) spreads it from a tight halo to a haze.
BLOOMS = [
    ("PostProcess.bloom()", PostProcess.bloom()),
    ("PostProcess.bloom(intensity=2.0)", PostProcess.bloom(intensity=2.0)),
    ("PostProcess.bloom(intensity=2.0, threshold=0.4)", PostProcess.bloom(intensity=2.0, threshold=0.4)),
    ("PostProcess.bloom(intensity=2.0, radius=0.1)", PostProcess.bloom(intensity=2.0, radius=0.1)),
    ("PostProcess.bloom(intensity=2.0, radius=1.0)", PostProcess.bloom(intensity=2.0, radius=1.0)),
]
for code, post in BLOOMS:
    glow(code, post)
    scene.wait(HOLD)

# Any of the three takes a Parameter, so the glow animates like any value.
pulse_start = scene.cursor
level = scene.viz.parameter(0.0)
scene.segment("pulse", post=PostProcess.bloom(intensity=level))
neon(page[0])
lines = page[1].split(Zones.rows(2))
scene.text("PostProcess.bloom(intensity=level)", role="code").place(lines[0], anchor="left")
scene.text("level.animate.set(4.0).repeat(2, yoyo=True)", role="code").place(lines[1], anchor="left")
scene.play([level.animate.set(4.0).duration(PULSE).repeat(2, yoyo=True)])
scene.wait(0.3)

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # The middle of every segment, then the pulse rising and at its peak.
    scene.snapshots(
        snapshots,
        [HOLD * (index + 0.5) for index in range(len(BLOOMS))] + [pulse_start + 0.4, pulse_start + PULSE],
    )
else:
    scene.render()

"""Motion blur: sub-frame averaging that smears motion by its speed, and drawables kept sharp.

``scene.canvas.motion_blur`` averages sub-frames over the time a film
shutter stays open, in exports and snapshots (the interactive preview stays
sharp). The first segment moves dots across their tracks at three speeds,
the second spins, scales and orbits shapes, and the last four give a
segment a shutter of its own. Each label is the code that moves its row or
panel, or that sets its segment's shutter.
"""

import math
import os

from gaanim import Easing, Scene, Zones

scene = Scene()
rule = scene.canvas.color("rule")

# A 360 degree shutter stays open for a whole frame (180 is the film
# standard, 720 the maximum), sampled 32 times; phase= moves the shutter
# relative to the frame time and None turns the blur off. It applies to
# every segment without a shutter of its own.
scene.canvas.motion_blur(360, samples=32)

# Every dot crosses its track at a constant speed: the faster it goes, the
# longer its smear. motion_blur(False) keeps one drawable sharp.
speeds_start = scene.cursor
scene.segment("speeds")
heading, body = scene.layout.zones(Zones.rows(["2fr", "5fr"], gap="24px", padding="page_padding"))
lines = heading.split(Zones.rows(2))
scene.text("scene.canvas.motion_blur(360, samples=32)", role="code").place(lines[0], anchor="left")
scene.text("move = dot.animate.shift_by(width, 0).easing(Easing.LINEAR)", role="code").place(lines[1], anchor="left")
rows = body.split(Zones.grid(rows=4, columns=["3fr", "5fr"], column_gap="48px"))
SPEEDS = [
    ("move.duration(1.6)", 1.6, True),
    ("move.duration(0.4)", 0.4, True),
    ("move.duration(0.2)", 0.2, True),
    ("dot.motion_blur(False)", 0.2, False),
]
moves = []
for index, (label, duration, blurred) in enumerate(SPEEDS):
    scene.text(label, role="code").place(rows[2 * index], anchor="right")
    track = rows[2 * index + 1]
    (start, y), (end, _) = track.point("left"), track.point("right")
    scene.geometry.line(start, y, end, y).stroke(rule, 0.02)
    dot = scene.geometry.dot(0.32).move_to(start, y).motion_blur(blurred)
    moves.append(dot.animate.shift_by(track.width, 0).easing(Easing.LINEAR).duration(duration))
scene.wait(0.3)
scene.play(moves)
scene.wait(0.3)


def spin(stage, size):
    """A bar across ``stage`` that turns once."""
    x, y = stage.center
    return scene.geometry.rect(size, size * 0.12).move_to(x, y).animate.rotate_by(math.tau)


def shrink(stage, size):
    """A ring filling ``stage`` that shrinks."""
    x, y = stage.center
    return scene.geometry.annulus(size / 2, size / 2 - 0.3).move_to(x, y).animate.scale_to(0.4)


def orbit(stage, size):
    """A dot that goes once around a circle filling ``stage``."""
    x, y = stage.center
    path = scene.geometry.circle(size / 2 - 0.3).no_fill().stroke(rule, 0.02).move_to(x, y)
    return scene.geometry.dot(0.32).move_to(x + size / 2 - 0.3, y).animate.move_along(path)


# Rotation, scale and travel along a curve (each label is the animate call of
# its shape) blur along the path each point follows within the frame.
shapes_start = scene.cursor
scene.segment("shapes")
panels = scene.layout.zones(Zones.columns(3, gap="48px", padding="page_padding"))
plays = []
for panel, (label, build) in zip(panels, [
    ("rotate_by(math.tau)", spin),
    ("scale_to(0.4)", shrink),
    ("move_along(orbit)", orbit),
]):
    stage, caption = panel.split(Zones.rows(["5fr", "1fr"]))
    scene.text(label, role="code").place(caption)
    plays.append(build(stage, min(stage.width, stage.height) * 0.9).duration(0.5).easing(Easing.LINEAR))
scene.wait(0.3)
scene.play(plays)
scene.wait(0.3)

# Segment.motion_blur gives one segment a shutter of its own, with the
# arguments of scene.canvas.motion_blur; None keeps it sharp. A ring kept
# sharp marks where the dot is at the frame time: a centered shutter smears
# around it, while phase=0 opens the shutter there and smears ahead.
SHUTTERS = [
    ("narrow", {"shutter_angle": 90}, "motion_blur(90, samples=32)"),
    ("wide", {"shutter_angle": 720}, "motion_blur(720, samples=32)"),
    ("ahead", {"shutter_angle": 360, "phase": 0}, "motion_blur(360, samples=32, phase=0)"),
    ("sharp", None, "motion_blur(None)"),
]
shutter_starts = []
for name, options, call in SHUTTERS:
    shutter_starts.append(scene.cursor)
    segment = scene.segment(name)
    if options is None:
        segment.motion_blur(None)
    else:
        segment.motion_blur(samples=32, **options)
    heading, body = scene.layout.zones(Zones.rows(["2fr", "5fr"], gap="24px", padding="page_padding"))
    scene.text(f'scene.segment("{name}").{call}', role="code").place(heading, anchor="left")
    track = body.inset(("0px", "96px"))
    (start, y), (end, _) = track.point("left"), track.point("right")
    scene.geometry.line(start, y, end, y).stroke(rule, 0.02)
    dot = scene.geometry.dot(0.32).move_to(start, y)
    ring = scene.geometry.circle(0.42).no_fill().stroke(rule, 0.03).move_to(start, y).motion_blur(False)
    scene.wait(0.3)
    scene.play([
        shape.animate.shift_by(track.width, 0).easing(Easing.LINEAR).duration(0.4) for shape in (dot, ring)
    ])
    scene.wait(0.3)

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # The dots while all four move and as the slowest goes on alone; the
    # shapes at a quarter and halfway through their motion; each shutter
    # halfway along the track.
    scene.snapshots(
        snapshots,
        [speeds_start + 0.4, speeds_start + 0.8, shapes_start + 0.42, shapes_start + 0.55]
        + [start + 0.5 for start in shutter_starts],
    )
else:
    scene.render()

"""Post-process chains, the scene default and its segment overrides, and custom WGSL passes with uniforms.

The first two segments chain the same two passes in both orders, the next
three inherit, replace and switch off the scene's post-process, and the last
three run a WGSL pass of our own whose uniform is a number, an animated
Parameter and seeded noise. Each label is the code that sets the segment's
post-process.
"""

import os

from gaanim import Brush, PostProcess, Scene, Zones

HOLD = 1.0

# gaanim_post receives the frame position (uv, 0..1 from the top left), its
# size in output pixels and the timeline time; gaanim_scene(uv) samples what
# was drawn and gaanim_uniforms holds the uniforms, as f32 fields.
WAVE = """
fn gaanim_post(uv: vec2<f32>, resolution: vec2<f32>, time: f32) -> vec4<f32> {
    let shift = sin(uv.y * 12.0 - time * 4.0) * gaanim_uniforms.amount;
    return gaanim_scene(uv + vec2<f32>(shift, 0.0));
}
"""

scene = Scene()
background = scene.canvas.color("background")
foreground = scene.canvas.color("foreground")
page = scene.layout.zones(Zones.rows(["5fr", "1fr"], gap="24px", padding="page_padding"))

# The scene's post-process applies to every segment that does not set its own.
scene.canvas.post = PostProcess.vignette(1.0, softness=0.3)


def test_card(zone):
    """Fill ``zone`` with four colored tiles, each holding a dark shape, above a gray ramp."""
    tiles, ramp = zone.split(Zones.rows(["4fr", "1fr"], gap="24px"))
    cells = tiles.split(Zones.columns(4, gap="24px"))
    roles = ["accent", "success", "danger", "chart"]
    shapes = [
        lambda: scene.geometry.circle(1.0),
        lambda: scene.geometry.regular_polygon(3, 1.0),
        lambda: scene.geometry.star(5, 1.0, 0.45),
        lambda: scene.geometry.annulus(1.0, 0.55),
    ]
    for cell, role, shape in zip(cells, roles, shapes):
        scene.geometry.rounded_rect(cell.width, cell.height, 0.2).fill(scene.canvas.color(role)).place(cell)
        shape().fill(background).place(cell.inset("48px"), fit="contain")
    (left, y), (right, _) = ramp.point("left"), ramp.point("right")
    gray = Brush.linear([background, foreground], start=(left, y), end=(right, y))
    scene.geometry.rounded_rect(ramp.width, ramp.height, 0.2).fill(gray).place(ramp)


def page_with(name, lines, **options):
    """Start segment ``name`` with ``options``, draw the test card and label it with ``lines`` of code."""
    start = scene.cursor
    scene.segment(name, **options)
    test_card(page[0])
    for line, zone in zip(lines, page[1].split(Zones.rows(len(lines)))):
        scene.text(line, role="code").place(zone, anchor="left")
    return start


# A list chains passes in order, each one sampling what the previous one
# wrote: aberration then gray leaves gray fringes, gray then aberration
# splits color out of gray tiles.
starts = []
starts.append(page_with("aberration first", [
    "post=[PostProcess.chromatic_aberration(0.015),",
    "      PostProcess.color_grade(saturation=0.0)]",
], post=[PostProcess.chromatic_aberration(0.015), PostProcess.color_grade(saturation=0.0)]))
scene.wait(HOLD)
starts.append(page_with("gray first", [
    "post=[PostProcess.color_grade(saturation=0.0),",
    "      PostProcess.chromatic_aberration(0.015)]",
], post=[PostProcess.color_grade(saturation=0.0), PostProcess.chromatic_aberration(0.015)]))
scene.wait(HOLD)

# Without post= a segment inherits scene.canvas.post; post= replaces it while
# the segment is active, and post=False (or an empty list) draws it plain.
DEFAULT = "scene.canvas.post = PostProcess.vignette(1.0, softness=0.3)"
starts.append(page_with("inherit", [DEFAULT, 'scene.segment("inherit")']))
scene.wait(HOLD)
starts.append(page_with("override", [DEFAULT, 'scene.segment("override", post=PostProcess.crt())'],
                        post=PostProcess.crt()))
scene.wait(HOLD)
starts.append(page_with("plain", [DEFAULT, 'scene.segment("plain", post=False)'], post=False))
scene.wait(HOLD)

# Uniforms are numbers, or a Parameter, Variable, Computed or scene.time of
# this scene, read every frame.
starts.append(page_with("number", ['PostProcess.shader(WAVE, uniforms={"amount": 0.02})'],
                        post=PostProcess.shader(WAVE, uniforms={"amount": 0.02})))
scene.wait(HOLD)
swell = scene.viz.parameter(0.0)
starts.append(page_with("parameter", [
    'PostProcess.shader(WAVE, uniforms={"amount": swell})',
    "swell.animate.set(0.04).repeat(2, yoyo=True)",
], post=PostProcess.shader(WAVE, uniforms={"amount": swell})))
scene.play([swell.animate.set(0.04).duration(HOLD / 2).repeat(2, yoyo=True)])
noise = scene.noise(frequency=2.0, amplitude=0.03, seed=4)
starts.append(page_with("noise", [
    'PostProcess.shader(WAVE, uniforms={"amount": noise})',
    "noise = scene.noise(frequency=2.0, amplitude=0.03, seed=4)",
], post=PostProcess.shader(WAVE, uniforms={"amount": noise})))
scene.wait(HOLD)

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # The middle of every segment, which is the Parameter's peak in its
    # segment, the Parameter on its way up, and the noise at two moments.
    scene.snapshots(snapshots, sorted([start + HOLD / 2 for start in starts] + [starts[6] + HOLD / 4, starts[7] + 0.9]))
else:
    scene.render()