Mechanics

"""Supports: every support kind, the directions they face, their sizes and their styleable parts.

``scene.mechanics.support_at`` draws the symbol of a structural support at a
connection point: a fixed plate, a pin, rollers, a guide, a cable or a spring.
The direction runs from the base toward the connection, the shortcuts fix the
kind, and each part (joint, body, rollers, guides, ground and hatching) can be
styled on its own. Each label is the code that builds the variant.
"""

import os

from gaanim import Direction, Scene, Zones

DRAW = 1.2
HOLD = 0.6

scene = Scene()
accent = scene.canvas.color("accent")
success = scene.canvas.color("success")


def cells(rows, columns):
    """A padded grid of ``rows`` x ``columns`` cells, each split into a label strip and a stage."""
    grid = scene.layout.zones(Zones.grid(rows=rows, columns=columns, gap="48px", padding="page_padding"))
    return [zone.split(Zones.rows([0.5, "1fr"], gap="12px")) for zone in grid]


def reveal(supports):
    """Draw the supports, then hold them."""
    scene.play([support.animate.create() for support in supports], duration=DRAW)
    scene.wait(HOLD)


def gallery(rows, columns, variants):
    """Label each cell with its code and place the support it builds in the cell's stage."""
    supports = []
    for (label_zone, stage), (label, support) in zip(cells(rows, columns), variants):
        scene.text(label, role="code").place(label_zone)
        supports.append(support.place(stage))
    reveal(supports)


def support(kind, direction=Direction.UP, size=1.6):
    """A support built at the origin, with a base wide enough for its size, to be placed in its cell."""
    return scene.mechanics.support_at((0, 0), kind=kind, direction=direction, size=size, ground_length=size * 1.7)


# support_at(point, kind=...) draws one of eight symbols, here at size=1.6
# with ground_length=2.7. Direction.UP puts the ground under the point.
scene.segment("kinds")
KINDS = ["fixed", "pin", "roller", "simple", "guided", "prismatic", "cable", "spring"]
gallery(2, 4, [(f'kind="{kind}"', support(kind)) for kind in KINDS])

# direction runs from the base toward the connection: around the frame, every
# roller (size=1.25) faces the centre. Direction.custom takes any vector.
scene.segment("directions")
COMPASS = ["DOWN_RIGHT", "DOWN", "DOWN_LEFT", "RIGHT", None, "LEFT", "UP_RIGHT", "UP", "UP_LEFT"]
gallery(3, 3, [
    (f"Direction.{name}", support("roller", getattr(Direction, name), size=1.25))
    if name
    else ("Direction.custom(1, 2)", support("roller", Direction.custom(1, 2), size=1.25))
    for name in COMPASS
])

# The shortcuts fix the kind. size scales the symbol, ground_length sets the
# base line, and color replaces the theme's foreground.
scene.segment("sizes")
SIZES = [
    ("pin_support(p, size=0.6)", scene.mechanics.pin_support((0, 0), size=0.6)),
    ("pin_support(p, size=1.2)", scene.mechanics.pin_support((0, 0), size=1.2)),
    ("fixed_support(p, size=1.2, ground_length=4)", scene.mechanics.fixed_support((0, 0), size=1.2, ground_length=4)),
    ("roller_support(p, size=1.2, color=accent)", scene.mechanics.roller_support((0, 0), size=1.2, color=accent)),
    ("guided_support(p, size=1.2, color=success)", scene.mechanics.guided_support((0, 0), size=1.2, color=success)),
]
grid = scene.layout.zones(Zones.grid(rows=len(SIZES), columns=[10, "1fr"], column_gap="48px", padding="page_padding"))
for index, (label, built) in enumerate(SIZES):
    scene.text(label, role="code").place(grid[2 * index], anchor="right")
    built.place(grid[2 * index + 1])
reveal([built for _, built in SIZES])

# Every part of a Support is a drawable of its own, styled like any other.
scene.segment("parts")
PARTS = [
    ("s.joint.stroke(accent, 0.12)", "roller", lambda s: s.joint.stroke(accent, 0.12)),
    ("s.body.fill(accent)", "roller", lambda s: s.body.fill(accent)),
    ("s.rollers.fill(accent)", "roller", lambda s: s.rollers.fill(accent)),
    ("s.ground.stroke(accent, 0.1)", "roller", lambda s: s.ground.stroke(accent, 0.1)),
    ("s.hatching.stroke(accent, 0.06)", "roller", lambda s: s.hatching.stroke(accent, 0.06)),
    ("s.guides.stroke(accent, 0.1)", "guided", lambda s: s.guides.stroke(accent, 0.1)),
]
styled = []
for label, kind, style in PARTS:
    part = support(kind, size=1.5)
    style(part)
    styled.append((label, part))
gallery(3, 2, styled)

if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    # Drawing and drawn, in each of the four segments.
    scene.snapshots(snapshots, [(DRAW + HOLD) * index + t for index in range(4) for t in (0.6, 1.6)])
else:
    scene.render()

"""Dimensions: linear dimensions with symbols and readouts, the side they sit on, their style and parts.

``scene.mechanics.dimension`` draws a fixed technical dimension;
``dimension_between`` follows two endpoints (points, drawables or anchor
points) and can read their distance, scaled and formatted, or any reactive
value. Every endpoint below moves, so the readouts change as they play.
Each label is the code that builds the variant.
"""

import os

from gaanim import Anchor, Easing, Scene, Zones, computed

MOVE = 2.0
HOLD = 0.6

scene = Scene()
M = scene.mechanics
fg = scene.canvas.color("foreground")
accent = scene.canvas.color("accent")
muted = scene.canvas.color("muted")
panel = scene.canvas.color("panel")


def label(zone, text, anchor="right"):
    """A code label in ``zone``; the lines of a long one align to the left."""
    scene.text(text, role="code", text_align="left").place(zone, anchor=anchor)


def table(entries, columns=("5fr", "3fr")):
    """One row per (label, build): ``build(a, b)`` dimensions a fixed point a and a point b that moves away."""
    grid = scene.layout.zones(Zones.grid(rows=len(entries), columns=list(columns), column_gap="64px", padding="page_padding"))
    reach = scene.viz.parameter(0.45)
    parts = []
    for index, (text, build) in enumerate(entries):
        label(grid[2 * index], text)
        stage = grid[2 * index + 1]
        y = stage.center[1] - 0.35
        a = (stage.left + 0.1, y)
        b = scene.geometry.dot(0.06).fill(muted).move_to(computed(lambda f, x=stage.left, w=stage.width: x + f * w - 0.1, inputs=[reach]), y)
        parts += [scene.geometry.dot(0.06).fill(muted).move_to(*a), b, build(a, b, stage)]
    scene.play([part.animate.fade_in() for part in parts], duration=0.4)
    scene.play([reach.animate.set(1.0).duration(MOVE).easing(Easing.SMOOTH)])
    scene.wait(HOLD)


def cells(rows, columns, label_height):
    """A padded grid of cells, each split into a label strip and a stage."""
    grid = scene.layout.zones(Zones.grid(rows=rows, columns=columns, gap="48px", padding="page_padding"))
    return [zone.split(Zones.rows([label_height, "1fr"], gap="16px")) for zone in grid]


# dimension() draws once between fixed coordinates; dimension_between follows
# its endpoints. label is a symbol, show_value reads the distance times
# scale, and value reads any number or Parameter instead.
scene.segment("linear")
length = scene.viz.parameter(120.0)
table([
    ("dimension(x1, y, x2, y, 0.5)", lambda a, b, stage: M.dimension(stage.left + 0.1, a[1], stage.right - 0.1, a[1], 0.5)),
    ("dimension_between(a, b, 0.5)", lambda a, b, stage: M.dimension_between(a, b, 0.5)),
    ('dimension_between(a, b, 0.5, label="$L$")', lambda a, b, stage: M.dimension_between(a, b, 0.5, label="$L$")),
    ('dimension_between(a, b, 0.5,\n    show_value=True, unit="m")', lambda a, b, stage: M.dimension_between(a, b, 0.5, show_value=True, unit="m")),
    (
        'dimension_between(a, b, 0.5,\n    label="$L$", value=length,\n    format=".0f", unit="mm")',
        lambda a, b, stage: M.dimension_between(a, b, 0.5, label="$L$", value=length, format=".0f", unit="mm"),
    ),
    (
        'dimension_between(a, b, 0.5,\n    show_value=True, scale=100,\n    format=".0f", unit="cm")',
        lambda a, b, stage: M.dimension_between(a, b, 0.5, show_value=True, scale=100, format=".0f", unit="cm"),
    ),
])

