#include "boid-playground.hpp"
#include "raycast.hpp"
#include "rprof.h"
#include "boid-list.hpp"

#include <immintrin.h>

static float vector2_atan2(Vector2 a) {
	return std::atan2(a.y, a.x);
}

static Vector2 vector2_mul_value(Vector2 v, float value) {
	return { v.x * value, v.y * value };
}

static Vector2 vector2_div_value(Vector2 v, float value) {
	return { v.x / value, v.y / value };
}

static Vector2 vector2_from_angle(float angle) {
	return { std::cos(angle), std::sin(angle) };
}

static Vector2 get_center_point(std::vector<Vector2> &points) {
	Vector2 center = { 0, 0 };
	for (int i = 0; i < points.size(); i++) {
		center.x += points[i].x;
		center.y += points[i].y;
	}
	center.x /= points.size();
	center.y /= points.size();
	return center;
}

static void fill_avoidance_ray_angles(float *rays, int ray_count, float ray_angle) {
	DEBUG_ASSERT(ray_count >= 1 && "Ray count must be at least 1");
	DEBUG_ASSERT(((ray_count - 1) % 2 == 0) && "Ray count must be a multiple of 2n+1");

	rays[0] = 0;

	int side_ray_count = ((ray_count-1)/2);
	float ray_angle_step = ray_angle / side_ray_count;
	for (int i = 0; i < side_ray_count; i++) {
		rays[2*i+0 + 1] = ray_angle_step * (i+1);
		rays[2*i+1 + 1] = -ray_angle_step * (i+1);
	}
}

static void fill_avoidance_ray_dirs(Vector2 *rays, int ray_count, float ray_angle) {
	float ray_angles[ray_count];
	fill_avoidance_ray_angles(ray_angles, ray_count, ray_angle);

	for (int i = 0; i < ray_count; i++) {
		rays[i].x = cos(ray_angles[i]);
		rays[i].y = sin(ray_angles[i]);
	}
}

static Vector2 get_collision_avoidance_dir(World *world, Boid *boid, Vector2 *directions, int direction_count) {
	int best_avoidance = -1;
	Vector2 avoidance_dir = { 0, 0 };
	bool got_hit = false;
	RayHitResult hit_results[direction_count];

	for (int i = 0; i < direction_count; i++) {
		Vector2 ray_dir;
		ray_dir.x = boid->dir.x * directions[i].x - boid->dir.y * directions[i].y;
		ray_dir.y = boid->dir.x * directions[i].y + boid->dir.y * directions[i].x;

		get_intersect_with_world(&hit_results[i], boid->pos, ray_dir, world);
		if (hit_results[i].hit != -1 && hit_results[i].hit <= world->collision_avoidance_distance) {
			got_hit = true;
		}

		if (hit_results[i].hit > hit_results[best_avoidance].hit || best_avoidance == -1) {
			avoidance_dir = ray_dir;
			best_avoidance = i;
		}
	}

	if (got_hit) {
		return avoidance_dir;
	} else {
		return { 0, 0 };
	}
}

static int count_out_of_bounds_boids(World *world) {
	int count = 0;
	for (int i = 0; i < world->boids.size(); i++) {
		Vector2 *pos = &world->boids[i].pos;

		bool x_out_of_bounds = (pos->x <= 0 || pos->x >= world->size.x);
		bool y_out_of_bounds = (pos->y <= 0 || pos->y >= world->size.y);
		if (x_out_of_bounds || y_out_of_bounds)  {
			count++;
		}
	}
	return count;
}

// -------------------- Init/Cleanup ------------------------

static void boid_rand_init(World *world, Boid *boid, float border) {
	float world_width = world->size.x;
	float world_height = world->size.y;
	boid->pos.x = GetRandomValue(border, world_width-border);
	boid->pos.y = GetRandomValue(border, world_height-border);

	float facing = GetRandomValue(0, 2*PI);
	boid->dir = Vector2Rotate({ 1, 0 }, facing);
	boid->speed = GetRandomValue(world->min_speed, world->max_speed);
}

static void world_init(World *world, float width, float height) {
	arena_init(&world->frame_arena, 1024 * 1024 * 256);
	world->size = { width, height };
}

static void world_free(World *world) {
	arena_free(&world->frame_arena);
}

// --------------------- Utils -----------------------

