Source code for simvx.core.coroutines

"""Coroutine helpers: parallel, wait, wait_until, wait_signal, next_frame.

Coroutines are driven by the scene tree with ``gen.send(dt)`` after the
first prime, so helpers that track time can read the per-tick delta via
``dt = yield`` and accumulate it. Helpers below follow that convention,
which means they correctly respect ``SceneTree.paused`` and time scale.
"""

from collections.abc import Callable
from typing import TYPE_CHECKING

from .descriptors import Coroutine

if TYPE_CHECKING:
    from .signals import Signal




def parallel(*coroutines: Coroutine) -> Coroutine:
    """Run multiple coroutines simultaneously, finish when all complete.

    Forwards the per-tick ``dt`` received via ``yield`` to each child
    coroutine, priming each on its first advance.
    """
    primed: dict[int, bool] = {id(co): False for co in coroutines}
    active = list(coroutines)
    dt = 0.0
    while active:
        finished = []
        for co in active:
            try:
                if primed[id(co)]:
                    co.send(dt)
                else:
                    next(co)
                    primed[id(co)] = True
            except StopIteration:
                finished.append(co)
        for co in finished:
            active.remove(co)
        if active:
            dt = yield





def wait(seconds: float) -> Coroutine:
    """Pause a coroutine for ``seconds`` of scene-tree time."""
    elapsed = 0.0
    while elapsed < seconds:
        dt = yield
        elapsed += dt or 0.0





def wait_until(condition: Callable[[], bool]) -> Coroutine:
    """Yield until condition() returns True."""
    while not condition():
        yield





def wait_signal(signal: Signal) -> Coroutine:
    """Yield until signal is emitted. Returns signal args."""
    received = [None]
    done = [False]

    def _on_signal(*args):
        received[0] = args
        done[0] = True
    signal.connect(_on_signal)
    try:
        while not done[0]:
            yield
    finally:
        signal.disconnect(_on_signal)
    return received[0]





def next_frame() -> Coroutine:
    """Yield for exactly one frame."""
    yield



# ============================================================================
# Damped-sine impulses ("juice") -- punch_position / punch_rotation
# ============================================================================


def _resolve_amplitude_vec2(amplitude):
    """Coerce a scalar/Vec2/tuple into a Vec2 amplitude."""
    from .math.types import Vec2

    if isinstance(amplitude, int | float):
        return Vec2(float(amplitude), float(amplitude))
    return Vec2(amplitude)




def punch_position(
    node,
    amplitude,
    duration: float,
    *,
    frequency: float = 12.0,
    decay: float = 6.0,
    attr: str = "position",
) -> Coroutine:
    """Damped-sine impulse on ``node.position``: arcade hit / screen-shake.

    Maps to the classic ``A * sin(2pi * f * t) * exp(-d * t)`` formula used by
    Balatro, Claustrowordia and Casual Crusade for card slams and tile hits.
    The node's ``position`` is offset around its current value each tick and
    restored exactly to the starting value on completion.

    Args:
        node: Node to shake. The named attribute must support both reads and
            tuple/Vec2 writes; ``position`` is the common case.
        amplitude: Peak displacement. Pass a scalar for a uniform shake
            (`Vec2(amp, amp)`), or a ``Vec2`` / 2-tuple to bias the axis
            (e.g. ``(0, 8)`` for vertical-only).
        duration: Total duration in seconds. After this many seconds the
            position is snapped back to the starting value.
        frequency: Oscillations per second (default 12 Hz -- arcade-feel).
        decay: Exponential decay rate per second (default 6.0).
        attr: Position attribute name on ``node`` (default ``"position"``).

    Yields:
        Per-tick ``dt`` from the coroutine driver.
    """
    import math as _math

    from .math.types import Vec2

    amp = _resolve_amplitude_vec2(amplitude)
    duration = max(float(duration), 0.0)
    base = Vec2(getattr(node, attr))

    elapsed = 0.0
    try:
        while elapsed < duration:
            envelope = _math.exp(-decay * elapsed)
            wave = _math.sin(2.0 * _math.pi * frequency * elapsed)
            setattr(node, attr, base + amp * (envelope * wave))
            dt = yield
            elapsed += dt or 0.0
    finally:
        setattr(node, attr, base)





def punch_rotation(
    node,
    amplitude: float,
    duration: float,
    *,
    frequency: float = 12.0,
    decay: float = 6.0,
    attr: str = "rotation",
) -> Coroutine:
    """Damped-sine impulse on ``node.rotation`` (radians).

    Twin of :func:`punch_position` for angular hits -- the camera shake
    rotation, the card-flip wobble, the tile-press lean. Pass ``amplitude``
    in **radians**.

    Args:
        node: Node to shake. The attribute must support float read/write.
        amplitude: Peak angular displacement in radians.
        duration: Total duration in seconds.
        frequency: Oscillations per second (default 12 Hz).
        decay: Exponential decay rate per second (default 6.0).
        attr: Rotation attribute name on ``node`` (default ``"rotation"``).

    Yields:
        Per-tick ``dt`` from the coroutine driver.
    """
    import math as _math

    amp = float(amplitude)
    duration = max(float(duration), 0.0)
    base = float(getattr(node, attr))

    elapsed = 0.0
    try:
        while elapsed < duration:
            envelope = _math.exp(-decay * elapsed)
            wave = _math.sin(2.0 * _math.pi * frequency * elapsed)
            setattr(node, attr, base + amp * envelope * wave)
            dt = yield
            elapsed += dt or 0.0
    finally:
        setattr(node, attr, base)