alife_dla.kt

@file:afield.runtime.completion.CompletionCategory("ALife")
@file:Suppress("MemberVisibilityCanBePrivate")

package afield.simulation.agents

import afield.data.collections.invoke
import afield.data.collections.r
import afield.data.collections.s
import afield.data.collections.times
import afield.data.graph.main
import afield.document.context.why
import afield.document.graph.minusAssign
import afield.document.graph.plusAssign
import afield.document.properties.compareTo
import afield.document.workspace.file
import afield.document.workspace.path
import afield.geometry.mesh.bounds
import afield.geometry.mesh.rounded
import afield.geometry.spatial.center
import afield.graphics.render.rangeTo
import afield.graphics.shader.compute
import afield.graphics.text.layout
import afield.io.osc.main
import afield.math.autodiff.clamp
import afield.math.autodiff.compareTo
import afield.math.autodiff.div
import afield.math.autodiff.mean
import afield.math.autodiff.min
import afield.math.autodiff.minus
import afield.math.autodiff.plus
import afield.math.autodiff.round
import afield.math.autodiff.times
import afield.math.control.main
import afield.math.linalg.invoke
import afield.math.linalg.remAssign
import afield.math.linalg.set
import afield.math.linalg.times
import afield.math.linalg.x
import afield.math.linalg.xy
import afield.math.statistics.main
import afield.math.units.div
import afield.math.units.minus
import afield.math.units.plus
import afield.math.units.times
import afield.media.image.channel
import afield.media.image.threshold
import afield.media.video.image
import afield.runtime.execution.start
import afield.runtime.lifecycle.once
import afield.simulation.automata.main
import afield.simulation.automata.step

import afield.audio.synth.*
import afield.data.collections.get
import afield.data.graph.fragment
import afield.document.properties.clear
import afield.graphics.render.pointSize
import afield.graphics.shader.attachShader
import afield.graphics.shader.fragment
import afield.graphics.shader.vertex
import afield.lang.stacksl.evolution.render
import afield.math.linalg.Vec4
import afield.math.linalg.vec
import afield.math.statistics.random
import afield.media.image.resize
import afield.document.graph.Box
import afield.app.protocol.Rect
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.Target
import afield.graphics.vulkan.buffers.FrameBufferGroup
import afield.graphics.vulkan.buffers.defaultBuffer
import afield.graphics.vulkan.device.CommandBuffer
import afield.graphics.render.Mesh
import afield.graphics.render.Module
import afield.graphics.render.ModulePointer
import afield.graphics.render.Parameters
import afield.graphics.render.before
import afield.document.context.fromContext
import afield.graphics.render.parameters
import afield.graphics.render.pointer
import afield.graphics.vulkan.targets.PingPongRenderToTexture
import afield.graphics.vulkan.targets.RenderTarget
import afield.graphics.texture.HasImageView
import afield.graphics.texture.HasImageViewArray
import afield.graphics.texture.ImageView
import afield.graphics.texture.TextureSampler
import afield.math.linalg.Vector2f
import org.lwjgl.system.MemoryStack
import org.lwjgl.vulkan.VK13
import org.lwjgl.vulkan.VK10.VK_ATTACHMENT_STORE_OP_STORE
import org.lwjgl.vulkan.VK10.VK_BLEND_FACTOR_ONE
import org.lwjgl.vulkan.VK10.VK_BLEND_OP_ADD
import org.lwjgl.vulkan.VK10.VK_FORMAT_R16G16B16A16_SFLOAT
import org.lwjgl.vulkan.VK13.VK_ACCESS_2_SHADER_READ_BIT
import org.lwjgl.vulkan.VK13.VK_ACCESS_2_SHADER_WRITE_BIT
import org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT
import org.lwjgl.vulkan.VK13.VK_QUEUE_FAMILY_IGNORED
import org.lwjgl.vulkan.VK13.VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2
import org.lwjgl.vulkan.VK13.VK_STRUCTURE_TYPE_DEPENDENCY_INFO
import org.lwjgl.vulkan.VkBufferMemoryBarrier2
import org.lwjgl.vulkan.VkDependencyInfo
import java.nio.ByteOrder
import kotlin.math.max
import kotlin.random.Random