void world_adjust_boid_count(World *world, int new_boid_count) {
	if (new_boid_count >= MAX_BOIDS) return;

	int boid_count = world->boids.size();
	int diff = new_boid_count - boid_count;
	if (diff > 0) {
		for (int i = 0; i < diff; i++) {
			Boid boid;
			boid_rand_init(world, &boid, 5);
			world->boids.push_back(boid);
		}
	} else if (diff < 0) {
		for (int i = 0; i < -diff ; i++) {
			world->boids.pop_back();
		}
	}
}

// --------------------- Update -----------------------

struct ChunkGrid {
	BoidList *data;
	int width;
	int height;
};

static size_t chunkgrid_get_idx(ChunkGrid *grid, int x, int y) {
	return y * grid->width + x;
}

static BoidList *chunkgrid_get(ChunkGrid *grid, int x, int y) {
	return &grid->data[chunkgrid_get_idx(grid, x, y)];
}

static void chunkgrid_init(MemoryArena *arena, ChunkGrid *grid, int width, int height) {
	grid->data = (BoidList*)arena_malloc(arena, width * height * sizeof(BoidList));
	grid->width = width;
	grid->height = height;

	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {
			boid_list_init(chunkgrid_get(grid, x, y));
		}
	}
}

static int g_prof_interactions = 0;

static int nearest_multiple(int num, int divisor) {
	return (num / divisor + (num % divisor > 0 ? 1 : 0)) * divisor;
}

// b2l = boid to (list of boids) comparison
static void assign_local_boids_b2l(World *world, BoidList *local_boids, uboid_t from_boid, uboid_t *to_boids, uboid_t to_boids_count) {
	Boid *boids = world->boids.data();

	// Simplified from: float dot_threshold = Vector2DotProduct(dir, Vector2Rotate(dir, world->view_angle/2));
	float dot_threshold = cosf(world->view_angle/2);
	float view_radius_sqr = world->view_radius * world->view_radius;

	for (int i = 0; i < to_boids_count; i++) {
		uint16_t to_boid = to_boids[i];
		assert(to_boid != from_boid);

		Vector2 offset = Vector2Subtract(boids[from_boid].pos, boids[to_boid].pos);
		float length_sqr = Vector2LengthSqr(offset);
		if (length_sqr > view_radius_sqr) continue;

		Vector2 normalized = offset;
		if (length_sqr != 0)
		{
			float ilength = 1.0f/sqrtf(length_sqr);
			normalized.x *= ilength;
			normalized.y *= ilength;
		}

		if (Vector2DotProduct(boids[from_boid].dir, Vector2Negate(normalized)) >= dot_threshold) {
			boid_list_append(&world->frame_arena, &local_boids[from_boid], to_boid);
			g_prof_interactions++;
		}
		if (Vector2DotProduct(boids[to_boid].dir, normalized) >= dot_threshold) {
			boid_list_append(&world->frame_arena, &local_boids[to_boid], from_boid);
			g_prof_interactions++;
		}
	}
}

#ifndef __EMSCRIPTEN__
static void print_m256_f32(__m256 value) {
	float *value_f32 = (float*)&value;
	printf("%f", value_f32[0]);
	for (int i = 1; i < 8; i++) {
		printf(",%f", value_f32[i]);
	}
	printf("\n");
}

static inline bool mm256_is_zero(__m256i x) {
	return _mm256_testz_si256(x, x);
}
static inline bool mm256_is_zero(__m256 x) {
	return mm256_is_zero((__m256i)x);
}

struct boid_pair {
	uboid_t from;
	uboid_t to;
};

