// https://github.com/frozein/VkGalaxy/blob/main/assets/shaders/particle_generate.comp
#version 450
#extension GL_EXT_debug_printf : enable

layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;

#define PI 3.14159

struct particle
{
    vec2 pos;
    float height;
    float angle;
    float tilt;
    float vel;
    float opacity;
    float temp;
};

struct draw_indirect_cmd
{
    uint vertex_count;
    uint instance_count;
    uint first_vertex;
    uint first_instance;
};

layout (binding = 0) uniform frustum_info
{
    vec4 planes[6];
}   frustum;

layout (std140, binding = 1) buffer stars
{
    particle particles[];
}   particles;

layout (std140, binding = 2) writeonly buffer indirect_draws
{
    draw_indirect_cmd draws[];
}   draws;

layout (binding = 3) writeonly buffer count_buffer
{
    uint draw_count;
}   count;

layout (push_constant) uniform galaxy
{
    uint star_count;
    float max_rad, bulge_rad;
    float angle, eccentricity;
    float base_height, height;
    float dust_temp, min_temp, max_temp;
    float min_star_opacity, max_star_opacity, min_dust_opacity, max_dust_opacity;
    float speed;
}   galaxy_params;

// http://www.gamedev.net/topic/592001-random-number-generation-based-on-time-in-hlsl/
#define RANDOM_IA 16807
#define RANDOM_IM 2147483647
#define RANDOM_AM 1.0 / float(RANDOM_IM)
#define RANDOM_IQ 127773
#define RANDOM_IR 2836
#define RANDOM_MASK 123459876

int _SEED = 0;

void srand_cycle(void)
{
    _SEED ^= RANDOM_MASK;
    int k = _SEED / RANDOM_IQ;
    _SEED = RANDOM_IA * (_SEED - k * RANDOM_IQ) - RANDOM_IR * k;

    if (_SEED < 0)
        _SEED += RANDOM_IM;

    _SEED ^= RANDOM_MASK;
}

void srand_set(int seed)
{
    _SEED = seed;
    srand_cycle();
}

float rand(void)
{
    srand_cycle();
    return RANDOM_AM * _SEED;
}

float ease_in_exp(float x)
{
    return x <= 0.0 ? 0.0 : pow(2, 10.0 * x - 10.0);
}

float ease_in_circ(float x)
{
    return x >= 1.0 ? 1.0 : 1.0 - sqrt(1.0 - x * x);
}

float rand_height(void)
{
    float r = ease_in_circ(rand());
    if (rand() < 0.5)
        r *= -1;

    return galaxy_params.base_height + (galaxy_params.height * 0.5) * r;
}

float rand_height_bulge(float rad)
{
    float r = ease_in_circ(rand());
    if (rand() < 0.5)
        r *= -1;

    float bound = (galaxy_params.height * 0.5) + (galaxy_params.height * 0.5) * cos(PI * rad / galaxy_params.bulge_rad);

    return galaxy_params.base_height + bound * r;
}

float rand_height_bulge_dust(float rad)
{
    float r = ease_in_circ(rand());
    if (rand() < 0.5)
        r *= -1;

    float bound = (galaxy_params.height * 0.5) * cos(PI * rad / galaxy_params.bulge_rad);

    return galaxy_params.base_height + bound * r;
}

bool frustum_check(vec4 pos, float radius)
{
    for (uint i = 0; i < 6; i++)
    {
        if (dot(pos, frustum.planes[i]) + radius < 0.0)
        {
            return false;
        }
    }

    return true;
}

void main(void)
{
    srand_set(int(gl_GlobalInvocationID.x));
    particle p;
    float speed = galaxy_params.speed;
    uint type = 0;

    if (gl_GlobalInvocationID.x < galaxy_params.star_count) // Star
    {
        float rad = ease_in_exp(rand()) * galaxy_params.max_rad;

        p.pos = vec2(rad, galaxy_params.eccentricity * rad);
        p.angle = rand() * 2.0 * PI;
        p.tilt = (rad / galaxy_params.max_rad) * galaxy_params.angle;
        p.vel = -speed * sqrt(1.0 / rad);
        p.opacity = galaxy_params.min_star_opacity + (galaxy_params.max_star_opacity - galaxy_params.min_star_opacity) * rand();
        p.temp = galaxy_params.min_temp + (galaxy_params.max_temp - galaxy_params.min_temp) * rand();

        if (rad < galaxy_params.bulge_rad)
            p.height = rand_height_bulge(rad);
        else
            p.height = rand_height();

        type = 0;
    }
    else // Dust
    {
        float rad;
        if (gl_GlobalInvocationID.x % 2 == 0)
            rad = rand() * galaxy_params.max_rad;
        else
            rad = ease_in_exp(rand()) * galaxy_params.max_rad;

        p.pos = vec2(rad, galaxy_params.eccentricity * rad);
        p.angle = rand() * 2.0 * PI;
        p.tilt = (rad / galaxy_params.max_rad) * galaxy_params.angle;
        p.vel = -speed * sqrt(1.0 / rad);
        p.temp = galaxy_params.dust_temp + 2.0 * rad;
        p.opacity = galaxy_params.min_dust_opacity + (galaxy_params.max_dust_opacity - galaxy_params.min_dust_opacity) * rand();

        if (rad < galaxy_params.bulge_rad)
            p.height = rand_height_bulge_dust(rad);
        else
            p.height = galaxy_params.base_height;

        type = 1;
    }

    particles.particles[gl_GlobalInvocationID.x] = p;

    vec4 pos = {p.pos.x, p.height, p.pos.y, 1.0};

    float size = 0.0;
    if (type == 0)  // Star
    {
        size = 10.0;
    }
    else    // Dust
    {
        size = 500.0;
    }

    debugPrintfEXT("Size: %f\n", size);

    // Cull against stars
    if (frustum_check(pos, size))
    {
        draws.draws[gl_GlobalInvocationID.x].vertex_count = 1;
        draws.draws[gl_GlobalInvocationID.x].instance_count = 1;
        draws.draws[gl_GlobalInvocationID.x].first_vertex = 0;
        draws.draws[gl_GlobalInvocationID.x].first_instance = 0;
        atomicAdd(count.draw_count, 1);
        debugPrintfEXT("Culled\n");
    }
    else
    {
        debugPrintfEXT("Not culled\n");
    }
}