# The sign of offset puts the dimension to the left (+) or right (-) of the
# a -> b direction; side= fixes it in scene terms even if the ends swap.
# Each box grows, and its corners carry the dimensions with it.
CORNERS = {"tl": Anchor.TOP_LEFT, "tr": Anchor.TOP_RIGHT, "bl": Anchor.BOTTOM_LEFT, "br": Anchor.BOTTOM_RIGHT}


def boxes(variants):
    """Label each cell and dimension a box in it whose corners are tl, tr, bl and br."""
    parts, grows = [], []
    for (label_zone, stage), (text, build) in zip(cells(1, 3, 2.2), variants):
        label(label_zone, text, anchor="bottom")
        box = scene.geometry.rect(2.6, 1.4).fill(panel).stroke(muted, 0.03).move_to(stage.center[0] - 0.3, stage.center[1])
        parts += [box, build({name: box.anchor_point(anchor) for name, anchor in CORNERS.items()})]
        grows.append(box.animate.scale_to(1.3).duration(MOVE).easing(Easing.SMOOTH))
    scene.play([part.animate.fade_in() for part in parts], duration=0.4)
    scene.play(grows)
    scene.wait(HOLD)


scene.segment("placement")
boxes([
    ("dimension_between(\n  tl, tr, 0.4,\n  show_value=True)", lambda c: M.dimension_between(c["tl"], c["tr"], 0.4, show_value=True)),
    ("dimension_between(\n  tl, tr, -0.4,\n  show_value=True)", lambda c: M.dimension_between(c["tl"], c["tr"], -0.4, show_value=True)),
    (
        'dimension_between(\n  br, bl, 0.4,\n  side="below",\n  show_value=True)',
        lambda c: M.dimension_between(c["br"], c["bl"], 0.4, side="below", show_value=True),
    ),
])

# A vertical or slanted dimension keeps its label upright, moved out by the
# label's half width; label_orientation="aligned" turns it with the line.
scene.segment("orientation")
boxes([
    (
        'dimension_between(\n  br, tr, 0.4,\n  side="right",\n  show_value=True)',
        lambda c: M.dimension_between(c["br"], c["tr"], 0.4, side="right", show_value=True),
    ),
    ("dimension_between(\n  bl, tr, 0.4,\n  show_value=True)", lambda c: M.dimension_between(c["bl"], c["tr"], 0.4, show_value=True)),
    (
        'dimension_between(\n  bl, tr, 0.4,\n  show_value=True,\n  label_orientation=\n    "aligned")',
        lambda c: M.dimension_between(c["bl"], c["tr"], 0.4, show_value=True, label_orientation="aligned"),
    ),
])

# Style: line width, dashed extension lines, color, type size and face. SHOWN
# holds the shared keywords: label="$L$", show_value=True, unit="m".
scene.segment("style")
SHOWN = {"label": "$L$", "show_value": True, "unit": "m"}
table(
    [
        ("dimension_between(a, b, 0.5, **SHOWN,\n    line_width=0.06)", lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, line_width=0.06)),
        (
            'dimension_between(a, b, 0.5, **SHOWN,\n    extension_style="dashed")',
            lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, extension_style="dashed"),
        ),
        (
            'dimension_between(a, b, 0.5, **SHOWN,\n    extension_style="dashed",\n    dash_length=0.06, gap_length=0.06)',
            lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, extension_style="dashed", dash_length=0.06, gap_length=0.06),
        ),
        ("dimension_between(a, b, 0.5, **SHOWN,\n    color=accent)", lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, color=accent)),
        ("dimension_between(a, b, 0.5, **SHOWN,\n    font_size=0.7)", lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, font_size=0.7)),
        (
            'dimension_between(a, b, 0.5, **SHOWN,\n    font="JetBrains Mono", weight=700)',
            lambda a, b, s: M.dimension_between(a, b, 0.5, **SHOWN, font="JetBrains Mono", weight=700),
        ),
    ],
    columns=("5fr", "3fr"),
)

# A Dimension's line, extensions, label, number and unit are drawables of
# their own.
scene.segment("parts")


def styled(style):
    """A dimension with SHOWN, restyled by ``style``."""

    def build(a, b, stage):
        dimension = M.dimension_between(a, b, 0.5, **SHOWN)
        style(dimension)
        return dimension

    return build


table([
    ("d.line.fill(accent)", styled(lambda d: d.line.fill(accent))),
    ("d.extensions.opacity(0.3)", styled(lambda d: d.extensions.opacity(0.3))),
    ("d.label.fill(accent)", styled(lambda d: d.label.fill(accent))),
    ("d.number.fill(accent)", styled(lambda d: d.number.fill(accent))),
    ("d.unit.fill(muted)", styled(lambda d: d.unit.fill(muted))),
])


# Every segment part-way through its motion and settled.
step = 0.4 + MOVE + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(snapshots, [step * index + t for index in range(5) for t in (0.4 + MOVE / 3, step - 0.1)])
else:
    scene.render()

"""Angle dimensions: angles between rays with symbols and readouts, in every sweep, with their options.

``scene.mechanics.angle_between(vertex, from_, to)`` draws the arc of an angle
with arrowheads and short extension lines; each ray is a point or a
``Direction``, and a ray that moves turns the arc with it. Here the ray to p
turns while the ray to q stays put. Each label is the code that builds the
variant.
"""

import os

from gaanim import Direction, Easing, Scene, Zones

TURN = 2.4
HOLD = 0.6

scene = Scene()
M = scene.mechanics
accent = scene.canvas.color("accent")
muted = scene.canvas.color("muted")


def angles(variants, end=2.2, length=2.6, centered=False):
    """Label each cell and measure, at its vertex o, the angle from q (to its right) to p, which turns to ``end``."""
    grid = scene.layout.zones(Zones.columns(len(variants), gap="32px", padding="page_padding"))
    theta = scene.viz.parameter(0.5)
    parts = []
    for zone, (text, build) in zip(grid, variants):
        label_zone, stage = zone.split(Zones.rows([2.0, "1fr"], gap="16px"))
        scene.text(text, role="code", text_align="left").place(label_zone, anchor="bottom")
        if centered:
            o = stage.center
        else:
            o = (stage.center[0] - 0.3 * length, stage.center[1] - 0.4 * length)
        q = (o[0] + length, o[1])
        p = scene.geometry.polar_point(o, length, theta)
        parts += [
            scene.geometry.line(o, q).stroke(muted, 0.03),
            scene.geometry.tracking_line(o, p).stroke(muted, 0.03),
            scene.geometry.dot(0.08).fill(muted).follow(p),
            build(o, q, p),
        ]
    scene.play([part.animate.fade_in() for part in parts], duration=0.4)
    scene.play([theta.animate.set(end).duration(TURN).easing(Easing.SMOOTH)])
    scene.wait(HOLD)


# By default the arc takes the minor angle, at radius 0.64, with arrowheads
# at both ends. A readout sits label_gap beyond the arc, measured to the
# middle of the text, so a wide one needs a wide gap.
scene.segment("angles")
angles([
    ("angle_between(o, q, p)", lambda o, q, p: M.angle_between(o, q, p)),
    ('angle_between(o, q, p,\n  label="$theta$")', lambda o, q, p: M.angle_between(o, q, p, label="$theta$")),
    (
        "angle_between(o, q, p,\n  show_value=True,\n  label_gap=0.9)",
        lambda o, q, p: M.angle_between(o, q, p, show_value=True, label_gap=0.9),
    ),
])

# Readouts in radians, a larger radius, and a ray given as a Direction.
scene.segment("readouts")
angles([
    (
        'angle_between(o, q, p,\n  show_value=True,\n  unit="rad",\n  label_gap=0.9)',
        lambda o, q, p: M.angle_between(o, q, p, show_value=True, unit="rad", label_gap=0.9),
    ),
    (
        'angle_between(o, q, p,\n  radius=1.6,\n  label="$theta$")',
        lambda o, q, p: M.angle_between(o, q, p, radius=1.6, label="$theta$"),
    ),
    (
        'angle_between(o,\n  Direction.UP, p,\n  radius=1.2,\n  label="$phi$")',
        lambda o, q, p: M.angle_between(o, Direction.UP, p, radius=1.2, label="$phi$"),
    ),
])