static void world_process_local_boid_pairs(World *world, BoidList *local_boids, boid_pair *b2b_cmps, int *b2b_cmps_count) {
	// Simplified from:
	//   float dot_threshold = Vector2DotProduct(dir, Vector2Rotate(dir, world->view_angle/2));
	//                         |
	//                         v
	//   float dot_threshold = cosf(world->view_angle/2);
	__m256 dot_threshold = _mm256_set1_ps(cosf(world->view_angle/2));

	float view_radius_sqr = world->view_radius * world->view_radius;
	__m256 view_radius = _mm256_set1_ps(view_radius_sqr);
	__m256 zero = _mm256_set1_ps(0);
	__m256 negative_one = _mm256_set1_ps(-1);
	Boid *boids = world->boids.data();

	int32_t *do_append_mask1_f32 = (int32_t*)arena_malloc(&world->frame_arena, sizeof(int32_t)*8, 32);
	int32_t *do_append_mask2_f32 = (int32_t*)arena_malloc(&world->frame_arena, sizeof(int32_t)*8, 32);

	int simd_iteration_count = nearest_multiple(*b2b_cmps_count, 8)/8;
	for (int i = 0; i < simd_iteration_count; i++) {
		__m256 from_pos_x = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].from].pos.x,
			boids[b2b_cmps[8*i+6].from].pos.x,
			boids[b2b_cmps[8*i+5].from].pos.x,
			boids[b2b_cmps[8*i+4].from].pos.x,
			boids[b2b_cmps[8*i+3].from].pos.x,
			boids[b2b_cmps[8*i+2].from].pos.x,
			boids[b2b_cmps[8*i+1].from].pos.x,
			boids[b2b_cmps[8*i+0].from].pos.x
		);

		__m256 from_pos_y = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].from].pos.y,
			boids[b2b_cmps[8*i+6].from].pos.y,
			boids[b2b_cmps[8*i+5].from].pos.y,
			boids[b2b_cmps[8*i+4].from].pos.y,
			boids[b2b_cmps[8*i+3].from].pos.y,
			boids[b2b_cmps[8*i+2].from].pos.y,
			boids[b2b_cmps[8*i+1].from].pos.y,
			boids[b2b_cmps[8*i+0].from].pos.y
		);

		__m256 to_pos_x = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].to].pos.x,
			boids[b2b_cmps[8*i+6].to].pos.x,
			boids[b2b_cmps[8*i+5].to].pos.x,
			boids[b2b_cmps[8*i+4].to].pos.x,
			boids[b2b_cmps[8*i+3].to].pos.x,
			boids[b2b_cmps[8*i+2].to].pos.x,
			boids[b2b_cmps[8*i+1].to].pos.x,
			boids[b2b_cmps[8*i+0].to].pos.x
		);

		__m256 to_pos_y = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].to].pos.y,
			boids[b2b_cmps[8*i+6].to].pos.y,
			boids[b2b_cmps[8*i+5].to].pos.y,
			boids[b2b_cmps[8*i+4].to].pos.y,
			boids[b2b_cmps[8*i+3].to].pos.y,
			boids[b2b_cmps[8*i+2].to].pos.y,
			boids[b2b_cmps[8*i+1].to].pos.y,
			boids[b2b_cmps[8*i+0].to].pos.y
		);

		__m256 sub_x = _mm256_sub_ps(from_pos_x, to_pos_x);
		__m256 sub_y = _mm256_sub_ps(from_pos_y, to_pos_y);

		__m256 length_sqr = _mm256_fmadd_ps(sub_y, sub_y, _mm256_mul_ps(sub_x, sub_x));
		__m256i in_range_mask = (__m256i)_mm256_cmp_ps(length_sqr, view_radius, _CMP_LE_OQ);
		if (mm256_is_zero(in_range_mask)) continue;

		__m256 from_dir_x = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].from].dir.x,
			boids[b2b_cmps[8*i+6].from].dir.x,
			boids[b2b_cmps[8*i+5].from].dir.x,
			boids[b2b_cmps[8*i+4].from].dir.x,
			boids[b2b_cmps[8*i+3].from].dir.x,
			boids[b2b_cmps[8*i+2].from].dir.x,
			boids[b2b_cmps[8*i+1].from].dir.x,
			boids[b2b_cmps[8*i+0].from].dir.x
		);

		__m256 from_dir_y = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].from].dir.y,
			boids[b2b_cmps[8*i+6].from].dir.y,
			boids[b2b_cmps[8*i+5].from].dir.y,
			boids[b2b_cmps[8*i+4].from].dir.y,
			boids[b2b_cmps[8*i+3].from].dir.y,
			boids[b2b_cmps[8*i+2].from].dir.y,
			boids[b2b_cmps[8*i+1].from].dir.y,
			boids[b2b_cmps[8*i+0].from].dir.y
		);

		__m256 to_dir_x = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].to].dir.x,
			boids[b2b_cmps[8*i+6].to].dir.x,
			boids[b2b_cmps[8*i+5].to].dir.x,
			boids[b2b_cmps[8*i+4].to].dir.x,
			boids[b2b_cmps[8*i+3].to].dir.x,
			boids[b2b_cmps[8*i+2].to].dir.x,
			boids[b2b_cmps[8*i+1].to].dir.x,
			boids[b2b_cmps[8*i+0].to].dir.x
		);

		__m256 to_dir_y = _mm256_set_ps(
			boids[b2b_cmps[8*i+7].to].dir.y,
			boids[b2b_cmps[8*i+6].to].dir.y,
			boids[b2b_cmps[8*i+5].to].dir.y,
			boids[b2b_cmps[8*i+4].to].dir.y,
			boids[b2b_cmps[8*i+3].to].dir.y,
			boids[b2b_cmps[8*i+2].to].dir.y,
			boids[b2b_cmps[8*i+1].to].dir.y,
			boids[b2b_cmps[8*i+0].to].dir.y
		);

		__m256 is_length_zero = (__m256)_mm256_cmpeq_epi32((__m256i)length_sqr, (__m256i)zero);
		__m256 ilength = _mm256_blendv_ps(_mm256_rsqrt_ps(length_sqr), zero, is_length_zero);

		__m256 x_norm = _mm256_mul_ps(sub_x, ilength);
		__m256 y_norm = _mm256_mul_ps(sub_y, ilength);

		__m256 x_neg_norm = _mm256_mul_ps(x_norm, negative_one);
		__m256 y_neg_norm = _mm256_mul_ps(y_norm, negative_one);

		__m256 dot_product1 = _mm256_fmadd_ps(from_dir_y, y_neg_norm, _mm256_mul_ps(from_dir_x, x_neg_norm));
		__m256 in_angle_mask1 = _mm256_cmp_ps(dot_product1, dot_threshold, _CMP_GE_OQ);
		__m256 do_append_mask1 = _mm256_and_ps(in_angle_mask1, (__m256)in_range_mask);

		__m256 dot_product2 = _mm256_fmadd_ps(to_dir_y, y_norm, _mm256_mul_ps(to_dir_x, x_norm));
		__m256 in_angle_mask2 = _mm256_cmp_ps(dot_product2, dot_threshold, _CMP_GE_OQ);
		__m256 do_append_mask2 = _mm256_and_ps(in_angle_mask2, (__m256)in_range_mask);

		_mm256_store_ps((float*)do_append_mask1_f32, do_append_mask1);
		_mm256_store_ps((float*)do_append_mask2_f32, do_append_mask2);
		for (int j = 0; j < 8; j++) {
			uboid_t cmp_idx = 8*i + j;
			if (cmp_idx >= *b2b_cmps_count) break;

			uboid_t from_boid = b2b_cmps[cmp_idx].from;
			uboid_t to_boid   = b2b_cmps[cmp_idx].to;
			if (do_append_mask1_f32[j]) {
				boid_list_append(&world->frame_arena, &local_boids[from_boid], to_boid);
				g_prof_interactions++;
			}
			if (do_append_mask2_f32[j]) {
				boid_list_append(&world->frame_arena, &local_boids[to_boid], from_boid);
				g_prof_interactions++;
			}
		}
	}

	*b2b_cmps_count = 0;
}

