pybevy 0.2.1

"""
Cube Wave Animation with ViewColumn Zero-Copy API + Parallel Execution

Demonstrates ViewColumn + Numba JIT with PARALLEL execution for maximum performance.
Updates 65,536+ cubes with sine wave patterns using multi-core CPU parallelization.

Performance (256x256 = 65,536 cubes):
- ViewColumn + Numba (parallel): ~0.05-0.2ms per frame
- ViewColumn + Numba (single-thread): ~0.1-0.5ms per frame
- Traditional Query: ~50-100ms per frame
- Speedup: 500-1000x faster than Query iteration

Parallelization:
- Uses numba.prange() for multi-threaded execution
- Distributes work across all CPU cores
- Zero-copy pointer access to ECS storage
- Additional 4-8x speedup from parallelization

Controls:
- Mouse drag to rotate camera
- Mouse wheel to zoom
- Press Ctrl+C to exit
"""

import math
from typing import TYPE_CHECKING

try:
    import numba  # type: ignore[import-untyped]
except ImportError:
    print("ERROR: Numba is required for this example.")
    print("Install with: poetry install --extras numba")
    print("Or: pip install numba")
    exit(1)

from pybevy.contrib import OrbitCamera, OrbitCameraPlugin
from pybevy.prelude import *

if TYPE_CHECKING:
    from pybevy.ecs import FieldExpr  # type: ignore[assignment]


# Configuration
GRID_SIZE = 387  # 256x256 = 65,536 cubes
CUBE_SPACING = 2.0
WAVE_SPEED = 2.0
WAVE_AMPLITUDE = 3.0



@numba.jit(nopython=True, parallel=True)
def animate_wave_kernel(
    pos_x: "FieldExpr",  # ViewColumn for Transform.translation.x
    pos_y: "FieldExpr",  # ViewColumn for Transform.translation.y
    pos_z: "FieldExpr",  # ViewColumn for Transform.translation.z
    time: float,
    wave_speed: float,
    amplitude: float,
) -> None:
    """Animate cubes in a sine wave pattern (PARALLEL).

    This function receives ViewColumn handles and accesses them directly
    with col[i] syntax. Numba compiles this to native pointer arithmetic
    and parallelizes across CPU cores using numba.prange.

    Args:
        pos_x, pos_y, pos_z: ViewColumn handles for position components
        time: Current elapsed time
        wave_speed: Speed of wave propagation
        amplitude: Wave height
    """
    n = len(pos_x)

    for i in numba.prange(n):  # Parallel loop across CPU cores
        # Read X and Z to calculate wave
        x = pos_x[i]
        z = pos_z[i]

        # Calculate distance from origin
        dist = math.sqrt(x * x + z * z)

        # Sine wave based on distance and time
        wave = math.sin(dist * 0.5 - time * wave_speed) * amplitude

        # Update Y position (direct write to ECS storage!)
        pos_y[i] = wave


@numba.jit(nopython=True, parallel=True)
def animate_spiral_kernel(
    pos_x: "FieldExpr",
    pos_y: "FieldExpr",
    pos_z: "FieldExpr",
    rot_y: "FieldExpr",  # ViewColumn for Transform.rotation.y (quat component)
    rot_w: "FieldExpr",  # ViewColumn for Transform.rotation.w (quat component)
    time: float,
) -> None:
    """Animate cubes in a spiral pattern with rotation (PARALLEL).

    Demonstrates updating both position and rotation via ViewColumn.
    Uses parallel execution across CPU cores.
    """
    n = len(pos_x)

    for i in numba.prange(n):  # Parallel loop
        x = pos_x[i]
        z = pos_z[i]

        # Spiral wave
        dist = math.sqrt(x * x + z * z)
        angle = math.atan2(z, x)
        wave = math.sin(dist * 0.3 - time * 2.0 + angle) * 2.0

        pos_y[i] = wave

        # Rotate cubes based on height
        rotation_angle = wave * 0.5
        half_angle = rotation_angle * 0.5
        rot_y[i] = math.sin(half_angle)
        rot_w[i] = math.cos(half_angle)