# sweep picks the minor (default) or major angle, or the one turning
# clockwise ("cw") or counterclockwise ("ccw") from the first ray. As p
# passes half a turn, minor and major swap sides; cw and ccw do not.
scene.segment("sweeps")
angles(
    [
        ("angle_between(\n  o, q, p)", lambda o, q, p: M.angle_between(o, q, p)),
        ('angle_between(\n  o, q, p,\n  sweep="major")', lambda o, q, p: M.angle_between(o, q, p, sweep="major")),
        ('angle_between(\n  o, q, p,\n  sweep="cw")', lambda o, q, p: M.angle_between(o, q, p, sweep="cw")),
        ('angle_between(\n  o, q, p,\n  sweep="ccw")', lambda o, q, p: M.angle_between(o, q, p, sweep="ccw")),
    ],
    end=5.2,
    length=1.5,
    centered=True,
)

# Arrowheads at the "start", the "end", "both" or "none"; extension lines
# can go; color paints the arc and its label; aligned labels follow the arc.
scene.segment("style")
angles([
    ('angle_between(o, q, p,\n  arrowheads="end")', lambda o, q, p: M.angle_between(o, q, p, arrowheads="end")),
    (
        'angle_between(o, q, p,\n  arrowheads="none",\n  show_extensions=False)',
        lambda o, q, p: M.angle_between(o, q, p, arrowheads="none", show_extensions=False),
    ),
    (
        'angle_between(o, q, p,\n  label="$theta$",\n  label_orientation=\n    "aligned",\n  color=accent)',
        lambda o, q, p: M.angle_between(o, q, p, label="$theta$", label_orientation="aligned", color=accent),
    ),
])

# Every segment part-way through the turn and settled; the sweeps before
# and after p passes half a turn.
step = 0.4 + TURN + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(snapshots, [step * index + t for index in range(4) for t in (0.4 + TURN / 3, step - 0.1)])
else:
    scene.render()

"""Forces: vectors with readouts, forces from a magnitude or from components, moments, frames and a free body.

``vector_between`` draws an arrow between two endpoints with an optional
symbol and a readout of its length; ``force_at`` and
``force_from_components`` take physical magnitudes in newtons and a
``visual_scale`` that turns them into scene units, so the readout reports the
force, not the drawn length. Every input may be a reactive value, and every
part may be styled. Each label is the code that builds the variant.
"""

import math
import os

from gaanim import Anchor, Direction, Easing, Scene, Zones, computed

GROW = 2.0
HOLD = 0.6

scene = Scene()
M = scene.mechanics
fg = scene.canvas.color("foreground")
accent = scene.canvas.color("accent")
success = scene.canvas.color("success")
panel = scene.canvas.color("panel")


def rows(count):
    """A padded grid of ``count`` rows: a label column and a wider stage column."""
    return scene.layout.zones(Zones.grid(rows=count, columns=["5fr", "3fr"], column_gap="64px", padding="page_padding"))


def label(zone, text, anchor="right"):
    """A code label in ``zone``; the lines of a long one align to the left."""
    scene.text(text, role="code", text_align="left").place(zone, anchor=anchor)


def table(entries):
    """One row per (label, build): ``build(tail, tip)`` draws an arrow from tail to a tip that grows."""
    grid = rows(len(entries))
    reach = scene.viz.parameter(0.4)
    arrows = []
    for index, (text, build) in enumerate(entries):
        label(grid[2 * index], text)
        stage = grid[2 * index + 1]
        tail = (stage.left, stage.center[1] - 0.3)
        tip_x = computed(lambda f, x=stage.left, w=stage.width: x + f * w, inputs=[reach])
        arrows.append(build(tail, scene.geometry.point_ref(tip_x, tail[1])))
    scene.play([arrow.animate.fade_in() for arrow in arrows], duration=0.4)
    scene.play([reach.animate.set(1.0).duration(GROW).easing(Easing.SMOOTH)])
    scene.wait(HOLD)


def cells(count, label_height):
    """A padded row of ``count`` cells, each split into a label strip and a stage."""
    grid = scene.layout.zones(Zones.columns(count, gap="48px", padding="page_padding"))
    return [zone.split(Zones.rows([label_height, "1fr"], gap="16px")) for zone in grid]


# vector_between(a, b) draws a shaft and a solid head; label is a symbol (math
# between $ signs) and show_value reads the length, times scale. A readout
# sits label_gap from the shaft, measured to the middle of the text.
scene.segment("vectors")
table([
    ("vector_between(a, b)", lambda a, b: M.vector_between(a, b)),
    ('vector_between(a, b, label="$F$")', lambda a, b: M.vector_between(a, b, label="$F$")),
    (
        "vector_between(a, b, show_value=True,\n    label_gap=0.35)",
        lambda a, b: M.vector_between(a, b, show_value=True, label_gap=0.35),
    ),
    (
        'vector_between(a, b, label="$F$",\n    show_value=True, unit="N",\n    label_gap=0.35)',
        lambda a, b: M.vector_between(a, b, label="$F$", show_value=True, unit="N", label_gap=0.35),
    ),
    (
        'vector_between(a, b, show_value=True,\n    scale=100, format=".0f", unit="cm",\n    label_gap=0.35)',
        lambda a, b: M.vector_between(a, b, show_value=True, scale=100, format=".0f", unit="cm", label_gap=0.35),
    ),
])

# The readout's gap, size and color are options, and the shaft, head,
# label, number and unit of a ForceVector are drawables of their own.
scene.segment("styles")


def readout(style):
    """A vector with a symbol and a reading in newtons, restyled by ``style``."""

    def build(a, b):
        vector = M.vector_between(a, b, label="$F$", show_value=True, unit="N", label_gap=0.35)
        style(vector)
        return vector

    return build


table([
    ('vector_between(a, b, label="$F$",\n    label_gap=0.7)', lambda a, b: M.vector_between(a, b, label="$F$", label_gap=0.7)),
    (
        'vector_between(a, b, label="$F$",\n    font_size=0.8, label_gap=0.45)',
        lambda a, b: M.vector_between(a, b, label="$F$", font_size=0.8, label_gap=0.45),
    ),
    ('vector_between(a, b, label="$F$",\n    color=success)', lambda a, b: M.vector_between(a, b, label="$F$", color=success)),
    ("v.shaft.stroke(accent, 0.1)", readout(lambda v: v.shaft.stroke(accent, 0.1))),
    ("v.head.fill(accent)", readout(lambda v: v.head.fill(accent))),
    ("v.number.fill(accent)", readout(lambda v: v.number.fill(accent))),
])

# force_at takes a magnitude and a direction in radians; force_from_components
# takes x and y components. visual_scale is scene units per newton, so 40 N
# at 0.05 is 2 units long, and any input may be a Parameter. READOUT holds
# the shared keywords.
scene.segment("forces")
READOUT = {"visual_scale": 0.05, "show_value": True, "label_gap": 0.45}
magnitude = scene.viz.parameter(20.0)
angle = scene.viz.parameter(0.0)
fx, fy = scene.viz.parameter(50.0), scene.viz.parameter(-10.0)
FORCES = [
    ("force_at(o, 40, **READOUT)", 0.06, lambda o: M.force_at(o, 40, **READOUT)),
    ("force_at(o, magnitude, **READOUT)", 0.06, lambda o: M.force_at(o, magnitude, **READOUT)),
    (
        'force_at(o, 30, direction=angle,\n    visual_scale=0.05, label="$F$")',
        0.3,
        lambda o: M.force_at(o, 30, direction=angle, visual_scale=0.05, label="$F$", label_gap=0.3),
    ),
    ("force_from_components(o, fx, fy, **READOUT)", 0.06, lambda o: M.force_from_components(o, fx, fy, **READOUT)),
]
grid = scene.layout.zones(Zones.grid(rows=len(FORCES), columns=[9.4, "1fr"], column_gap="64px", padding="page_padding"))
parts = []
for index, (text, place, build) in enumerate(FORCES):
    label(grid[2 * index], text)
    stage = grid[2 * index + 1]
    body = scene.geometry.circle(0.2).fill(panel).stroke(fg, 0.04)
    body.move_to(stage.left + place * stage.width, stage.center[1] - 0.5)
    parts += [body, build(body)]
scene.play([part.animate.fade_in() for part in parts], duration=0.4)
scene.play([
    magnitude.animate.set(80.0).duration(GROW),
    angle.animate.set(1.3).duration(GROW),
    fx.animate.set(40.0).duration(GROW),
    fy.animate.set(25.0).duration(GROW),
])
scene.wait(HOLD)

# moment_about draws a curved arrow around a centre, counterclockwise unless
# direction="cw". It follows its centre: here every body moves right.
scene.segment("moments")
MOMENTS = [
    ("moment_about(c, 1.2)", lambda c: M.moment_about(c, 1.2)),
    ('moment_about(c, 1.2,\n  direction="cw")', lambda c: M.moment_about(c, 1.2, direction="cw")),
    ('moment_about(c, 1.2,\n  label="$M$",\n  color=accent)', lambda c: M.moment_about(c, 1.2, label="$M$", color=accent)),
]