data class ShaderDlaSettings(
    /**
     * Prefix for everything this system registers by name: its two modules, their meshes, the
     * particle buffer, the barrier, and the entries it makes in a render target's pre-target and
     * module maps. Two systems in one renderer need two names.
     */
    var name: String = "dla",
    /**
     * Width of the simulation texture in pixels, which is also the coordinate system the walkers step
     * in: positions are pixels, not normalised.
     */
    var width: Int = 1024,
    /** Height of the simulation texture in pixels. Must be positive, as must [width]. */
    var height: Int = 1024,
    /**
     * How many walkers are in flight at once, one 16-byte record each. Because a stuck walker is
     * respawned immediately this is a growth rate rather than a bound on the cluster, and it trades
     * off against [width] and [height]: too few walkers for the area and the cluster takes thousands
     * of frames to become visible, too many and the whole field sticks in the first few.
     */
    var particleCount: Int = 65_536,
    /** `gl_PointSize` for the drawn walkers and splats. 1 draws a single pixel each. */
    var pointSize: Float = 1f,
    /** How far a walker moves each frame, along one of the four axis directions. */
    var walkStepPx: Float = 1f,
    /**
     * How far from a walker the four neighbour probes are taken, rounded and floored to one pixel.
     * Raising it lets walkers stick before they touch, which thickens the branches.
     */
    var sampleStepPx: Float = 1f,
    /**
     * The red value at which a pixel counts as solid. Both the neighbour probe and the age channel
     * ask this question, so it is one threshold for sticking and for colouring.
     */
    var solidThreshold: Float = 0.5f,
    /**
     * The green a walker adds to the pixel it is standing on. Nothing reads green back, and the
     * preserve pass does not carry it forward, so this is a live view of the diffusing cloud.
     */
    var walkerGreen: Float = 0.12f,
    /**
     * The red a walker adds at the moment it sticks. This is the only thing that grows the cluster,
     * and it is added to what is already there, so a pixel several walkers have found reads above 1
     * and takes proportionally longer to decay away.
     */
    var stuckRed: Float = 1f,
    /** The blue added alongside [stuckRed], so a new branch tip starts with some age rather than none. */
    var stuckBlue: Float = 0.35f,
    /**
     * Centre of the disc the preserve pass holds solid, in pixels; `null` puts it at the middle of the
     * texture. Something has to be solid before anything can stick to it, so this is where the cluster
     * starts. Resolved once in the constructor, unlike the rest of these.
     */
    var seedPositionPx: Vector2f? = null,
    /** Radius of that disc in pixels. */
    var seedRadiusPx: Float = 1.5f,
    /**
     * How much red every pixel outside the seed disc loses per frame. At 0 the cluster is permanent;
     * at the default of 0.001 a branch fades once walkers stop reinforcing it, which is what keeps a
     * long-running field from filling in.
     */
    var solidDecayPerFrame: Float = 0.001f,
    /** How much blue a pixel below [solidThreshold] loses per frame. */
    var ageDecayPerFrame: Float = 0.004f,
    /** How much blue a solid pixel gains per frame, up to 1. At the default that is 100 frames to full age. */
    var ageGrowPerFrame: Float = 0.01f,
    /** The blue the seed disc is held at, so the oldest part of the cluster reads as old from frame one. */
    var seedBlue: Float = 0.15f,
    /**
     * Colour format of the ping-pong target. A float format is doing work here: `solidDecayPerFrame`
     * at its default is smaller than one step of an 8-bit channel, so an 8-bit target would round each
     * frame's decay away and nothing would ever fade.
     */
    var format: Int = VK_FORMAT_R16G16B16A16_SFLOAT,
    /**
     * Clear value for the target's colour attachment. The preserve pass covers every pixel, so this is
     * what the first frame starts from rather than a background that stays visible behind the cluster.
     */
    var clearColor: Vec4 = vec(0, 0, 0, 1),
    /** Seed for the initial scatter of walkers, handed to [ShaderDlaSystem.resetParticles] at construction. */
    var randomSeed: Long = 1L
)

    init {
        require(settings.width > 0) { "width must be positive" }
        require(settings.height > 0) { "height must be positive" }
        require(settings.particleCount > 0) { "particleCount must be positive" }

        val seedPosition = settings.seedPositionPx ?: Vector2f(settings.width * 0.5f, settings.height * 0.5f)

        preserveModule = Module(
            parameters(renderer) {
                +this@ShaderDlaSystem.simulation
                name = "${settings.name}-preserve"
                depthTest = false
                blend = false
                vertex = FULLSCREEN_TRIANGLE_VERTEX
                fragment = PRESERVE_FRAGMENT
                textureResolver = sourceTextureResolver()
            }
        ).pointer()
        Mesh(renderer, preserveModule).apply {
            resize(3, 0, preserveModule)
            push.bind("seedPositionPx") { seedPosition }
            push.bind("seedRadiusPx") { settings.seedRadiusPx }
            push.bind("solidThreshold") { settings.solidThreshold }
            push.bind("solidDecayPerFrame") { settings.solidDecayPerFrame }
            push.bind("ageDecayPerFrame") { settings.ageDecayPerFrame }
            push.bind("ageGrowPerFrame") { settings.ageGrowPerFrame }
            push.bind("seedBlue") { settings.seedBlue }
            preserveModule.current.meshes["${settings.name}-preserve-quad"] = this
        }

        particleModule = Module(
            parameters(renderer) {
                +this@ShaderDlaSystem.simulation
                name = "${settings.name}-particles"
                depthTest = false
                blend = true
                srcColor = VK_BLEND_FACTOR_ONE
                dstColor = VK_BLEND_FACTOR_ONE
                colorBlendOp = VK_BLEND_OP_ADD
                srcAlpha = VK_BLEND_FACTOR_ONE
                dstAlpha = VK_BLEND_FACTOR_ONE
                alphaBlendOp = VK_BLEND_OP_ADD
                primitiveType = points
                vertex = PARTICLE_VERTEX
                fragment = PARTICLE_FRAGMENT
                textureResolver = sourceTextureResolver()
            }
        ).pointer()
        particleModule.current.before["${settings.name}-particle-ssbo-barrier"] = before { command, _ ->
            graphicsToGraphicsBarrier(command, particleBuffer.get())
        }.target

        particleMesh = Mesh(renderer, particleModule).apply {
            resize(settings.particleCount, 0, particleModule)
            push.bind("particleBuffer") { particleBuffer }
            push.bind("pointSize") { settings.pointSize }
            push.bind("walkStepPx") { settings.walkStepPx }
            push.bind("sampleStepPx") { settings.sampleStepPx }
            push.bind("solidThreshold") { settings.solidThreshold }
            push.bind("walkerGreen") { settings.walkerGreen }
            push.bind("stuckRed") { settings.stuckRed }
            push.bind("stuckBlue") { settings.stuckBlue }
        }
        particleModule.current.meshes["${settings.name}-points"] = particleMesh

        simulation += "preserve" to preserveModule
        simulation += "particles" to particleModule

        resetParticles(settings.randomSeed)
        Box.context.get()?.attachShader(particleModule)
    }

    companion object {
        private const val PARTICLE_STRIDE_BYTES = 16

        private val FULLSCREEN_TRIANGLE_VERTEX = """
            #version 450

            vec2 fullscreenVertex(uint index) {
                if (index == 0u) return vec2(-1.0, -1.0);
                if (index == 1u) return vec2(3.0, -1.0);
                return vec2(-1.0, 3.0);
            }

            void main() {
                gl_Position = vec4(fullscreenVertex(uint(gl_VertexIndex)), 0.0, 1.0);
            }
        """.trimIndent()

        private val PRESERVE_FRAGMENT = """
            #version 450

            layout(set = 0, binding = 0) uniform sampler2D sourceTex;

            layout(push_constant) uniform Push {
                vec2 seedPositionPx;
                float seedRadiusPx;
                float solidThreshold;
                float solidDecayPerFrame;
                float ageDecayPerFrame;
                float ageGrowPerFrame;
                float seedBlue;
            } push;

            layout(location = 0) out vec4 outColor;

            void main() {
                ivec2 dims = textureSize(sourceTex, 0);
                ivec2 pixel = clamp(ivec2(gl_FragCoord.xy), ivec2(0), dims - 1);
                vec4 previous = texelFetch(sourceTex, pixel, 0);
                bool inSeed = distance(gl_FragCoord.xy, push.seedPositionPx) <= push.seedRadiusPx;

                float decayedSolid = max(previous.r - push.solidDecayPerFrame, 0.0);
                float solid = max(decayedSolid, inSeed ? 1.0 : 0.0);

                float decayedAge = max(previous.b - push.ageDecayPerFrame, 0.0);
                float age = solid >= push.solidThreshold
                    ? min(decayedAge + push.ageGrowPerFrame, 1.0)
                    : decayedAge;
                if (inSeed) age = max(age, push.seedBlue);

                outColor = vec4(solid, 0.0, age, 1.0);
            }
        """.trimIndent()

        private val PARTICLE_VERTEX = """
            #version 450
            #extension GL_EXT_buffer_reference : require
            #extension GL_EXT_scalar_block_layout : require

            struct Particle {
                vec2 positionPx;
                uint rngState;
                uint mode;
            };

            layout(std430, buffer_reference, scalar) buffer ParticleBuffer {
                Particle particles[];
            };

            layout(set = 0, binding = 0) uniform sampler2D sourceTex;

            layout(push_constant) uniform Push {
                ParticleBuffer particleBuffer;
                float pointSize;
                float walkStepPx;
                float sampleStepPx;
                float solidThreshold;
                float walkerGreen;
                float stuckRed;
                float stuckBlue;
            } push;

            layout(location = 0) out vec4 vColor;

            uint nextRandom(inout uint state) {
                state = state * 747796405u + 2891336453u;
                uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
                return (word >> 22u) ^ word;
            }

            float random01(inout uint state) {
                return float(nextRandom(state)) / 4294967295.0;
            }

            vec2 randomPosition(inout uint state, ivec2 dims) {
                return vec2(
                    random01(state) * float(max(dims.x - 1, 0)),
                    random01(state) * float(max(dims.y - 1, 0))
                );
            }

            vec2 randomWalkStep(inout uint state, float stepPx) {
                uint choice = nextRandom(state) & 3u;
                if (choice == 0u) return vec2(stepPx, 0.0);
                if (choice == 1u) return vec2(-stepPx, 0.0);
                if (choice == 2u) return vec2(0.0, stepPx);
                return vec2(0.0, -stepPx);
            }

            bool isSolidAt(ivec2 pixel, ivec2 dims, float threshold) {
                ivec2 clampedPixel = clamp(pixel, ivec2(0), dims - 1);
                return texelFetch(sourceTex, clampedPixel, 0).r >= threshold;
            }

            void main() {
                ivec2 dims = textureSize(sourceTex, 0);
                vec2 maxPos = vec2(max(dims.x - 1, 0), max(dims.y - 1, 0));

                Particle particle = push.particleBuffer.particles[gl_VertexIndex];
                uint rng = particle.rngState ^ (uint(gl_VertexIndex) * 2246822519u + 3266489917u);
                vec2 positionPx = clamp(particle.positionPx, vec2(0.0), maxPos);
                uint mode = particle.mode;
                bool becameSolid = false;

                if (mode != 0u) {
                    positionPx = randomPosition(rng, dims);
                    mode = 0u;
                } else {
                    ivec2 center = clamp(ivec2(round(positionPx)), ivec2(0), dims - 1);
                    int probe = max(int(round(push.sampleStepPx)), 1);
                    bool solidNeighbor =
                        isSolidAt(center + ivec2(probe, 0), dims, push.solidThreshold) ||
                        isSolidAt(center + ivec2(-probe, 0), dims, push.solidThreshold) ||
                        isSolidAt(center + ivec2(0, probe), dims, push.solidThreshold) ||
                        isSolidAt(center + ivec2(0, -probe), dims, push.solidThreshold);

                    if (solidNeighbor) {
                        mode = 1u;
                        becameSolid = true;
                    } else {
                        positionPx = clamp(positionPx + randomWalkStep(rng, push.walkStepPx), vec2(0.0), maxPos);
                    }
                }

                push.particleBuffer.particles[gl_VertexIndex].positionPx = positionPx;
                push.particleBuffer.particles[gl_VertexIndex].rngState = max(rng, 1u);
                push.particleBuffer.particles[gl_VertexIndex].mode = mode;

                vec2 uv = (positionPx + vec2(0.5)) / vec2(max(dims.x, 1), max(dims.y, 1));
                vec2 clip = uv * 2.0 - 1.0;
                gl_Position = vec4(clip.x, -clip.y, 0.0, 1.0);
                gl_PointSize = push.pointSize;
                vColor = becameSolid
                    ? vec4(push.stuckRed, 0.0, push.stuckBlue, 0.0)
                    : vec4(0.0, push.walkerGreen, 0.0, 0.0);
            }
        """.trimIndent()

        private val PARTICLE_FRAGMENT = """
            #version 450

            layout(location = 0) in vec4 vColor;
            layout(location = 0) out vec4 outColor;

            void main() {
                outColor = vColor;
            }
        """.trimIndent()

        private fun graphicsToGraphicsBarrier(command: CommandBuffer, buffer: afield.graphics.vulkan.buffers.VulkanBuffer) {
            MemoryStack.stackPush().use { stack ->
                val barrier = VkBufferMemoryBarrier2.calloc(1, stack)
                barrier[0]
                    .sType(VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2)
                    .srcStageMask(VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT)
                    .srcAccessMask(VK_ACCESS_2_SHADER_WRITE_BIT)
                    .dstStageMask(VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT)
                    .dstAccessMask(VK_ACCESS_2_SHADER_READ_BIT or VK_ACCESS_2_SHADER_WRITE_BIT)
                    .srcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
                    .dstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
                    .buffer(buffer.buffer)
                    .offset(0L)
                    .size(buffer.requestedSize)

                val dependency = VkDependencyInfo.calloc(stack)
                    .sType(VK_STRUCTURE_TYPE_DEPENDENCY_INFO)
                    .pBufferMemoryBarriers(barrier)

                VK13.vkCmdPipelineBarrier2(command.vkCommandBuffer, dependency)
            }
        }
    }
}