static void world_compute_local_boids_simd(BoidList *local_boids, World *world, ChunkGrid *chunks) {
	Boid *boids = world->boids.data();
	int boid_count = world->boids.size();
	MemoryArena *arena = &world->frame_arena;

	int b2b_padding = 8;
	int b2b_capacity = 2048*10;
	boid_pair b2b_cmps[b2b_capacity + b2b_padding];
	int b2b_cmps_count = 0;
	for (int i = 0; i < b2b_padding; i++) {
		memset(&b2b_cmps[b2b_capacity + i], 0, sizeof(boid_pair));
	}

	RPROF_START("Move chunk data to static arrays");
	uboid_t *static_chunks[chunks->width * chunks->height];
	for (int y = 0; y < chunks->height; y++) {
		for (int x = 0; x < chunks->width; x++) {
			size_t chunk_idx = chunkgrid_get_idx(chunks, x, y);
			BoidList *chunk = &chunks->data[chunk_idx];
			static_chunks[chunk_idx] = (uboid_t*)arena_malloc(arena, sizeof(uboid_t) * chunk->count);
			boid_list_to_array(static_chunks[chunk_idx], chunk);
		}
	}
	RPROF_STOP();

	RPROF_START("Calc dot products and ranges (simd)");
	for (int y = 0; y < chunks->height; y++) {
		for (int x = 0; x < chunks->width; x++) {
			size_t chunk_idx = chunkgrid_get_idx(chunks, x, y);
			BoidList *chunk = &chunks->data[chunk_idx];
			if (chunk->count == 0) continue;

			uboid_t *chunk_boids = static_chunks[chunk_idx];
			for (int i = 0; i < chunk->count-1; i++) {
				uboid_t from_boid = chunk_boids[i];
				uboid_t *to_boids = &chunk_boids[i+1];
				uboid_t to_boids_count = chunk->count-i-1;

				for (int j = 0; j < to_boids_count; j++) {
					// TODO:
					// DEBUG_ASSERT(b2b_cmps_count < b2b_capacity-1);
					b2b_cmps[b2b_cmps_count++] = {
						.from = from_boid,
						.to = to_boids[j]
					};
				}
			}

			Vector2 neighbours[] = { { 1, 0 }, { 0, 1 }, { 1, 1 }, { -1, 1 } };
			for (int neighbour_i = 0; neighbour_i < ARRAY_LEN(neighbours); neighbour_i++) {
				int chunk_y = y + neighbours[neighbour_i].y;
				if (chunk_y < 0 || chunk_y >= chunks->height) continue;
				int chunk_x = x + neighbours[neighbour_i].x;
				if (chunk_x < 0 || chunk_x >= chunks->width) continue;

				size_t neighbour_idx = chunkgrid_get_idx(chunks, chunk_x, chunk_y);
				BoidList *neighbour_chunk = &chunks->data[neighbour_idx];
				if (neighbour_chunk->count == 0) continue;

				uboid_t *neighbour_boids = static_chunks[neighbour_idx];
				for (int i = 0; i < chunk->count; i++) {
					for (int j = 0; j < neighbour_chunk->count; j++) {
						// TODO:
						// DEBUG_ASSERT(b2b_cmps_count < b2b_capacity-1);
						b2b_cmps[b2b_cmps_count++] = {
							.from = chunk_boids[i],
							.to = neighbour_boids[j]
						};
					}
				}
			}

			if (b2b_cmps_count > 2048*3) {
				world_process_local_boid_pairs(world, local_boids, b2b_cmps, &b2b_cmps_count);
			}
		}
	}

	if (b2b_cmps_count > 0) {
		world_process_local_boid_pairs(world, local_boids, b2b_cmps, &b2b_cmps_count);
	}

	RPROF_STOP();
}
#endif

