'''
Example of using shader storage buffer in compute shader.
We read from a buffer and write the result to another buffer.
Every frame we swap the buffers around transforming positions
of balls.

Buffer.bind_to_storage_buffer is used to bind a buffer as storage buffer
to a specific binding point specified in the compute program.

In addition we render the balls using a geometry shader to easily
batch draw them all in one render call.

author: minu jeong
modified by: einarf
'''
import math
import random
import numpy as np
from _example import Example


items_vertex_shader_code = """
    #version 430

    in vec4 in_vert;
    in vec4 in_col;

    out vec4 v_color;

    void main()
    {
        gl_Position = in_vert; // x, y, 0, radius
        v_color = in_col;
    }
    """

# Geometry shader turning the points into triangle strips.
# This can also be done with point sprites.
items_geo_shader = """
    #version 330

    layout(points) in;
    layout(triangle_strip, max_vertices=4) out;

    in vec4 v_color[];
    out vec2 uv;
    out vec4 color;

    void main() {
        float radius = gl_in[0].gl_Position.w;
        vec2 pos = gl_in[0].gl_Position.xy;

        // Emit the triangle strip creating a "quad"
        // Lower left
        gl_Position = vec4(pos + vec2(-radius, -radius), 0, 1);
        color = v_color[0];
        uv = vec2(0, 0);
        EmitVertex();

        // upper left
        gl_Position = vec4(pos + vec2(-radius, radius), 0, 1);
        color = v_color[0];
        uv = vec2(0, 1);
        EmitVertex();

        // lower right
        gl_Position = vec4(pos + vec2(radius, -radius), 0, 1);
        color = v_color[0];
        uv = vec2(1, 0);
        EmitVertex();

        // upper right
        gl_Position = vec4(pos + vec2(radius, radius), 0, 1);
        color = v_color[0];
        uv = vec2(1, 1);
        EmitVertex();

        EndPrimitive();
    }
"""

items_fragment_shader_code = """
    #version 430

    in vec2 uv;
    in vec4 color;
    out vec4 out_color;
    void main()
    {
        // Calculate the length from the center of the "quad"
        // using texture coordinates discarding fragments
        // further away than 0.5 creating a circle.
        if (length(vec2(0.5, 0.5) - uv.xy) > 0.5)
        {
            discard;
        }
        out_color = color;
    }
"""

# calc position with compute shader
compute_worker_shader_code = """
#version 430
#define GROUP_SIZE %COMPUTE_SIZE%

layout(local_size_x=GROUP_SIZE) in;

// All values are vec4s because of block alignment rules (keep it simple).
// We could also declare all values as floats to make it tightly packed.
// See : https://www.khronos.org/opengl/wiki/Interface_Block_(GLSL)#Memory_layout
struct Ball
{
    vec4 pos; // x, y, 0, radius
    vec4 vel; // x, y (velocity)
    vec4 col; // r, g, b (color)
};

layout(std430, binding=0) buffer balls_in
{
    Ball balls[];
} In;
layout(std430, binding=1) buffer balls_out
{
    Ball balls[];
} Out;

void main()
{
    int x = int(gl_GlobalInvocationID);

    Ball in_ball = In.balls[x];

    vec4 p = in_ball.pos.xyzw;
    vec4 v = in_ball.vel.xyzw;

    p.xy += v.xy;

    float rad = p.w * 0.5;
    if (p.x - rad <= -1.0)
    {
        p.x = -1.0 + rad;
        v.x *= -0.98;
    }
    else if (p.x + rad >= 1.0)
    {
        p.x = 1.0 - rad;
        v.x *= -0.98;
    }

    if (p.y - rad <= -1.0)
    {
        p.y = -1.0 + rad;
        v.y *= -0.98;
    }
    else if (p.y + rad >= 1.0)
    {
        p.y = 1.0 - rad;
        v.y *= -0.98;
    }
    v.y += -0.001;

    Ball out_ball;
    out_ball.pos.xyzw = p.xyzw;
    out_ball.vel.xyzw = v.xyzw;

    vec4 c = in_ball.col.xyzw;
    out_ball.col.xyzw = c.xyzw;

    Out.balls[x] = out_ball;
}
"""

class ComputeShaderSSBO(Example):
    title = "Compute Shader SSBO"
    gl_version = 4, 3  # Required opengl version
    window_size = 600, 600  # Initial window size
    aspect_ratio = 1.0  # Force viewport aspect ratio (regardless of window size)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.COUNT = 256  # number of balls
        self.STRUCT_SIZE = 12  # number of floats per item/ball

        # Program for drawing the balls / items
        self.program = self.ctx.program(
            vertex_shader=items_vertex_shader_code,
            geometry_shader=items_geo_shader,
            fragment_shader=items_fragment_shader_code
        )

        # Load compute shader
        compute_shader_code_parsed = compute_worker_shader_code.replace("%COMPUTE_SIZE%", str(self.COUNT))
        self.compute_shader = self.ctx.compute_shader(compute_shader_code_parsed)

        # Create the two buffers the compute shader will write and read from
        compute_data = np.fromiter(self.gen_initial_data(), dtype="f4")
        self.compute_buffer_a = self.ctx.buffer(compute_data)
        self.compute_buffer_b = self.ctx.buffer(compute_data)

        # Prepare vertex arrays to drawing balls using the compute shader buffers are input
        # We use 4x4 (padding format) to skip the velocity data (not needed for drawing the balls)
        self.balls_a = self.ctx.vertex_array(
            self.program, [self.compute_buffer_a.bind('in_vert', 'in_col', layout='4f 4x4 4f')],
        )
        self.balls_b = self.ctx.vertex_array(
            self.program, [self.compute_buffer_b.bind('in_vert', 'in_col', layout='4f 4x4 4f')],
        )

    def gen_initial_data(self):
        """Generator function creating the initial buffer data"""
        for i in range(self.COUNT):
            _angle = (i / self.COUNT) * math.pi * 2.0
            _dist = 0.125
            radius = random.random() * 0.01 + 0.01
            # position and radius (vec4)
            yield math.cos(_angle) * _dist
            yield math.sin(_angle) * _dist
            yield 0.0
            yield radius
            # velocity (vec4)
            _v = random.random() * 0.005 + 0.01
            yield math.cos(_angle) * _v
            yield math.sin(_angle) * _v
            yield 0.0
            yield 0.0
            # color (vec4)
            yield 1.0 * random.random()
            yield 1.0 * random.random()
            yield 1.0 * random.random()
            yield 1.0


    def render(self, time, frame_time):
        # Calculate the next position of the balls with compute shader
        self.compute_buffer_a.bind_to_storage_buffer(0)
        self.compute_buffer_b.bind_to_storage_buffer(1)
        self.compute_shader.run(group_x=self.STRUCT_SIZE)

        # Batch draw the balls
        self.balls_b.render(mode=self.ctx.POINTS)

        # Swap the buffers and vertex arrays around for next frame
        self.compute_buffer_a, self.compute_buffer_b = self.compute_buffer_b, self.compute_buffer_a
        self.balls_a, self.balls_b = self.balls_b, self.balls_a


if __name__ == "__main__":
    ComputeShaderSSBO.run()