def follow_bodies(variants):
    """Label each cell, put a body in it with its symbol, then move every body right."""
    parts, moves = [], []
    for (label_zone, stage), (text, build) in zip(cells(len(variants), 1.4), variants):
        label(label_zone, text, anchor="bottom")
        body = scene.geometry.circle(0.15).fill(fg).move_to(stage.center[0] - 0.6, stage.center[1])
        parts += [body, build(body)]
        moves.append(body.animate.shift_by(1.2, 0).duration(GROW).easing(Easing.SMOOTH))
    scene.play([part.animate.fade_in() for part in parts], duration=0.4)
    scene.play(moves)
    scene.wait(HOLD)


follow_bodies(MOMENTS)

# coordinate_frame_at draws an x axis toward x_direction and a y axis a
# quarter turn from it, of the given length.
scene.segment("frames")
follow_bodies([
    ("coordinate_frame_at(c,\n  Direction.RIGHT)", lambda c: M.coordinate_frame_at(c, Direction.RIGHT)),
    ("coordinate_frame_at(c,\n  Direction.UP_RIGHT,\n  length=1.6)", lambda c: M.coordinate_frame_at(c, Direction.UP_RIGHT, length=1.6)),
    ("coordinate_frame_at(c,\n  Direction.UP,\n  color=success)", lambda c: M.coordinate_frame_at(c, Direction.UP, length=1.6, color=success)),
])

# A free body: a pull that varies on its own (Parameter.add_updater_fn) drags
# the block, and the normal force and friction are computed from it. The
# normal force pushes up on the block from a point a reactive distance below.
scene.segment("free_body")
page = scene.layout.zones(Zones.columns([7.8, "1fr"], gap="48px", padding="page_padding"))
scene.text(
    "\n".join([
        "pull = scene.viz.parameter(40)",
        "pull.add_updater_fn(swing)",
        "n = computed(normal, inputs=[pull])",
        "f = computed(friction, inputs=[n])",
        "force_at(corner, pull, direction=A)",
        "force_from_components(below, 0, n)",
    ]),
    role="code",
    text_align="left",
    line_spacing=1.6,
).place(page[0], anchor="left")
stage = page[1]
W, A, MU, SCALE = 60.0, 0.35, 0.4, 0.04
block = scene.geometry.rounded_rect(2.6, 1.8, 0.1).fill(panel).stroke(fg, 0.05)
block.move_to(stage.left + 2.3, stage.center[1] + 0.8)
pull = scene.viz.parameter(40.0)


def swing(_value, _dt, elapsed):
    """A pull that rises and falls around 40 N."""
    return 40.0 + 15.0 * math.sin(2.5 * elapsed)


def normal(p):
    """The floor carries the weight less the lift of the pull."""
    return W - p * math.sin(A)


def friction(n):
    """Sliding friction, proportional to the normal force."""
    return MU * n


n = computed(normal, inputs=[pull])
f = computed(friction, inputs=[n])
contact = block.anchor_point(Anchor.BOTTOM, offset=(-0.7, 0))
below = scene.geometry.offset_point(contact, 0, computed(lambda value: -SCALE * value, inputs=[n]))
corner = block.anchor_point(Anchor.TOP_RIGHT)
loads = [
    M.force_at(corner, pull, direction=A, visual_scale=SCALE, label="$P$", show_value=True, label_gap=0.4, color=accent),
    M.force_at(block.anchor_point(Anchor.BOTTOM_LEFT), f, direction=math.pi, visual_scale=SCALE, label="$f$", label_gap=0.3),
    M.force_from_components(block, 0, -W, visual_scale=SCALE, label="$W$", label_gap=0.3),
    M.force_from_components(below, 0, n, visual_scale=SCALE, label="$N$", label_gap=0.3),
]
scene.play([block.animate.fade_in()] + [load.animate.fade_in() for load in loads], duration=0.5)
pull.add_updater_fn(swing)
scene.play([block.animate.shift_by(0.8, 0).duration(2 * GROW).easing(Easing.SMOOTH)])
pull.remove_updater()
scene.wait(HOLD)

# Each segment growing or moving, and settled; then the free body at
# three moments of its pull.
step = 0.4 + GROW + HOLD
body = 5 * step
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(
        snapshots,
        [step * index + t for index in range(5) for t in (0.4 + GROW / 2, step - 0.1)]
        + [body + 0.5 + 2 * GROW * f for f in (0.2, 0.5, 0.9)],
    )
else:
    scene.render()

"""Mechanisms: bars and joints between moving points, a slider-crank and a four-bar linkage.

``bar_between`` and ``joint_at`` follow their endpoints every frame, so a
mechanism is a few reactive points driven by one angle: ``polar_point`` puts
the crank pin on a circle, ``computed`` solves the other joints from the same
angle, and ``move_to`` and ``rotate_to`` bind rigid bodies to those values.
Labels are the code that builds each part.
"""

import math
import os

from gaanim import Direction, Easing, Scene, Zones, computed

HOLD = 0.6
TURN = 4.0

scene = Scene()
fg = scene.canvas.color("foreground")
accent = scene.canvas.color("accent")
muted = scene.canvas.color("muted")
panel = scene.canvas.color("panel")


def legend(zone, lines):
    """Set code ``lines`` as one left-aligned block at the left of ``zone``."""
    scene.text("\n".join(lines), role="code", text_align="left", line_spacing=1.6).place(zone, anchor="left")


def turn(theta):
    """Turn ``theta`` through one revolution at a steady speed, then hold."""
    scene.play([theta.animate.set(math.tau).duration(TURN).easing(Easing.LINEAR)])
    scene.wait(HOLD)


# bar_between draws a round-capped bar between two endpoints (give it a color
# and width with stroke), and joint_at draws a joint that follows one; both
# track b as it rises and falls. Rows with joints add them to a bar like the
# first row's.
scene.segment("bars_and_joints")
ROWS = [
    "bar_between(a, b).stroke(fg, 0.12)",
    "bar_between(a, b).stroke(accent, 0.3)",
    "joint_at(a), joint_at(b)",
    "joint_at(b, size=0.8, color=accent)",
    'joint_at(b, kind="prismatic", size=0.8,\n         axis=Direction.UP)',
]
grid = scene.layout.zones(Zones.grid(rows=len(ROWS), columns=["5fr", "4fr"], column_gap="64px", padding="page_padding"))
lift = scene.viz.parameter(0.0)
parts = []
for index, label in enumerate(ROWS):
    scene.text(label, role="code", text_align="left").place(grid[2 * index], anchor="right")
    stage = grid[2 * index + 1].inset((0, 0.4))
    (left, y), (right, _) = stage.point("left"), stage.point("right")
    a = (left, y)
    b = scene.geometry.offset_point((right, y), 0, lift)
    width, color = (0.3, accent) if index == 1 else (0.12, fg)
    parts.append(scene.mechanics.bar_between(a, b).stroke(color, width))
    if index == 2:
        parts += [scene.mechanics.joint_at(a), scene.mechanics.joint_at(b)]
    elif index == 3:
        parts.append(scene.mechanics.joint_at(b, size=0.8, color=accent))
    elif index == 4:
        # The prismatic joint slides along a guide drawn through b's travel.
        parts.append(scene.geometry.line(right, y - 0.75, right, y + 0.75).stroke(muted, 0.03))
        parts.append(scene.mechanics.joint_at(b, kind="prismatic", size=0.8, axis=Direction.UP))
scene.play([part.animate.fade_in() for part in parts], duration=0.4)
for height in (0.6, -0.6, 0.0):
    scene.play([lift.animate.set(height).duration(1.0).easing(Easing.SMOOTH)])
scene.wait(HOLD)

# A slider-crank: the crank pin turns on a circle and the slider sits where the
# connecting rod of length L reaches the rail. The wheel turns with rotate_to.
scene.segment("slider_crank")
page = scene.layout.zones(Zones.rows([1.9, "1fr"], gap="16px", padding="page_padding"))
legend(page[0], [
    "tip = polar_point(o, R, theta)",
    "x = computed(slider_x, inputs=[theta])",
    "block.move_to(x, y)    wheel.rotate_to(theta)",
])
stage = page[1]
R, L = 1.8, 5.4
# Centre the span from the wheel's left edge to the block's furthest reach.
ox, oy = stage.center[0] - (L + 0.45) / 2, stage.center[1] + 0.6
theta = scene.viz.parameter(0.0)
o = (ox, oy)


def slider_x(angle):
    """Where a rod of length L from the crank pin meets the rail through o."""
    return ox + R * math.cos(angle) + math.sqrt(L**2 - (R * math.sin(angle)) ** 2)