static void world_compute_local_boids_scalar(BoidList *local_boids, World *world, ChunkGrid *chunks) {
	RPROF_START("Calc dot products and ranges (scalar)");
	for (int y = 0; y < chunks->height; y++) {
		for (int x = 0; x < chunks->width; x++) {
			BoidList *chunk = chunkgrid_get(chunks, x, y);
			if (chunk->count == 0) continue;

			uboid_t chunk_boids[chunk->count];
			boid_list_to_array(chunk_boids, chunk);
			for (int i = 0; i < chunk->count-1; i++) {
				uboid_t from_boid = chunk_boids[i];
				uboid_t *to_boids = &chunk_boids[i+1];
				uboid_t to_boids_count = chunk->count-i-1;
				assign_local_boids_b2l(world, local_boids, from_boid, to_boids, to_boids_count);
			}

			Vector2 neighbours[] = { { 1, 0 }, { 0, 1 }, { 1, 1 }, { -1, 1 } };
			for (int i = 0; i < ARRAY_LEN(neighbours); i++) {
				int chunk_y = y + neighbours[i].y;
				int chunk_x = x + neighbours[i].x;
				if (chunk_y < 0 || chunk_y >= chunks->height) continue;
				if (chunk_x < 0 || chunk_x >= chunks->width) continue;

				BoidList *neighbour_chunk = chunkgrid_get(chunks, chunk_x, chunk_y);
				if (neighbour_chunk->count == 0) continue;

				uboid_t neighbour_ids[neighbour_chunk->count];
				boid_list_to_array(neighbour_ids, neighbour_chunk);

				uboid_t boid1;
				BoidsListNodeIterator it1 = boid_list_get_iterator(chunk);
				while (boid_list_iterator_next(&it1, &boid1)) {
					assign_local_boids_b2l(world, local_boids, boid1, neighbour_ids, neighbour_chunk->count);
				}
			}
		}
	}
	RPROF_STOP();
}