@numba.jit(nopython=True, parallel=True)
def animate_ripple_kernel(
    pos_x: "FieldExpr",
    pos_y: "FieldExpr",
    pos_z: "FieldExpr",
    time: float,
) -> None:
    """Multiple overlapping ripples emanating from different points (PARALLEL)."""
    n = len(pos_x)

    # Ripple sources
    sources = [
        (0.0, 0.0),
        (50.0, 50.0),
        (-50.0, 50.0),
        (50.0, -50.0),
        (-50.0, -50.0),
    ]

    for i in numba.prange(n):  # Parallel loop
        x = pos_x[i]
        z = pos_z[i]

        # Sum ripples from all sources
        height = 0.0
        for sx, sz in sources:
            dx = x - sx
            dz = z - sz
            dist = math.sqrt(dx * dx + dz * dz)
            ripple = math.sin(dist * 0.3 - time * 3.0) * 2.0 / (1.0 + dist * 0.05)
            height += ripple

        pos_y[i] = height



@component
class Cube(Component):
    """Marker component for animated cubes."""



def setup_scene(
    commands: Commands,
    meshes: ResMut[Assets[Mesh]],
    materials: ResMut[Assets[StandardMaterial]],
) -> None:
    """Set up 10,000 cubes in a grid."""

    # Materials and meshes
    cube_material = materials.add(Color.srgb(0.3, 0.7, 0.9))
    cube_mesh = meshes.add(Cuboid(1.0, 1.0, 1.0))

    for row in range(GRID_SIZE):
        for col in range(GRID_SIZE):
            x = (col - GRID_SIZE / 2) * CUBE_SPACING
            z = (row - GRID_SIZE / 2) * CUBE_SPACING

            commands.spawn(
                Cube(),
                Mesh3d(cube_mesh),
                MeshMaterial3d(cube_material),
                Transform.from_xyz(x, 0.0, z),
            )

    # Lighting
    commands.spawn(
        DirectionalLight(
            illuminance=10000.0,
            color=Color.WHITE,
            shadows_enabled=False,  # Disable shadows for performance
        ),
        Transform.IDENTITY.looking_at(Vec3(-1.0, -2.5, -1.0), Vec3.Y),
    )

    commands.insert_resource(GlobalAmbientLight(brightness=300.0, color=Color.WHITE))

    # Camera
    camera_distance = GRID_SIZE * CUBE_SPACING * 0.8
    camera_height = camera_distance * 0.6

    target = Vec3(0.0, 0.0, 0.0)
    initial_pitch = math.atan2(camera_height, camera_distance)

    commands.spawn(
        Camera3d(),
        Transform.from_xyz(0.0, camera_height, camera_distance).looking_at(
            target, Vec3.Y
        ),
        OrbitCamera(
            distance=math.sqrt(camera_height**2 + camera_distance**2),
            yaw=0.0,
            pitch=initial_pitch,
            target=target,
        ),
    )



def cube_animation_system(
    view: View[Mut[Transform], With[Cube]],
    time: Res[Time],
) -> None:
    """Animate all cubes using ViewColumn zero-copy API.

    This system processes 10,000 transforms in ~0.1ms using Numba JIT,
    compared to ~10ms with traditional Query iteration (100x speedup).
    """
    t = time.elapsed_secs()

    # Choose animation pattern based on time
    pattern = int(t / 10.0) % 3

    for batch in view.iter_batches():
        transform = batch.column_mut(Transform)

        if pattern == 0:
            animate_wave_kernel(
                transform.translation.x,
                transform.translation.y,
                transform.translation.z,
                t,
                WAVE_SPEED,
                WAVE_AMPLITUDE,
            )
        elif pattern == 1:
            animate_spiral_kernel(
                transform.translation.x,
                transform.translation.y,
                transform.translation.z,
                transform.rotation.y,
                transform.rotation.w,
                t,
            )
        else:
            animate_ripple_kernel(
                transform.translation.x,
                transform.translation.y,
                transform.translation.z,
                t,
            )


class FPSCounter:
    def __init__(self) -> None:
        self.frame_count: int = 0


def fps_system(time: Res[Time], counter: Local[FPSCounter]) -> None:
    counter.frame_count += 1
    t = time.elapsed_secs()
    if t > 0 and counter.frame_count % 60 == 0:
        print(f"FPS: ~{counter.frame_count / max(t, 0.001):.1f}")


@entrypoint
def main(app: App) -> App:
    return (
        app.add_plugins(DefaultPlugins)
        .add_plugins(OrbitCameraPlugin())  # type: ignore[arg-type]
        .add_systems(Startup, setup_scene)
        .add_systems(Update, (cube_animation_system, fps_system))
    )


if __name__ == "__main__":
    print(f"View+Numba parallel cube wave: {GRID_SIZE * GRID_SIZE:,} cubes. Drag mouse to rotate camera.")
    main().run()