tip = scene.geometry.polar_point(o, R, theta)
x = computed(slider_x, inputs=[theta])
pin = scene.geometry.point_ref(x, oy)
rail = scene.mechanics.fixed_support((ox + L, oy - 0.45), size=0.9, ground_length=2 * R + 1.8)
wheel = scene.geometry.circle(R + 0.3).fill(panel).stroke(muted, 0.04).move_to(ox, oy)
spokes = [scene.geometry.line(ox - R - 0.3, oy, ox + R + 0.3, oy), scene.geometry.line(ox, oy - R - 0.3, ox, oy + R + 0.3)]
hub = scene.geometry.group([wheel, *[spoke.stroke(muted, 0.04) for spoke in spokes]])
hub.rotate_to(theta)
block = scene.geometry.rounded_rect(1.5, 0.9, 0.1).fill(panel).stroke(fg, 0.05).move_to(x, oy)
mechanism = [
    scene.mechanics.pin_support(o, size=0.9, ground_length=1.6),
    scene.mechanics.bar_between(o, tip).stroke(accent, 0.2),
    scene.mechanics.bar_between(tip, pin).stroke(fg, 0.14),
    scene.mechanics.joint_at(tip, size=0.5),
    scene.mechanics.joint_at(pin, size=0.5),
    scene.mechanics.angle_between(o, Direction.RIGHT, tip, radius=0.7, sweep="ccw", label="$theta$", label_gap=0.25, arrowheads="end"),
    scene.mechanics.dimension_between(o, pin, R + 0.6, side="below", label="$x$", show_value=True, label_gap=0.15),
]
scene.play([hub.animate.fade_in(), rail.animate.create(), block.animate.fade_in()] + [part.animate.fade_in() for part in mechanism], duration=0.5)
turn(theta)

# A four-bar linkage: the rocker pin is where circles about the crank pin and
# the right pivot cross, and a pen on the coupler draws its coupler curve.
scene.segment("four_bar")
page = scene.layout.zones(Zones.columns([6.0, "1fr"], gap="32px", padding="page_padding"))
legend(page[0], [
    "b = polar_point(a, r, theta)",
    "c = point_ref(cx, cy)",
    "pen = dot().move_to(px, py)",
    "traced_path(pen)",
])
stage = page[1]
CRANK, COUPLER, ROCKER, GROUND = 1.7, 5.1, 4.0, 5.7
ax, ay = stage.left + CRANK + 0.2, stage.bottom + 2.4
dx = ax + GROUND
theta = scene.viz.parameter(0.0)


def joints(angle):
    """Crank pin, rocker pin and coupler pen of the linkage at ``angle``."""
    bx, by = ax + CRANK * math.cos(angle), ay + CRANK * math.sin(angle)
    ex, ey = dx - bx, ay - by
    span = math.hypot(ex, ey)
    along = (COUPLER**2 - ROCKER**2 + span**2) / (2 * span)
    across = math.sqrt(max(COUPLER**2 - along**2, 0.0))
    cx = bx + (along * ex - across * ey) / span
    cy = by + (along * ey + across * ex) / span
    # The pen sits above the middle of the coupler, like a triangular plate.
    ux, uy = (cx - bx) / COUPLER, (cy - by) / COUPLER
    return (cx, cy), (bx + (cx - bx) * 0.5 - uy * 1.5, by + (cy - by) * 0.5 + ux * 1.5)


cx, cy = (computed(lambda angle, i=i: joints(angle)[0][i], inputs=[theta]) for i in (0, 1))
px, py = (computed(lambda angle, i=i: joints(angle)[1][i], inputs=[theta]) for i in (0, 1))
a, d = (ax, ay), (dx, ay)
b = scene.geometry.polar_point(a, CRANK, theta)
c = scene.geometry.point_ref(cx, cy)
pen = scene.geometry.dot(0.12).fill(accent).move_to(px, py)
trail = scene.geometry.traced_path(pen).no_fill().stroke(accent, 0.04)
linkage = [
    scene.mechanics.pin_support(a, size=0.9, ground_length=1.6),
    scene.mechanics.pin_support(d, size=0.9, ground_length=1.6),
    scene.mechanics.bar_between(a, b).stroke(accent, 0.2),
    scene.mechanics.bar_between(b, c).stroke(fg, 0.14),
    scene.mechanics.bar_between(b, pen).stroke(fg, 0.08),
    scene.mechanics.bar_between(c, pen).stroke(fg, 0.08),
    scene.mechanics.bar_between(d, c).stroke(fg, 0.14),
    scene.mechanics.joint_at(b, size=0.5),
    scene.mechanics.joint_at(c, size=0.5),
]
scene.play([part.animate.fade_in() for part in linkage + [trail]] + [pen.animate.fade_in()], duration=0.5)
turn(theta)

# The bars rising, falling and level; then each mechanism at three points of
# its turn, the four-bar's last one with its coupler curve closed.
slider_start = 0.4 + 3.0 + HOLD
four_bar_start = slider_start + 0.5 + TURN + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(
        snapshots,
        [1.4, 2.4, 3.6]
        + [slider_start + 0.5 + TURN * f for f in (0.15, 0.5, 0.8)]
        + [four_bar_start + 0.5 + TURN * f for f in (0.3, 0.6, 1.0)],
    )
else:
    scene.render()

"""Springs: helical springs between moving endpoints, with their coils, radius, crossings and straight ends.

``scene.mechanics.spring_between`` regenerates a helix every frame between
two endpoints, keeping its radius while its pitch stretches with the
distance. Endpoints may be points, drawables or anchor points of moving and
rotating bodies. Each label is the code that builds the variant.
"""

import os

from gaanim import Anchor, Direction, Easing, Scene, Zones

STRETCH = 1.2
HOLD = 0.6

scene = Scene()
M = scene.mechanics
fg = scene.canvas.color("foreground")
panel = scene.canvas.color("panel")


def block(size=0.8):
    """A mass: a square block with the theme's panel fill."""
    return scene.geometry.rounded_rect(size, size, 0.08).fill(panel).stroke(fg, 0.04)


# Every row hangs a spring from a wall to the left face of a block that
# moves out, in and back. amplitude is the helix radius, crossing (0 to 1)
# folds each turn back into a loop, and the straight lengths sit at the ends.
scene.segment("options")
OPTIONS = [
    ("spring_between(wall, face)", {}),
    ("spring_between(wall, face, coils=4)", {"coils": 4}),
    ("spring_between(wall, face, coils=14)", {"coils": 14}),
    ("spring_between(wall, face, amplitude=0.3)", {"amplitude": 0.3}),
    ("spring_between(wall, face, crossing=0.5)", {"crossing": 0.5}),
    ("spring_between(wall, face, crossing=1.0)", {"crossing": 1.0}),
    ("spring_between(wall, face,\n    start_straight=0.8, end_straight=0.4)", {"start_straight": 0.8, "end_straight": 0.4}),
]
grid = scene.layout.zones(Zones.grid(rows=len(OPTIONS), columns=[9.0, "1fr"], column_gap="64px", padding="page_padding"))
parts, blocks = [], []
for index, (text, options) in enumerate(OPTIONS):
    scene.text(text, role="code", text_align="left").place(grid[2 * index], anchor="right")
    stage = grid[2 * index + 1]
    wall = (stage.left + 0.3, stage.center[1])
    mass = block(0.7).move_to(stage.left + 3.0, stage.center[1])
    parts += [
        M.fixed_support(wall, direction=Direction.RIGHT, size=0.7, ground_length=0.9),
        mass,
        M.spring_between(wall, mass.anchor_point(Anchor.LEFT), **options),
    ]
    blocks.append(mass)
scene.play([part.animate.fade_in() for part in parts], duration=0.4)
for shift in (1.4, -2.2, 0.8):
    scene.play([mass.animate.shift_by(shift, 0).duration(STRETCH).easing(Easing.SMOOTH) for mass in blocks])
scene.wait(HOLD)

# Springs in systems: a block that drops and bounces from a ceiling, two
# blocks between walls (w) that move apart, and a spring that holds the tip
# of a lever as it turns. Anchor points keep each end on its body.
scene.segment("systems")
cells = scene.layout.zones(Zones.columns(3, gap="48px", padding="page_padding"))
labels = [
    "spring_between(\n  ceiling,\n  m.anchor_point(\n    Anchor.TOP))",
    "spring_between(w, a)\nspring_between(a, b)\nspring_between(b, w)",
    "spring_between(\n  ground,\n  lever.anchor_point(\n    Anchor.RIGHT))",
]
stages = []
for zone, text in zip(cells, labels):
    label_zone, stage = zone.split(Zones.rows([2.0, "1fr"], gap="24px"))
    scene.text(text, role="code", text_align="left").place(label_zone, anchor="bottom")
    stages.append(stage)

