#include "ParticleEmitter.h"

ParticleEmitter::ParticleEmitter(glm::vec3 initialPosition) : 
	render_shader("shaders\\particles.vert", "shaders\\particles.frag", "shaders\\particles.geom"),
	compute_shader("shaders\\particles.comp")
{
	startPosition = initialPosition;
	is_emitting = false;
	reset(true);
}

ParticleEmitter::~ParticleEmitter()
{
	glDeleteBuffers(2, &SSBO_positions[0]);
	glDeleteBuffers(2, &SSBO_velocities[0]);
	glDeleteBuffers(2, &VAOs[0]);
}

void ParticleEmitter::simulate(float timeDelta)
{
	compute_shader.use();

	compute_shader.setUInt("LastCount", particle_count);
	compute_shader.setUInt("MaximumCount", static_cast<unsigned int>(NUM_PARTICLES));
	compute_shader.setFloat("DeltaT", timeDelta);

	// set SSBO and atomic counters...
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, SSBO_positions[current_vao_index]);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, SSBO_velocities[current_vao_index]);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, SSBO_positions[!current_vao_index]);
	glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 3, SSBO_velocities[!current_vao_index]);
	glBindBufferBase(GL_ATOMIC_COUNTER_BUFFER, 4, atomic_counter);
	current_vao_index = !current_vao_index;
	
	// Execute compute shader with a 256 x 1 x 1 work group size
	GLuint groups = (NUM_PARTICLES + 256 - 1) / 256;

	glDispatchCompute(groups, 1, 1);
	glMemoryBarrier(GL_ATOMIC_COUNTER_BARRIER_BIT);

	glBindBuffer(GL_ATOMIC_COUNTER_BUFFER, atomic_counter);
	GLuint* counterValue = (GLuint*)glMapBufferRange(GL_ATOMIC_COUNTER_BUFFER, 0, sizeof(GLuint), GL_MAP_READ_BIT | GL_MAP_WRITE_BIT);
	particle_count = counterValue[0];
	
	counterValue[0] = 0;
	glUnmapBuffer(GL_ATOMIC_COUNTER_BUFFER);
	glMemoryBarrier(GL_VERTEX_ATTRIB_ARRAY_BARRIER_BIT);
}

void ParticleEmitter::draw()
{
 	enableBlendMode();

	// shader with VS -> GS -> FS
	render_shader.use();

	// set color
	render_shader.setVec3F("color", current_color);

	glBindVertexArray(VAOs[current_vao_index]); // bind VAO
	glDrawArrays(GL_POINTS, 0, NUM_PARTICLES);
	glBindVertexArray(0);
	glUseProgram(0);

	disableBlendMode();
}

void ParticleEmitter::update_shader(glm::mat4 viewMatrix, glm::mat4 projectionMatrix)
{
	render_shader.use();

	glm::mat4 ModelMatrix(1.0f);
	render_shader.setMat4F("modelMatrix", ModelMatrix);
	render_shader.setMat4F("viewMatrix", viewMatrix);
	render_shader.setMat4F("projectionMatrix", projectionMatrix);
}

void ParticleEmitter::enableBlendMode() {
	glEnable(GL_BLEND);
	glDepthMask(GL_FALSE);
	glBlendFunc(GL_SRC_COLOR, GL_SRC_COLOR);
	glBlendEquation(GL_MAX);
}

void ParticleEmitter::disableBlendMode() {
	glDisable(GL_BLEND);
	glDepthMask(GL_TRUE);
	glBlendEquation(GL_FUNC_ADD);
}

void ParticleEmitter::reset(bool generate_buffers) 
{
	//if (delete_buffers) 
	//{
	//	glDeleteBuffers(2, &SSBO_positions[0]);
	//	glDeleteBuffers(2, &SSBO_velocities[0]);
	//	glDeleteBuffers(2, &VAOs[0]);
	//}

	particle_count = 0;
	std::mt19937 rng(time(0)); // random generator seed
	std::normal_distribution<double> distribution(0, 1);

	// create instance of each particle
	for (int i = 0; i < NUM_PARTICLES; i++)
	{
		float velocity_x = distribution(rng);
		float velocity_y = distribution(rng);
		float velocity_z = distribution(rng);

		positions.push_back(glm::vec4(startPosition, TTL));
		velocities.push_back(glm::vec4(velocity_x, velocity_y, velocity_z, 0.0f));
	}

	// generate *two* SSBOs for positions and velocities
	if (generate_buffers)
		glGenBuffers(2, &SSBO_positions[0]);
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_positions[0]);
	glBufferData(GL_SHADER_STORAGE_BUFFER, NUM_PARTICLES * sizeof(glm::vec4), NULL, GL_DYNAMIC_DRAW);
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_positions[1]);
	glBufferData(GL_SHADER_STORAGE_BUFFER, NUM_PARTICLES * sizeof(glm::vec4), NULL, GL_DYNAMIC_DRAW);

	if (generate_buffers)
		glGenBuffers(2, &SSBO_velocities[0]);
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_velocities[0]);
	glBufferData(GL_SHADER_STORAGE_BUFFER, NUM_PARTICLES * sizeof(glm::vec4), NULL, GL_DYNAMIC_DRAW);
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_velocities[1]);
	glBufferData(GL_SHADER_STORAGE_BUFFER, NUM_PARTICLES * sizeof(glm::vec4), NULL, GL_DYNAMIC_DRAW);

	// copy data -> only one SSBO for positions and vertices is enough
	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_positions[0]);
	glBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, NUM_PARTICLES * sizeof(positions[0]), &positions[0]);

	glBindBuffer(GL_SHADER_STORAGE_BUFFER, SSBO_velocities[0]);
	glBufferSubData(GL_SHADER_STORAGE_BUFFER, 0, NUM_PARTICLES * sizeof(velocities[0]), &velocities[0]);

	// generate VAOs for each position SSBO
	if (generate_buffers)
		glGenVertexArrays(2, &VAOs[0]);

	glBindVertexArray(VAOs[0]);
	glBindBuffer(GL_ARRAY_BUFFER, SSBO_positions[0]);
	glEnableVertexAttribArray(position_location);
	glVertexAttribPointer(position_location, 4, GL_FLOAT, GL_FALSE, 0, NULL);

	glBindVertexArray(VAOs[1]);
	glBindBuffer(GL_ARRAY_BUFFER, SSBO_positions[1]);
	glEnableVertexAttribArray(position_location);
	glVertexAttribPointer(position_location, 4, GL_FLOAT, GL_FALSE, 0, NULL);

	// generate atomic counter
	glGenBuffers(1, &atomic_counter);
	glBindBuffer(GL_ATOMIC_COUNTER_BUFFER, atomic_counter);
	glBufferData(GL_ATOMIC_COUNTER_BUFFER, sizeof(GLuint), NULL, GL_DYNAMIC_DRAW);

	// initialize value of atomic counter to zero
	GLuint value = 0;
	glBufferSubData(GL_ATOMIC_COUNTER_BUFFER, 0, sizeof(GLuint), &value);
	glBindBuffer(GL_ATOMIC_COUNTER_BUFFER, 0);

	// set color
	current_color = rocket_colors[color_index];
	color_index = (color_index + 1) % rocket_colors.size();
}

