"""
Example using indirect rendering using simple buffer arrays.

Indirect rendering is simply describing one or multiple
rendering calls in a Buffer. This buffer can be written
to from the python side or generated by a compute shader.

The advantage of indirect rendering is that you can pack
all your objects in the same buffer(s) and the indirect
rendering call can draw all of them at once.

The Geometry.geometry_indirect method takes an additional
parameter (a Buffer) containing the info for each rendering call
to issue.

The format of this buffer depends if we are rendering with
or without index buffer.

If Python and Arcade are installed, this example can be run from the command line with:
python -m arcade.examples.gl.render_indirect
"""

from random import (
    randint,
    random,
    uniform,
)
from array import array

import arcade
from arcade.gl import BufferDescription

WINDOW_WIDTH = 800
WINDOW_HEIGHT = 600


class RenderIndirect(arcade.Window):

    def __init__(self):
        super().__init__(WINDOW_WIDTH, WINDOW_HEIGHT, "Render Indirect", gl_version=(4, 3))
        # We don't want vsync when using a query to measure performance
        self.set_vsync(False)
        self.frame = 0

        # Let's generate some simple geometry
        self.num_objects = 10_000
        # Draw parameters. These can be adjusted to render parts of the geometry.
        self.start = 0
        self.count = self.num_objects

        # Query object to measure the rendering calls
        self.query = self.ctx.query()

        # Generate:
        # geometry: Triangles for lots of squares
        # indices: Optional indices for an index buffer
        # array_draw_commands: List of draw commands for rendering each individual
        # square (without index buffer)
        # indexed_draw_commands: List of draw commands for rendering each individual
        # square (with index buffer)
        geometry, indices, array_draw_commands, indexed_draw_commands = self.gen_data(
            self.num_objects
        )

        # Crate OpenGL buffers of the data
        self.draw_command_array_buffer = self.ctx.buffer(data=array("I", array_draw_commands))
        self.draw_command_indexed_buffer = self.ctx.buffer(data=array("I", indexed_draw_commands))
        self.vbo = self.ctx.buffer(data=array("f", geometry))
        self.ibo = self.ctx.buffer(data=array("i", indices))

        # Our geometry without index buffer
        self.geometry_array = self.ctx.geometry(
            [BufferDescription(self.vbo, "2f 4f", ("in_pos", "in_color"))],
            mode=self.ctx.TRIANGLES,
        )
        # Our geometry with index buffer
        self.geometry_indexed = self.ctx.geometry(
            [BufferDescription(self.vbo, "2f 4f", ("in_pos", "in_color"))],
            index_buffer=self.ibo,
            mode=self.ctx.TRIANGLES,
        )

        # Simple program rendering colored geometry
        self.program = self.ctx.program(
            vertex_shader="""
            #version 330

            // Piggyback on pyglet's global projection matrix
            uniform WindowBlock {
                mat4 projection;
                mat4 view;
            } window;

            in vec2 in_pos;
            in vec4 in_color;
            out vec4 color;

            void main() {
                gl_Position = window.projection * window.view * vec4(in_pos, 0.0, 1.0);
                color = in_color;
            }
            """,
            fragment_shader="""
            #version 330

            in vec4 color;
            out vec4 fragColor;

            void main() {
                fragColor = vec4(color);
            }
            """,
        )

    def gen_data(self, count):
        # Simple unit quad consisting of 2 rectangles
        quad = (
            # First triangle
            (-0.5, 0.5),  # Top left
            (-0.5, -0.5),  # Bottom left
            (0.5, 0.5),  # Top right
            # Second triangle
            (0.5, 0.5),  # Top right
            (-0.5, -0.5),  # Bottom left
            (0.5, -0.5),  # Bottom right
        )

        geometry = []
        index_buffer = []
        array_commands = []
        indexed_commands = []

        # Create count quads
        for quad_index in range(count):
            # Apply random offset, scale and colors
            offset_x = randint(0, WINDOW_WIDTH)
            offset_y = randint(0, WINDOW_HEIGHT)
            scale = randint(5, 15)
            r, g, b, a = random(), random(), random(), uniform(0.2, 0.8)

            # Emit a quad
            for index, (x, y) in enumerate(quad):
                geometry.append(x * scale + offset_x)
                geometry.append(y * scale + offset_y)
                # Color needs to be specified per vertex
                geometry.append(r)
                geometry.append(g)
                geometry.append(b)
                geometry.append(a)
                # Add to index buffer
                # NOTE: This is a rather silly index buffer not reusing indices
                index_buffer.append(index)

            # Create a rendering command for array rendering (no index buffer)
            # Expected data structure:
            # typedef  struct {
            #     uint  count;
            #     uint  instanceCount;
            #     uint  first;
            #     uint  baseInstance;
            # } DrawArraysIndirectCommand;
            array_commands.extend(
                (
                    6,  # count: How many vertices to render
                    1,  # instanceCount: How many instances. Typically 1.
                    quad_index * 6,  # first: What vertex index the object starts on
                    0,  # baseInstance: The base instance id (ignore when instanceCount = 1)
                )
            )
            # Create a rendering command for indexed rendering (with index buffer)
            # Expected data structure:
            # typedef  struct {
            #     GLuint  count;
            #     GLuint  instanceCount;
            #     GLuint  firstIndex;
            #     GLuint  baseVertex;
            #     GLuint  baseInstance;
            # } DrawElementsIndirectCommand;
            indexed_commands.extend(
                (
                    6,  # count: How many indices to render
                    1,  # instanceCount: How many instances to render. Typically 1
                    quad_index * 6,  # firstIndex: Position of the first index (ibo)
                    quad_index * 6,  # baseVertex: Position of the first vertex (vbo)
                    0,  # baseInstance: What value gl_InstanceID should start on
                )
            )

        return geometry, index_buffer, array_commands, indexed_commands

    def on_draw(self):
        self.clear()
        method = ""

        self.ctx.enable(self.ctx.BLEND)

        # Cycle all three rendering methods:
        # * Direct rendering of all quads
        # * Indirectly draw each individual quad (no index buffer)
        # * Indirectly draw each individual quad (with index buffer)
        with self.query:
            if self.frame % 3 == 0:
                # Draw the rectangles as one object
                # NOTE: first and vertices are not necessary if we are drawing all the geometry.
                method = "direct           "
                self.geometry_array.render(
                    self.program,
                    first=self.start * 6,
                    vertices=self.count * 6,
                )
            elif self.frame % 3 == 1:
                # Draw the rectangles as separate objects
                # NOTE: first and vertices are not necessary if we are drawing all the geometry.
                method = "indirect(array)  "
                self.geometry_array.render_indirect(
                    self.program,
                    self.draw_command_array_buffer,
                    first=self.start,
                    count=self.count,
                )
            else:
                method = "indirect(indexed)"
                # Draw rectangles as separate objects with index buffer
                # NOTE: first and vertices are not necessary if we are drawing all the geometry.
                self.geometry_indexed.render_indirect(
                    self.program,
                    self.draw_command_indexed_buffer,
                    first=self.start,
                    count=self.count,
                )

        # We want to prove by using a query that the following
        # is identical for all the rendering methods:
        # * The number of pixels written (rgba / 4)
        # * The number of primitives generate
        # We also compare the nanoseconds each methods is using
        # In addition we can tweak the start and count value
        # to prove that partial rendering also works.
        # NOTE: These values can be skewed if vsync is enabled
        print(

                f"[{method}] "
                f"Pixels written = {self.query.samples_passed // 4}, "
                f"Primitives Generated = {self.query.primitives_generated}, "
                f"time = {self.query.time_elapsed / 1_000_000_000}s"

        )

    def on_update(self, delta_time: float):
        self.frame += 1


RenderIndirect().run()