# A block hanging from a ceiling, dropped with a springy easing.
stage = stages[0]
cx = stage.center[0]
ceiling = (cx, stage.top - 0.2)
hanging = block(1.0).move_to(cx, stage.top - 2.6)
hang = [
    M.fixed_support(ceiling, direction=Direction.DOWN, size=0.9, ground_length=2.4),
    hanging,
    M.spring_between(ceiling, hanging.anchor_point(Anchor.TOP), coils=10, amplitude=0.3),
]

# Two blocks between walls, joined in series.
stage = stages[1]
y = stage.center[1] + 0.6
left, right = (stage.left + 0.2, y), (stage.right - 0.2, y)
a = block(0.7).move_to(stage.left + 1.6, y)
b = block(0.7).move_to(stage.right - 1.6, y)
series = [
    M.fixed_support(left, direction=Direction.RIGHT, size=0.7, ground_length=1.0),
    M.fixed_support(right, direction=Direction.LEFT, size=0.7, ground_length=1.0),
    a,
    b,
    M.spring_between(left, a.anchor_point(Anchor.LEFT), coils=5, amplitude=0.2),
    M.spring_between(a.anchor_point(Anchor.RIGHT), b.anchor_point(Anchor.LEFT), coils=5, amplitude=0.2),
    M.spring_between(b.anchor_point(Anchor.RIGHT), right, coils=5, amplitude=0.2),
]

# A lever on a pin, its tip held by a spring to the ground below.
stage = stages[2]
pivot = (stage.left + 0.6, stage.center[1] + 0.6)
lever = scene.geometry.rounded_rect(3.2, 0.24, 0.12).fill(panel).stroke(fg, 0.04)
lever.move_to(pivot[0] + 1.5, pivot[1]).pivot(*pivot)
ground = (pivot[0] + 3.0, stage.bottom + 0.6)
lever_parts = [
    M.pin_support(pivot, size=0.8, ground_length=1.2),
    M.fixed_support(ground, size=0.8, ground_length=1.6),
    lever,
    M.joint_at(pivot, size=0.5),
    M.spring_between(ground, lever.anchor_point(Anchor.RIGHT), coils=8, amplitude=0.25, crossing=1.0),
]

scene.play([part.animate.fade_in() for part in hang + series + lever_parts], duration=0.4)
scene.play([
    hanging.animate.shift_by(0, -1.4).duration(2 * STRETCH).easing(Easing.spring(bounce=0.6)),
    a.animate.shift_by(-0.4, 0).duration(STRETCH).easing(Easing.THERE_AND_BACK),
    b.animate.shift_by(0.4, 0).duration(STRETCH).easing(Easing.THERE_AND_BACK),
    lever.animate.rotate_by(-0.5).duration(STRETCH).easing(Easing.SMOOTH),
])
scene.play([lever.animate.rotate_by(0.8).duration(STRETCH).easing(Easing.SMOOTH)])
scene.wait(HOLD)

# The blocks out, in and back, then the systems mid-swing, at the turn
# of the lever, and settled.
options_end = 0.4 + 3 * STRETCH + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(
        snapshots,
        [0.4 + STRETCH, 0.4 + 2 * STRETCH, options_end - 0.1]
        + [options_end + 0.4 + t for t in (STRETCH / 2, STRETCH, 2 * STRETCH + HOLD - 0.1)],
    )
else:
    scene.render()

"""Transmissions: gears and their options, gear trains, racks and pinions, cams with followers and contact.

``scene.mechanics.gear``, ``rack`` and ``cam_profile`` draw editorial
silhouettes; the motion between them is a visual coupling, not a solver.
``bind_rotation_from`` turns a gear from another with a ratio and a phase,
``bind_translation_from_rotation`` slides a rack as a pinion turns, and a cam
bound to an angle with ``rotate_to`` lifts a follower placed by ``computed``.
Each label is the code that builds the variant.
"""

import math
import os

from gaanim import Direction, Easing, Scene, Zones, computed

HOLD = 0.6
SPIN = 3.0

scene = Scene()
fg = scene.canvas.color("foreground")
accent = scene.canvas.color("accent")
muted = scene.canvas.color("muted")
panel = scene.canvas.color("panel")


def cells(rows, columns, label_height=0.5):
    """A padded grid of cells, each split into a label strip and a stage."""
    grid = scene.layout.zones(Zones.grid(rows=rows, columns=columns, gap="48px", padding="page_padding"))
    return [zone.split(Zones.rows([label_height, "1fr"], gap="16px")) for zone in grid]


def legend(zone, lines):
    """Set code ``lines`` as one left-aligned block at the left of ``zone``."""
    scene.text("\n".join(lines), role="code", text_align="left", line_spacing=1.6).place(zone, anchor="left")


def radius(teeth):
    """Gears that mesh share a tooth pitch, so their radii follow their tooth counts."""
    return 0.07 * teeth


def mesh_phase(source_teeth, teeth, angle):
    """The phase that puts a tooth of the source in a gap of a gear at ``angle`` from it."""
    return angle + math.pi + (source_teeth * angle - 2.5 * math.pi) / teeth


# gear(radius, teeth) draws a toothed rim and a bore. The teeth add an eighth
# of the radius outside it; bore_radius and color are options.
scene.segment("gears")
GEARS = [
    ("gear(1.1, 12)", lambda: scene.mechanics.gear(1.1, 12)),
    ("gear(1.1, 24)", lambda: scene.mechanics.gear(1.1, 24)),
    ("gear(1.1, 48)", lambda: scene.mechanics.gear(1.1, 48)),
    ("gear(1.1, 24,\n     bore_radius=0.6)", lambda: scene.mechanics.gear(1.1, 24, bore_radius=0.6)),
    ("gear(1.1, 24,\n     color=accent)", lambda: scene.mechanics.gear(1.1, 24, color=accent)),
    ("gear(0.55, 12)", lambda: scene.mechanics.gear(0.55, 12)),
]
gears = []
for (label_zone, stage), (label, build) in zip(cells(2, 3, label_height=0.9), GEARS):
    scene.text(label, role="code", text_align="left").place(label_zone, anchor="bottom")
    gears.append(build().move_to(*stage.center))
scene.play([gear.animate.create() for gear in gears], duration=SPIN / 2)
scene.wait(HOLD)

# bind_rotation_from copies the source's rotation times ratio, plus phase.
# Meshing gears turn the other way at the ratio of their tooth counts, and
# the phase sets a tooth of one in a gap of the other.
scene.segment("gear_train")
page = scene.layout.zones(Zones.rows([2.0, "1fr"], gap="16px", padding="page_padding"))
TEETH = {"a": 24, "b": 12, "c": 36}
ratio_b, phase_b = -TEETH["a"] / TEETH["b"], mesh_phase(TEETH["a"], TEETH["b"], 0.0)
ratio_c, phase_c = -TEETH["b"] / TEETH["c"], mesh_phase(TEETH["b"], TEETH["c"], 0.0)
legend(page[0], [
    "a.animate.rotate_by(math.tau)",
    f"b.bind_rotation_from(a, ratio={ratio_b:g}, phase={phase_b:.3f})",
    f"c.bind_rotation_from(b, ratio={ratio_c:.3f}, phase={phase_c:.3f})",
])
stage = page[1]
spans = {name: radius(teeth) * 1.065 for name, teeth in TEETH.items()}
width = 2 * sum(spans.values())
x = stage.center[0] - width / 2
centers = {}
for name in "abc":
    centers[name] = (x + spans[name], stage.center[1])
    x += 2 * spans[name]
a = scene.mechanics.gear(radius(TEETH["a"]), TEETH["a"], bore_radius=0.32, color=accent).move_to(*centers["a"])
b = scene.mechanics.gear(radius(TEETH["b"]), TEETH["b"], bore_radius=0.32).move_to(*centers["b"])
c = scene.mechanics.gear(radius(TEETH["c"]), TEETH["c"], bore_radius=0.32).move_to(*centers["c"])
b.bind_rotation_from(a, ratio=ratio_b, phase=phase_b)
c.bind_rotation_from(b, ratio=ratio_c, phase=phase_c)
names = [scene.text(name, role="code").move_to(*centers[name]) for name in "abc"]
scene.play([gear.animate.fade_in() for gear in (a, b, c, *names)], duration=0.4)
scene.play([a.animate.rotate_by(math.tau).duration(2 * SPIN).easing(Easing.LINEAR)])
scene.wait(HOLD)

