#include <cuda_runtime.h>

#include "shared_binning.cuh"

_USE_MATH_DEFINES
#include <math.h>

#include "main.h"

/* apparently the compiler really hates when you try declare these as either
 * extern or declare them in another file and try to set them from a different
 * file, silently failing. :( */
__constant__ int c_bin_size;
__constant__ int c_bin_dim_x;
__constant__ int c_bin_dim_y;

void
set_bin_constants(int bin_size, int bin_dim_x, int bin_dim_y)
{
	cudaMemcpyToSymbol(c_bin_size, &bin_size, sizeof(int));
	cudaMemcpyToSymbol(c_bin_dim_x, &bin_dim_x, sizeof(int));
	cudaMemcpyToSymbol(c_bin_dim_y, &bin_dim_y, sizeof(int));
}

/* Applies / searches neighbors to apply gravity with one thread-block mapping
 * to one bin Actual updating of movement occurs below. Projectiles that do not
 * fit into a bin simply keep their last computed velocity until they can find a
 * bin (handled below), it was either this, or skip updating them. */
__global__ void
update_binned_shared(Projectile* projs, Projectile* projs_snapshot, int* bins,
                     int* bin_counts, float delta_t)

{
	int bin_idx = blockIdx.x;
	int slot = threadIdx.x;
	int total_bins = c_bin_dim_x * c_bin_dim_y;

	if (bin_idx >= total_bins)
	{
		return;
	}

	extern __shared__ float4 sh_neighborhood[];
	/* + sizeof(int) * (9 + 9 + 1 + 1)*/
	__shared__ int sh_neighbor_ids[9];
	__shared__ int sh_neighbor_counts[9];
	__shared__ int sh_num_neighbors;
	__shared__ int sh_own_neighbor_idx;


	if (slot == 0)
	{
		int bx = bin_idx % c_bin_dim_x;
		int by = bin_idx / c_bin_dim_x;
		int n = 0;
		sh_own_neighbor_idx = 0;

		for (int nby = by - 1; nby <= by + 1; nby++)
		{
			for (int nbx = bx - 1; nbx <= bx + 1; nbx++)
			{
				if (nbx < 0 || nbx >= c_bin_dim_x || nby < 0 ||
				    nby >= c_bin_dim_y)
				{
					continue;
				}

				int nb = nby * c_bin_dim_x + nbx;
				sh_neighbor_ids[n] = nb;
				sh_neighbor_counts[n] = min(bin_counts[nb], c_bin_size);

				if (nb == bin_idx)
				{
					sh_own_neighbor_idx = n;
				}

				n++;
			}
		}
		sh_num_neighbors = n;
	}
	__syncthreads();

	/* thread n loads neighbor bin n entirely into shared memory */
	if (slot < sh_num_neighbors)
	{
		int nb = sh_neighbor_ids[slot];
		int nb_count = sh_neighbor_counts[slot];
		int base = slot * c_bin_size;

		for (int s = 0; s < c_bin_size; s++)
		{
			if (s < nb_count)
			{
				int j = bins[nb * c_bin_size + s];
				Projectile* q = &projs_snapshot[j];
				sh_neighborhood[base + s].x = q->pos.x;
				sh_neighborhood[base + s].y = q->pos.y;
				sh_neighborhood[base + s].z = q->gravity;
				sh_neighborhood[base + s].w = __int_as_float(j);
			}
			else
			{
				sh_neighborhood[slot * c_bin_size + s] =
					make_float4(0.0f, 0.0f, 0.0f, __int_as_float(-1));
			}
		}
	}
	__syncthreads();

	int own_count = sh_neighbor_counts[sh_own_neighbor_idx];
	if (slot >= own_count)
	{
		return;
	}

	int own_base = sh_own_neighbor_idx * c_bin_size;
	int my_proj_idx = __float_as_int(sh_neighborhood[own_base + slot].w);

	if (my_proj_idx < 0 || !projs_snapshot[my_proj_idx].active)
	{
		return;
	}

	Projectile p = projs_snapshot[my_proj_idx];
	float normalized_dt = delta_t * 60.0f;
	float grav_x = 0.0f;
	float grav_y = 0.0f;

	int total_slots = sh_num_neighbors * c_bin_size;
	for (int s = 0; s < total_slots; s++)
	{
		int j = __float_as_int(sh_neighborhood[s].w);
		if (j < 0 || j == my_proj_idx || sh_neighborhood[s].z == 0.0f)
		{
			continue;
		}

		float dx = sh_neighborhood[s].x - p.pos.x;
		float dy = sh_neighborhood[s].y - p.pos.y;
		float dist_sq = dx * dx + dy * dy;

		if (dist_sq < 1.0f)
		{
			continue;
		}

		float dist = sqrtf(dist_sq);
		float force = sh_neighborhood[s].z / dist_sq;

		grav_x += (dx / dist) * force;
		grav_y += (dy / dist) * force;
	}

	projs[my_proj_idx].velocity.x = p.velocity.x + grav_x * normalized_dt;
	projs[my_proj_idx].velocity.y = p.velocity.y + grav_y * normalized_dt;
}

/* second kernel handles updating the movement for ALL of the active projectiles
 * -- including those that weren't binned. Previously those that weren't binned
 * got ignored entirely, which while *technically* correct behavior, I
 * personally thought that was unacceptable, so I decided we can hide it by
 * splitting the movement updating into two kernels, a more complex one that can
 * take better advantage of shared memory to update gravity, and another to just
 * add the velocity like before gravity was implemented. */

/* We also don't need a second buffer as there is no race condition with every
 * projectile just updating it's own velocity */
__global__ void
update_velocity_shared(Projectile* projs, float delta_t)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	if (i >= MAX_PROJS)
	{
		return;
	}

	Projectile p = projs[i];

	if (!p.active)
	{
		return;
	}

	float normalized_dt = delta_t * 60.0f;


	p.velocity.x = p.velocity.x;
	p.velocity.y = p.velocity.y;

	update_standard_projectile_velocity(&p, 0, 0, delta_t, normalized_dt);

	// /* despawn if out of bounds */
	// if (p.pos.x < 0.0f || p.pos.x > max_width || p.pos.y < 0.0f ||
	//     p.pos.y > max_height)
	// {
	// 	p.active = 0;
	// }

	projs[i] = p;
}