#include <cuda_runtime.h>
#define _USE_MATH_DEFINES
#include <math.h>

#include "naive_update.h"
#include "main.h"
#include "utils.h"
#include "naive_binning.cuh"

/* This file is responsible for naive binning implementation (naive meaning not
 * taking advantage of the memory hierarchy, as it seems to do that I would need
 * to completely redo all of this. At least for accuracy's sake; attempting to
 * do shared memory like this initially seems simple, but there lies an issue
 * with the projectiles where if there were too many in a bin (likely because
 * they all attract each other with positive gravity), projectiles that weren't
 * within the initial BIN_SIZE would just not render whatsoever).
 *
 * See shared_binning.cu for how this was (hopefully by the time you're reading
 * this) resolved. )*/

__global__ void
bin_projectiles(Projectile* projs, int* bins, int* bin_counts, int bin_size,
                int bin_dim_x, int bin_dim_y, float bin_w, float bin_h)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	if (i >= MAX_PROJS)
	{
		return;
	}

	Projectile p = projs[i];

	if (!p.active)
	{
		return;
	}

	/* compute which bin this projectile falls into */
	int bx = (int)(p.pos.x / bin_w);
	int by = (int)(p.pos.y / bin_h);

	/* clamp to grid bounds */
	bx = max(0, min(bx, bin_dim_x - 1));
	by = max(0, min(by, bin_dim_y - 1));

	int bin_idx = by * bin_dim_x + bx;

	int slot = atomicAdd(&bin_counts[bin_idx], 1);

	if (slot < bin_size)
	{
		bins[bin_idx * bin_size + slot] = i;
	}
}

__global__ void
update_binned(Projectile* projs, Projectile* projs_snapshot, int* bins,
              int* bin_counts, int bin_size, float delta_t, int bin_dim_x,
              int bin_dim_y, float bin_w, float bin_h)
{
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	if (i >= MAX_PROJS)
	{
		return;
	}

	if (!projs_snapshot[i].active)
	{
		return;
	}

	Projectile p = projs_snapshot[i];
	float normalized_dt = delta_t * 60.0f;

	/* find which bin this projectile is in */
	int bx = (int)(p.pos.x / bin_w);
	int by = (int)(p.pos.y / bin_h);
	bx = max(0, min(bx, bin_dim_x - 1));
	by = max(0, min(by, bin_dim_y - 1));

	float grav_x = 0.0f;
	float grav_y = 0.0f;

	/* visit 3x3 halo of surrounding bins per slide 90 */
	for (int nby = by - 1; nby <= by + 1; nby++)
	{
		for (int nbx = bx - 1; nbx <= bx + 1; nbx++)
		{
			/* oob protection */
			if (nbx < 0 || nbx >= bin_dim_x || nby < 0 || nby >= bin_dim_y)
			{
				continue;
			}

			int neighbor_bin = nby * bin_dim_x + nbx;
			int n_in_bin = min(bin_counts[neighbor_bin], bin_size);

			/* actual gravity computation */
			for (int s = 0; s < n_in_bin; s++)
			{
				int j = bins[neighbor_bin * bin_size + s];

				if (j == i || !projs_snapshot[j].active ||
				    projs_snapshot[j].gravity == 0)
				{
					continue;
				}

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

				if (dist_sq < 1.0f)
				{
					continue;
				}

				float dist = sqrtf(dist_sq);
				float force = projs_snapshot[j].gravity / dist_sq;

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

	update_standard_projectile_velocity(&p, grav_x, grav_y, delta_t,
	                                    normalized_dt);

	projs[i] = p;
}