static BoidList* world_compute_local_boids(World *world) {
	Boid *boids = world->boids.data();
	int boid_count = world->boids.size();
	MemoryArena *arena = &world->frame_arena;

	RPROF_START("Alloc groups");
	BoidList *all_local_boids = (BoidList*)arena_malloc(arena, boid_count * sizeof(BoidList));
	for (int i = 0; i < boid_count; i++) {
		boid_list_init(&all_local_boids[i]);
	}
	RPROF_STOP();

	float chunk_size = std::max(world->view_radius, 1.0f);
	int chunks_wide = std::ceil(world->size.x / chunk_size) + 1;
	int chunks_high = std::ceil(world->size.y / chunk_size) + 1;

	RPROF_START("Alloc chunks");
	ChunkGrid chunks;
	chunkgrid_init(arena, &chunks, chunks_wide, chunks_high);
	RPROF_STOP();

	RPROF_START("Assign boids to chunks");
	for (int i = 0; i < boid_count; i++) {
		Boid *boid = &boids[i];
		int x = boid->pos.x / chunk_size;
		int y = boid->pos.y / chunk_size;

		boid_list_append(arena, chunkgrid_get(&chunks, x, y), i);
	}
	RPROF_STOP();

	RPROF_START("world_compute_local_boids()");
#ifdef __EMSCRIPTEN__
	// TODO: Rewrite simd version to only use SSE, not AVX2
	world_compute_local_boids_scalar(all_local_boids, world, &chunks);
#else
	world_compute_local_boids_simd(all_local_boids, world, &chunks);
#endif

	RPROF_STOP();

	return all_local_boids;
}