# bind_translation_from_rotation slides a drawable along an axis by the
# source's turn times scale: at the pitch radius r, rack and pinion roll
# together. The pinion turns one way and then back.
scene.segment("rack_and_pinion")
page = scene.layout.zones(Zones.columns(2, gap="32px", padding="page_padding"))
legend(page[0], [
    "pinion.animate.rotate_by(0.9)",
    "low.bind_translation_from_rotation(",
    "    pinion, axis=Direction.RIGHT, scale=r)",
    "side.bind_translation_from_rotation(",
    "    pinion, axis=Direction.UP, scale=r)",
])
TEETH_PINION, RACK_TEETH, TURN = 24, 9, 0.9
pitch_radius = radius(TEETH_PINION) * 1.065
pitch = math.tau * pitch_radius / TEETH_PINION
height = 1.25 * pitch
gx, gy = page[1].center[0] + 0.3, page[1].center[1] - 0.7
# A gap of the pinion faces down and right, where a tooth of each rack sits;
# the racks start two teeth back so that they stay in the frame.
pinion = scene.mechanics.gear(radius(TEETH_PINION), TEETH_PINION, bore_radius=0.32, color=accent)
pinion.move_to(gx, gy).rotate_to(-math.pi / 2 - 7 * math.pi / (4 * TEETH_PINION))
# A rack's tips lie 0.725 of its height from its centre.
tips = radius(TEETH_PINION) + 0.01 + 0.725 * height
low = scene.mechanics.rack(RACK_TEETH * pitch, RACK_TEETH).move_to(gx - 2 * pitch, gy - tips)
side = scene.mechanics.rack(RACK_TEETH * pitch, RACK_TEETH).rotate_to(math.pi / 2).move_to(gx + tips, gy + 2 * pitch)
low.bind_translation_from_rotation(pinion, axis=Direction.RIGHT, scale=pitch_radius)
side.bind_translation_from_rotation(pinion, axis=Direction.UP, scale=pitch_radius)
scene.play([part.animate.fade_in() for part in (pinion, low, side)], duration=0.4)
scene.play([pinion.animate.rotate_by(TURN).duration(SPIN / 2).easing(Easing.SMOOTH)])
scene.play([pinion.animate.rotate_by(-TURN).duration(SPIN / 2).easing(Easing.SMOOTH)])
scene.wait(HOLD)

# cam_profile takes (angle, radius) samples. Each cam turns with rotate_to(turn)
# and its follower rests on the profile straight above the shaft, through a
# prismatic guide.
scene.segment("cams")
SAMPLES = 120
ANGLES = [math.tau * i / SAMPLES for i in range(SAMPLES)]


def eccentric(angle):
    """A disc of radius 1 whose centre sits 0.4 off the shaft."""
    return 0.4 * math.cos(angle) + math.sqrt(1 - (0.4 * math.sin(angle)) ** 2)


def lobes(angle):
    """Two lobes: the follower rises and falls twice per turn."""
    return 0.95 + 0.35 * math.cos(2 * angle)


def heart(angle):
    """A rise and fall at constant speed."""
    return 0.7 + 0.7 * (1 - abs(angle - math.pi) / math.pi)


CAMS = [("cam_profile(ECCENTRIC)", eccentric), ("cam_profile(LOBES)", lobes), ("cam_profile(HEART)", heart)]
SCALE = 1.4
turn = scene.viz.parameter(0.0)
parts = []
for (label_zone, stage), (label, profile) in zip(cells(1, 3), CAMS):
    scene.text(label, role="code").place(label_zone)
    cx, cy = stage.center[0], stage.bottom + 2.2
    cam = scene.mechanics.cam_profile([(angle, SCALE * profile(angle)) for angle in ANGLES], bore_radius=0.2)
    cam.move_to(cx, cy).rotate_to(turn)
    # The profile under the follower is the one a quarter turn ahead of the cam's angle.
    lift = computed(lambda t, f=profile: SCALE * f((math.pi / 2 - t) % math.tau), inputs=[turn])
    rod = scene.geometry.rounded_rect(0.3, 2.4, 0.08).fill(panel).stroke(fg, 0.05)
    rod.move_to(cx, computed(lambda r, y=cy: y + r + 1.2, inputs=[lift]))
    guide = scene.mechanics.joint_at((cx, cy + 2.8), kind="prismatic", axis=Direction.UP, size=0.9)
    parts += [cam, rod, guide]
scene.play([part.animate.fade_in() for part in parts], duration=0.4)
scene.play([turn.animate.set(math.tau).duration(2 * SPIN).easing(Easing.LINEAR)])
scene.wait(HOLD)

# contact_on_curve marks a point on a curve with its tangent and normal, at
# a fraction t of the way round it; the lengths are options.
scene.segment("contact")
OUTLINE = [(2.2 * math.cos(angle) * lobes(angle), 2.2 * math.sin(angle) * lobes(angle)) for angle in ANGLES]
CONTACTS = [
    ("contact_on_curve(outline, t)", {}),
    (
        "contact_on_curve(outline, t,\n    tangent_length=2.4,\n    normal_length=1.6)",
        {"tangent_length": 2.4, "normal_length": 1.6},
    ),
]
t = scene.viz.parameter(0.0)
parts = []
for (label_zone, stage), (label, options) in zip(cells(1, 2, label_height=1.6), CONTACTS):
    scene.text(label, role="code", text_align="left").place(label_zone, anchor="bottom")
    outline = scene.geometry.polygon([(stage.center[0] + px, stage.center[1] + py) for px, py in OUTLINE])
    outline.no_fill().stroke(muted, 0.04)
    parts += [outline, scene.mechanics.contact_on_curve(outline, t, **options)]
scene.play([part.animate.fade_in() for part in parts], duration=0.4)
scene.play([t.animate.set(1.0).duration(2 * SPIN).easing(Easing.LINEAR)])
scene.wait(HOLD)

# Gears being drawn and drawn; the train at a third and two thirds of a
# turn; the racks out and back; the cams and the contact point at a
# quarter, a half and most of the way round.
gears_end = SPIN / 2 + HOLD
train_end = gears_end + 0.4 + 2 * SPIN + HOLD
racks_end = train_end + 0.4 + SPIN + HOLD
cams_end = racks_end + 0.4 + 2 * SPIN + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(
        snapshots,
        [SPIN / 4, gears_end - 0.1]
        + [gears_end + 0.4 + 2 * SPIN * f for f in (0.33, 0.66)]
        + [train_end + 0.4 + SPIN * f for f in (0.5, 1.0)]
        + [racks_end + 0.4 + 2 * SPIN * f for f in (0.25, 0.5, 0.8)]
        + [cams_end + 0.4 + 2 * SPIN * f for f in (0.25, 0.6)],
    )
else:
    scene.render()

"""Simulations: a damped pendulum, a double pendulum and a mass on a spring, stepped at a fixed rate.

``add_updater_fn(step, reset=reset, fixed_dt=...)`` integrates state that
cannot be written as a function of time. Gaanim calls ``reset`` and replays
the fixed steps whenever it seeks or exports, so every frame is the same in
the editor, in snapshots and in a video. Rods, springs, angles, forces and
traces built on the simulated bodies follow them in the same frame. The code
on the left of each segment is how it is built.
"""

import math
import os

from gaanim import Anchor, Direction, Scene, Zones

STEP = 1 / 240
HOLD = 0.4

scene = Scene()
M = scene.mechanics
fg = scene.canvas.color("foreground")
accent = scene.canvas.color("accent")
muted = scene.canvas.color("muted")
panel = scene.canvas.color("panel")


def page(lines):
    """Set the code ``lines`` in a left column and return the stage to its right."""
    zones = scene.layout.zones(Zones.columns([7.2, "1fr"], gap="48px", padding="page_padding"))
    scene.text("\n".join(lines), role="code", text_align="left", line_spacing=1.6).place(zones[0], anchor="left")
    return zones[1]


# A damped pendulum: step() advances the angle and its speed by one fixed
# step and returns the bob's position; reset() restores the start.
scene.segment("pendulum")
stage = page([
    "bob.add_updater_fn(step,",
    "    reset=reset, fixed_dt=1 / 240)",
    "bar_between(hinge, bob)",
    "angle_between(hinge,",
    "    Direction.DOWN, bob)",
    "traced_path(bob,",
    "    dissipating_time=1.5)",
])
LENGTH, GRAVITY, DAMPING, START = 3.2, 9.81, 0.25, 1.1
hinge = (stage.center[0], stage.top - 1.2)
pendulum = {"angle": START, "speed": 0.0}