static void world_update(World *world, float dt) {
	if (world->freeze) return;

	arena_clear(&world->frame_arena);

	Boid *boids = world->boids.data();
	int boid_count = world->boids.size();
	assert(boid_count <= MAX_BOIDS);

	BoidList *list_of_local_boids = world_compute_local_boids(world);

	RPROF_START("Apply forces");
	int avoidance_ray_count = world->collision_avoidance_ray_count * 2 + 1;
	Vector2 avoidance_ray_dirs[avoidance_ray_count];
	fill_avoidance_ray_dirs(avoidance_ray_dirs, avoidance_ray_count, world->collision_avoidance_ray_angle);

	for (int i = 0; i < boid_count; i++) {
		Boid *boid = &world->boids[i];
		Vector2 acc = { 0, 0 };

		BoidList *local_boids = &list_of_local_boids[i];

		if (local_boids->count > 0) {
			Vector2 separation_force = { 0, 0 };
			Vector2 flock_center  = { 0, 0 };
			Vector2 flock_heading = { 0, 0 };

			uboid_t local_boid_id;
			BoidsListNodeIterator it = boid_list_get_iterator(local_boids);
			while (boid_list_iterator_next(&it, &local_boid_id)) {
				Boid *local_boid = &boids[local_boid_id];
				flock_heading = Vector2Add(flock_heading, local_boid->dir);
				flock_center  = Vector2Add(flock_center , local_boid->pos);

				Vector2 pos_diff = Vector2Subtract(boid->pos, local_boid->pos);
				float dist_sqr = Vector2LengthSqr(pos_diff);
				if (dist_sqr <= world->separation_radius * world->separation_radius) {
					separation_force = Vector2Add(separation_force, vector2_div_value(pos_diff, dist_sqr));
				}
			}
			flock_center = vector2_div_value(flock_center, local_boids->count);

			Vector2 alignment_force = Vector2Normalize(flock_heading);
			acc = Vector2Add(acc, vector2_mul_value(alignment_force, world->alignment_strength));

			Vector2 cohesion_force = Vector2Normalize(Vector2Subtract(flock_center, boid->pos));
			acc = Vector2Add(acc, vector2_mul_value(cohesion_force, world->cohesion_strength));

			separation_force = Vector2Normalize(separation_force);
			acc = Vector2Add(acc, vector2_mul_value(separation_force, world->separation_strength));
		}

		// Apply obstacle avoidance to accelaration
		Vector2 collision_avoidance = get_collision_avoidance_dir(world, boid, avoidance_ray_dirs, avoidance_ray_count);
		acc = Vector2Add(acc, vector2_mul_value(collision_avoidance, world->collision_avoidance_strength));

		acc = vector2_mul_value(acc, world->max_speed);

		// Clamp accelaration
		Vector2 clamped_acc = acc;
		float acc_size = Vector2Length(acc);
		if (acc_size > world->max_steer_speed) {
			clamped_acc = vector2_mul_value(Vector2Normalize(acc), world->max_steer_speed);
		}

		// Apply accelaration
		Vector2 velocity = Vector2Multiply(boid->dir, { boid->speed, boid->speed });
		velocity = Vector2Add(velocity, vector2_mul_value(clamped_acc, dt));

		boid->dir = Vector2Normalize(velocity);
		boid->speed = Vector2Length(velocity);

		boid->speed = Clamp(boid->speed, world->min_speed, world->max_speed);
		Vector2 step = vector2_mul_value(boid->dir, boid->speed * dt);
		Vector2 target_pos = Vector2Add(boid->pos, step);

		// Check collisions
		RayHitResult hit_result;
		get_intersect_with_world(&hit_result, target_pos, step, world);
		if (hit_result.hit == -1 || hit_result.hit > 2) {
			boid->pos = target_pos;
		}

		if (world->looping_walls) {
			if (boid->pos.x >= world->size.x) {
				boid->pos.x -= world->size.x;
			} else if (boid->pos.x < 0) {
				boid->pos.x += world->size.x;
			}
			if (boid->pos.y >= world->size.y) {
				boid->pos.y -= world->size.y;
			} else if (boid->pos.y < 0) {
				boid->pos.y += world->size.y;
			}
		} else {
			if (boid->pos.x >= world->size.x) {
				boid->pos.x = world->size.x-1;
			} else if (boid->pos.x < 0) {
				boid->pos.x = 0;
			}
			if (boid->pos.y >= world->size.y) {
				boid->pos.y = world->size.y-1;
			} else if (boid->pos.y < 0) {
				boid->pos.y = 0;
			}
		}
	}
	RPROF_STOP();
}

// --------------------- Draw ------------------------

static void draw_obstacle(Obstacle *obstacle, Color color) {
	std::vector<Vector2> *points = &obstacle->points;
	int point_count = points->size();

	rlBegin(RL_TRIANGLES);
	{
		rlColor4ub(color.r, color.g, color.b, color.a);
		for (int j = 0; j < point_count-1; j++) {
			Vector2 *point1 = &(*points)[j];
			Vector2 *point2 = &(*points)[j+1];
			rlVertex2f(point1->x, point1->y);
			rlVertex2f(obstacle->center.x, obstacle->center.y);
			rlVertex2f(point2->x, point2->y);
		}

		rlVertex2f((*points)[point_count-1].x, (*points)[point_count-1].y);
		rlVertex2f(obstacle->center.x, obstacle->center.y);
		rlVertex2f((*points)[0].x, (*points)[0].y);
	}
	rlEnd();
}

static void draw_obstacle_avoidance_rays(Visuals *visuals, World *world, Boid *boid) {
	Vector2 pos = boid->pos;

	int ray_count = world->collision_avoidance_ray_count * 2 + 1;
	float ray_angles[ray_count];
	fill_avoidance_ray_angles(ray_angles, ray_count, world->collision_avoidance_ray_angle);

	float facing = std::atan2(boid->dir.y, boid->dir.x);
	for (int i = 0; i < ray_count; i++) {
		Vector2 ray_dir = {
			std::cos(facing + ray_angles[i]),
			std::sin(facing + ray_angles[i])
		};

		RayHitResult hit_result;
		get_intersect_with_world(&hit_result, pos, ray_dir, world);
		bool hit_obstacle = (hit_result.hit != -1 && hit_result.hit <= world->collision_avoidance_distance);

		Color ray_color = GREEN;
		float ray_length = world->collision_avoidance_distance;
		if (hit_obstacle) {
			ray_length = hit_result.hit;
			ray_color = BLUE;
		}

		Vector2 hit_pos = Vector2Add(pos, Vector2Multiply(ray_dir, { ray_length, ray_length }));
		DrawLine(pos.x, pos.y, hit_pos.x, hit_pos.y, ray_color);
		if (hit_obstacle) {
			DrawCircle(hit_pos.x, hit_pos.y, visuals->boid_edge_size * 0.05, ray_color);
		}
	}
}

static void draw_circle_sector(Vector2 center, float radius, float start_angle, float end_angle, int segments, Color color) {
	rlBegin(RL_TRIANGLES);
	float angle_step = (end_angle - start_angle) / segments;
	for (int i = 0; i < segments; i++)
	{
		rlColor4ub(color.r, color.g, color.b, color.a);
		float angle = start_angle + i * angle_step;
		float nextAngle = start_angle + (i+1) * angle_step;

		rlVertex2f(center.x, center.y);
		rlVertex2f(center.x + cosf(nextAngle)*radius, center.y + sinf(nextAngle)*radius);
		rlVertex2f(center.x + cosf(angle)    *radius, center.y + sinf(angle)    *radius);
	}
	rlEnd();
}

static void draw_boids(World *world, Visuals *visuals) {
	int boid_count = world->boids.size();

	float boid_length = visuals->boid_edge_size * std::sqrt(3)/2;
	float boid_width = visuals->boid_edge_size * 0.6;
	Color color = visuals->boid_color;

	rlBegin(RL_TRIANGLES);
	for (int i = 0; i < boid_count; i++) {
		Boid *boid = &world->boids[i];

		Vector2 triangle[] = {
			{  boid_length*2/3.0f, 0 },
			{ -boid_length*1/3.0f, -boid_width/2 },
			{ -boid_length*1/3.0f, boid_width/2 },
		};
		float facing = std::atan2(boid->dir.y, boid->dir.x);
		float facing_cos = cos(facing);
		float facing_sin = sin(facing);
		for (int i = 0; i < 3; i++) {
			Vector2 new_pos = boid->pos;
			new_pos.x += triangle[i].x * facing_cos - triangle[i].y * facing_sin;
			new_pos.y += triangle[i].x * facing_sin + triangle[i].y * facing_cos;
			triangle[i] = new_pos;
		}

		rlColor4ub(color.r, color.g, color.b, color.a);
		rlVertex2f(triangle[0].x, triangle[0].y);
		rlVertex2f(triangle[1].x, triangle[1].y);
		rlVertex2f(triangle[2].x, triangle[2].y);
	}

	rlEnd();
}

static void world_draw(World *world, Visuals *visuals) {
	for (int i = 0; i < world->obstacles.size(); i++) {
		draw_obstacle(&world->obstacles[i], GRAY);
	}

	if (visuals->draw_view_cone) {
		Color view_cone_color = Fade(GRAY, 0.4);
		for (int i = 0; i < world->boids.size(); i++) {
			Boid *boid = &world->boids[i];
			Vector2 pos = boid->pos;
			float facing = std::atan2(boid->dir.y, boid->dir.x);

			float view_angle = world->view_angle;
			float segments = 16;

			draw_circle_sector(pos, world->view_radius, facing - view_angle/2, facing + view_angle/2, segments, view_cone_color);
		}
	}

	int boid_count = world->boids.size();
	for (int i = 0; i < boid_count; i++) {
		Boid *boid = &world->boids[i];
		if (visuals->draw_collision_avoidance_rays) {
			draw_obstacle_avoidance_rays(visuals, world, boid);
		}
		if (visuals->draw_separation_radius) {
			DrawCircleLines(boid->pos.x, boid->pos.y, world->separation_radius, MAGENTA);
		}
	}

	draw_boids(world, visuals);

	for (int i = 0; i < boid_count; i++) {
		Boid *boid = &world->boids[i];
		if (visuals->draw_boid_direction) {
			DrawCircle(boid->pos.x, boid->pos.y, visuals->boid_edge_size * 0.05, RED);
			Vector2 look_pos = Vector2Add(boid->pos, vector2_mul_value(boid->dir, visuals->boid_edge_size*1.5));
			DrawLine(boid->pos.x, boid->pos.y, look_pos.x, look_pos.y, RED);
		}
	}
}