def bob_at(angle):
    """The bob's position for a rod at ``angle`` from the vertical."""
    return (hinge[0] + LENGTH * math.sin(angle), hinge[1] - LENGTH * math.cos(angle), 0.0)


def reset_pendulum():
    """Restore the starting angle and speed."""
    pendulum.update(angle=START, speed=0.0)


def step_pendulum(_position, dt, _elapsed):
    """One semi-implicit Euler step of a damped pendulum."""
    pendulum["speed"] += dt * (-GRAVITY / LENGTH * math.sin(pendulum["angle"]) - DAMPING * pendulum["speed"])
    pendulum["angle"] += dt * pendulum["speed"]
    return bob_at(pendulum["angle"])


bob = scene.geometry.circle(0.3).fill(accent).move_to(*bob_at(START)[:2])
parts = [
    M.pin_support(hinge, direction=Direction.DOWN, size=0.9, ground_length=1.8),
    M.bar_between(hinge, bob).stroke(fg, 0.08),
    M.angle_between(hinge, Direction.DOWN, bob, radius=1.2, label="$theta$", label_gap=0.25),
    scene.geometry.traced_path(bob, dissipating_time=1.5).no_fill().stroke(accent, 0.04),
]
bob.add_updater_fn(step_pendulum, reset=reset_pendulum, fixed_dt=STEP)
scene.play([part.animate.fade_in() for part in parts] + [bob.animate.fade_in()], duration=0.3)
scene.wait(5.0)
bob.remove_updater()
scene.wait(HOLD)

# A double pendulum: only the upper bob steps the shared state; the lower
# one reads it in the same frame, so the system is integrated once.
scene.segment("double_pendulum")
stage = page([
    "upper.add_updater_fn(step,",
    "    reset=reset, fixed_dt=1 / 240)",
    "lower.add_updater_fn(lower_bob)",
    "bar_between(hinge, upper)",
    "bar_between(upper, lower)",
    "traced_path(lower,",
    "    dissipating_time=3.0)",
])
ARM, DOUBLE_START = 1.65, (2.0, 2.7)
pivot = (stage.center[0], stage.center[1] + 0.3)
double = {}


def reset_double():
    """Both arms at rest at their starting angles."""
    double.update(a1=DOUBLE_START[0], w1=0.0, a2=DOUBLE_START[1], w2=0.0)


def joints():
    """The positions of both bobs from the shared state."""
    x1 = pivot[0] + ARM * math.sin(double["a1"])
    y1 = pivot[1] - ARM * math.cos(double["a1"])
    return (x1, y1, 0.0), (x1 + ARM * math.sin(double["a2"]), y1 - ARM * math.cos(double["a2"]), 0.0)


def step_double(_position, dt, _elapsed):
    """One semi-implicit Euler step of two equal masses on equal arms."""
    a1, w1, a2, w2 = double["a1"], double["w1"], double["a2"], double["w2"]
    delta = a1 - a2
    den = 3.0 - math.cos(2.0 * delta)
    alpha1 = (-3.0 * GRAVITY * math.sin(a1) - GRAVITY * math.sin(a1 - 2.0 * a2)
              - 2.0 * math.sin(delta) * (w2 * w2 * ARM + w1 * w1 * ARM * math.cos(delta))) / (ARM * den)
    alpha2 = (2.0 * math.sin(delta) * (2.0 * w1 * w1 * ARM + 2.0 * GRAVITY * math.cos(a1)
              + w2 * w2 * ARM * math.cos(delta))) / (ARM * den)
    double["w1"] = w1 + dt * (alpha1 - 0.02 * w1)
    double["w2"] = w2 + dt * (alpha2 - 0.02 * w2)
    double["a1"] = a1 + dt * double["w1"]
    double["a2"] = a2 + dt * double["w2"]
    return joints()[0]


def lower_bob(_position, _dt, _elapsed):
    """Read the lower bob from the state the upper bob just stepped."""
    return joints()[1]


reset_double()
start_upper, start_lower = joints()
upper = scene.geometry.circle(0.22).fill(fg).move_to(*start_upper[:2])
lower = scene.geometry.circle(0.28).fill(accent).move_to(*start_lower[:2])
parts = [
    M.pin_support(pivot, size=0.8, ground_length=1.4),
    M.bar_between(pivot, upper).stroke(fg, 0.07),
    M.bar_between(upper, lower).stroke(fg, 0.07),
    scene.geometry.traced_path(lower, dissipating_time=3.0).no_fill().stroke(accent, 0.035),
]
upper.add_updater_fn(step_double, reset=reset_double, fixed_dt=STEP)
lower.add_updater_fn(lower_bob)
scene.play([part.animate.fade_in() for part in parts] + [upper.animate.fade_in(), lower.animate.fade_in()], duration=0.3)
scene.wait(6.0)
upper.remove_updater()
lower.remove_updater()
scene.wait(HOLD)

# A mass on a spring: the spring force is a Parameter stepped from the same
# state, drawn as a force, and a pen plots the displacement over time.
scene.segment("mass_on_spring")
stage = page([
    "block.add_updater_fn(step,",
    "    reset=reset, fixed_dt=1 / 240)",
    "fx = scene.viz.parameter(0)",
    "fx.add_updater_fn(spring_force)",
    "force_from_components(top, fx, 0)",
    "pen.add_updater_fn(plot)",
    "traced_path(pen)",
])
STIFFNESS, FRICTION, PULL = 6.0, 0.35, 1.3
wall = (stage.left + 0.3, stage.center[1] + 1.6)
rest = wall[0] + 3.4
oscillator = {}


def reset_oscillator():
    """The block pulled out and let go."""
    oscillator.update(x=PULL, v=0.0)


def step_oscillator(position, dt, _elapsed):
    """One semi-implicit Euler step of a damped mass on a spring."""
    oscillator["v"] += dt * (-STIFFNESS * oscillator["x"] - FRICTION * oscillator["v"])
    oscillator["x"] += dt * oscillator["v"]
    return (rest + oscillator["x"], position[1], position[2])


def spring_force(_value, _dt, _elapsed):
    """Hooke's law on the displacement the block has just reached."""
    return -STIFFNESS * oscillator["x"]


origin = (stage.left + 0.6, stage.bottom + 1.4)


def plot(_position, _dt, elapsed):
    """Time to the right, displacement up."""
    return (origin[0] + 1.1 * elapsed, origin[1] + 0.9 * oscillator["x"], 0.0)


reset_oscillator()
block = scene.geometry.rounded_rect(1.0, 1.0, 0.08).fill(panel).stroke(fg, 0.05).move_to(rest + PULL, wall[1])
fx = scene.viz.parameter(-STIFFNESS * PULL)
# The force is drawn above the block, from a point that moves with it.
top = block.anchor_point(Anchor.TOP, offset=(0, 0.9))
pen = scene.geometry.dot(0.08).fill(accent).move_to(origin[0], origin[1] + 0.9 * PULL)
parts = [
    M.fixed_support(wall, direction=Direction.RIGHT, size=0.9, ground_length=1.4),
    M.fixed_support((rest, wall[1] - 0.5), size=0.8, ground_length=5.2),
    M.spring_between(wall, block.anchor_point(Anchor.LEFT), coils=9, amplitude=0.25),
    M.force_from_components(top, fx, 0, visual_scale=0.2, label="$F$", show_value=True, label_gap=0.45, color=accent),
    scene.geometry.line(origin[0], origin[1], stage.right, origin[1]).stroke(muted, 0.02),
    scene.geometry.line(origin[0], origin[1] - 1.3, origin[0], origin[1] + 1.3).stroke(muted, 0.02),
    scene.geometry.traced_path(pen).no_fill().stroke(accent, 0.04),
]
block.add_updater_fn(step_oscillator, reset=reset_oscillator, fixed_dt=STEP)
fx.add_updater_fn(spring_force)
pen.add_updater_fn(plot)
scene.play([part.animate.fade_in() for part in parts] + [block.animate.fade_in(), pen.animate.fade_in()], duration=0.3)
scene.wait(5.6)
for driven in (block, fx, pen):
    driven.remove_updater()
scene.wait(HOLD)

# Three moments of each simulation, the last one near its end.
pendulum_end = 0.3 + 5.0 + HOLD
double_end = pendulum_end + 0.3 + 6.0 + HOLD
if snapshots := os.environ.get("GAANIM_SNAPSHOTS"):
    scene.snapshots(
        snapshots,
        [0.3 + t for t in (0.7, 2.2, 4.8)]
        + [pendulum_end + 0.3 + t for t in (1.0, 3.0, 5.8)]
        + [double_end + 0.3 + t for t in (0.6, 2.0, 5.4)],
    )
else:
    scene.render()