package physics

import "base:runtime"
import "bvh"
import "collision"
import "common:name"
import "core:container/bit_array"
import "core:fmt"
import "core:log"
import "core:math"
import lg "core:math/linalg"
import "core:math/rand"
import "core:slice"
import "game:assets"
import "game:debug"
import he "game:halfedge"
import "libs:tracy"

_ :: name
_ :: log
_ :: rand
_ :: math
_ :: fmt
_ :: slice
_ :: he
_ :: debug

Solver_Config :: struct {
	// Will automatically do fixed timestep
	timestep:            f32,
	gravity:             Vec3,
	substreps_minus_one: i32,
}

Solver_State :: struct {
	accumulated_time:                 f32,
	// Incremented when simulate is called (not simulate_step)
	simulation_frame:                 u32,

	// Number of immediate bodies referenced this frame
	immediate_bodies:                 Immediate_Container(Body_Handle),
	immediate_suspension_constraints: Immediate_Container(Suspension_Constraint_Handle),
	immediate_engines:                Immediate_Container(Engine_Handle),
	immediate_level_geoms:            Immediate_Container(Level_Geom_Handle),
}

copy_solver_state :: proc(dst, src: ^Solver_State) {
	dst.accumulated_time = src.accumulated_time
	dst.simulation_frame = src.simulation_frame

	immediate_container_copy(&dst.immediate_bodies, &src.immediate_bodies)
	immediate_container_copy(
		&dst.immediate_suspension_constraints,
		&src.immediate_suspension_constraints,
	)
	immediate_container_copy(&dst.immediate_engines, &src.immediate_engines)
	immediate_container_copy(&dst.immediate_level_geoms, &src.immediate_level_geoms)
}

destroy_solver_state :: proc(state: ^Solver_State) {
	immediate_container_destroy(&state.immediate_bodies)
	immediate_container_destroy(&state.immediate_suspension_constraints)
	immediate_container_destroy(&state.immediate_engines)
	immediate_container_destroy(&state.immediate_level_geoms)
}

MAX_STEPS :: 10

Step_Mode :: enum {
	Accumulated_Time,
	Single,
}

Static_TLAS :: struct {
	bvh_tree:         bvh.BVH,
	level_geom_aabbs: []bvh.AABB,
	built:            bool,
}

// Top Level Acceleration Structure
Dynamic_TLAS :: struct {
	bvh_tree:   bvh.BVH,
	body_aabbs: []bvh.AABB,
	built:      bool,
}

build_static_tlas :: proc(sim_state: ^Sim_State, out_tlas: ^Static_TLAS) {
	tracy.Zone()

	static_tlas_destroy(out_tlas)

	level_geom_aabbs := make([]bvh.AABB, len(sim_state.level_geoms))
	level_geom_indices := make([]u16, len(sim_state.level_geoms), context.temp_allocator)

	{
		for i in 0 ..< len(sim_state.level_geoms) {
			level_geom := &sim_state.level_geoms[i]

			if level_geom.alive {
				aabb := &level_geom_aabbs[i]
				level_geom_indices[i] = u16(i)

				aabb.min = level_geom.aabb.center - level_geom.aabb.extent
				aabb.max = level_geom.aabb.center + level_geom.aabb.extent
			}
		}
	}

	sim_state_bvh := bvh.build_bvh_from_aabbs(level_geom_aabbs, level_geom_indices)

	out_tlas^ = Static_TLAS {
		built            = true,
		bvh_tree         = sim_state_bvh,
		level_geom_aabbs = level_geom_aabbs,
	}
}

static_tlas_destroy :: proc(static_tlas: ^Static_TLAS) {
	if static_tlas.built {
		bvh.destroy_bvh(&static_tlas.bvh_tree)
		delete(static_tlas.level_geom_aabbs)
		static_tlas^ = {}
	}
}

// TODO: free intermediate temp allocs
// Creates TLAS using temp allocator
build_dynamic_tlas :: proc(
	sim_state: ^Sim_State,
	config: Solver_Config,
	out_tlas: ^Dynamic_TLAS,
	allocator := context.allocator,
) {
	tracy.Zone()

	dynamic_tlas_destroy(out_tlas)

	body_aabbs := make([]bvh.AABB, len(sim_state.bodies_slice), allocator)
	body_indices := make([]u16, len(sim_state.bodies_slice), context.temp_allocator)

	{
		for i in 0 ..< len(sim_state.bodies_slice) {
			body := &sim_state.bodies_slice[i]

			if body.alive {
				aabb := &body_aabbs[i]
				body_indices[i] = u16(i)

				phys_aabb := body_get_aabb(body)

				EXPAND_K :: 10
				expand := lg.abs(EXPAND_K * config.timestep * body.v) + 0.1
				phys_aabb.extent += expand * 0.5

				aabb.min = phys_aabb.center - phys_aabb.extent
				aabb.max = phys_aabb.center + phys_aabb.extent
			}
		}
	}

	sim_state_bvh := bvh.build_bvh_from_aabbs(body_aabbs, body_indices, allocator)

	out_tlas^ = Dynamic_TLAS {
		bvh_tree   = sim_state_bvh,
		body_aabbs = body_aabbs,
		built      = true,
	}
}

dynamic_tlas_destroy :: proc(dyn_tlas: ^Dynamic_TLAS) {
	if dyn_tlas.built {
		bvh.destroy_bvh(&dyn_tlas.bvh_tree)
		delete(dyn_tlas.body_aabbs)
		dyn_tlas^ = {}
	}
}


raycasts_level :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	tlas: ^Static_TLAS,
	origin, dir: Vec3,
	distance := max(f32),
) -> (
	t: f32,
	normal: Vec3,
	hit: bool,
) {
	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	t = distance

	ray: bvh.Ray
	ray.origin = origin
	ray.dir = dir
	ray.dir_inv = 1.0 / dir
	it := bvh.iterator_intersect_leaf_ray(&tlas.bvh_tree, ray, distance)

	for leaf_node in bvh.iterator_intersect_leaf_next(&it) {
		for j in 0 ..< leaf_node.prim_len {
			level_geom_handle := index_to_level_geom(
				int(tlas.bvh_tree.primitives[leaf_node.child_or_prim_start + j]),
			)
			level_geom := get_level_geom(sim_state, level_geom_handle)
			blas, vertices, indices := get_level_geom_data(sim_state, sim_cache, level_geom_handle)

			// TODO: transform ray into blas space and back

			inv_q := lg.quaternion_inverse(level_geom.q)
			local_ray := ray
			local_ray.origin = lg.quaternion_mul_vector3(inv_q, ray.origin - level_geom.x)
			local_ray.dir = lg.quaternion_mul_vector3(inv_q, ray.dir)
			local_ray.dir_inv = 1.0 / local_ray.dir

			blas_it := bvh.iterator_intersect_leaf_ray(&blas, local_ray, distance)
			for blas_leaf_node in bvh.iterator_intersect_leaf_next(&blas_it) {
				for k in 0 ..< blas_leaf_node.prim_len {
					tri_idx := int(blas.primitives[blas_leaf_node.child_or_prim_start + k])

					tri := get_triangle(vertices, indices, tri_idx)

					hit_t, tmp_normal, _, ok := collision.intersect_ray_triangle(
						{local_ray.origin, local_ray.origin + local_ray.dir},
						tri,
					)

					if ok && (!hit || hit_t < t) {
						t = hit_t
						normal = lg.quaternion_mul_vector3(level_geom.q, tmp_normal)
						hit = true
					}
				}
			}
		}
	}

	return
}

raycast_bodies :: proc(
	sim_state: ^Sim_State,
	tlas: ^Dynamic_TLAS,
	origin, dir: Vec3,
	distance := max(f32),
) -> (
	t: f32,
	normal: Vec3,
	body_handle: Body_Handle,
	hit: bool,
) {
	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	t = distance

	ray: bvh.Ray
	ray.origin = origin
	ray.dir = dir
	ray.dir_inv = 1.0 / dir
	it := bvh.iterator_intersect_leaf_ray(&tlas.bvh_tree, ray, distance)

	for leaf_node in bvh.iterator_intersect_leaf_next(&it) {
		for j in 0 ..< leaf_node.prim_len {
			body_idx := tlas.bvh_tree.primitives[leaf_node.child_or_prim_start + j]
			tmp_body_handle := index_to_body_handle(int(body_idx))

			body := get_body(sim_state, tmp_body_handle)

			it := shapes_iterator(sim_state, body.shape)
			for shape in shapes_iterator_next(&it) {
				mesh := body_get_convex_shape_world(
					sim_state,
					body,
					shape^,
					context.temp_allocator,
				)

				hit_t, _, tmp_normal, _, ok := collision.ray_vs_convex(
					mesh,
					ray.origin,
					ray.dir,
					distance,
				)

				if ok && (!hit || hit_t < t) {
					t = hit_t
					normal = tmp_normal
					body_handle = tmp_body_handle
					hit = true
				}
			}
		}
	}

	if hit {
		normal = lg.normalize0(normal)
	}

	return
}

raycast :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	static_tlas: ^Static_TLAS,
	dyn_tlas: ^Dynamic_TLAS,
	origin, dir: Vec3,
	distance := max(f32),
) -> (
	t: f32,
	normal: Vec3,
	hit_body: Body_Handle,
	hit: bool,
) {
	t = distance

	t1, normal1, hit1 := raycasts_level(sim_state, sim_cache, static_tlas, origin, dir, t)
	t2, normal2, body, hit2 := raycast_bodies(sim_state, dyn_tlas, origin, dir, t)

	hit = hit1 || hit2

	if hit1 && hit2 {
		if t1 < t2 {
			t = t1
			normal = normal1
		} else {
			t = t2
			normal = normal2
			hit_body = body
		}
	} else if hit1 {
		t = t1
		normal = normal1
	} else if hit2 {
		t = t2
		normal = normal2
		hit_body = body
	}
	if hit {
		normal = lg.normalize0(normal)
	}

	return
}

// Cache used during simulation to avoid complex computations
Sim_Cache :: struct {
	// Looking up bvh can be expensive because assetman touches the filesystem each time. We assume that during simulation it cannot change
	// so it's safe to cache
	level_geom_asset_bvh: map[Level_Geom_Handle]assets.Loaded_BVH,
}


get_triangle :: proc(vertices: []Vec3, indices: []u16, tri_idx: int) -> (tri: [3]Vec3) {
	i1, i2, i3 := indices[tri_idx * 3 + 0], indices[tri_idx * 3 + 1], indices[tri_idx * 3 + 2]
	tri[0], tri[1], tri[2] = vertices[i1], vertices[i2], vertices[i3]
	return
}

get_transformed_triangle :: proc(
	vertices: []Vec3,
	indices: []u16,
	tri_idx: int,
	position: Vec3,
	rotation: Quat,
) -> (
	tri: [3]Vec3,
) {
	tri = get_triangle(vertices, indices, tri_idx)

	tri[0] = lg.quaternion_mul_vector3(rotation, tri[0]) + position
	tri[1] = lg.quaternion_mul_vector3(rotation, tri[1]) + position
	tri[2] = lg.quaternion_mul_vector3(rotation, tri[2]) + position
	return
}

get_triangle_aabb :: proc(tri: [3]Vec3) -> (aabb: bvh.AABB) {
	aabb.min = lg.min(tri[0], lg.min(tri[1], tri[2]))
	aabb.max = lg.max(tri[0], lg.max(tri[1], tri[2]))

	return
}

remove_invalid_contacts :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	static_tlas: Static_TLAS,
	dyn_tlas: Dynamic_TLAS,
) {
	tracy.Zone()

	i := 0
	for i < len(sim_state.contact_container.contacts) {
		contact := sim_state.contact_container.contacts[i]

		should_remove := false

		if contact.type == .Body_vs_Body {
			should_remove |= !get_body(sim_state, contact.a).alive
			should_remove |= !get_body(sim_state, Body_Handle(contact.b)).alive

			if !should_remove {
				aabb_a := dyn_tlas.body_aabbs[body_handle_to_index(contact.a)]
				aabb_b := dyn_tlas.body_aabbs[body_handle_to_index(Body_Handle(contact.b))]

				should_remove |= !bvh.test_aabb_vs_aabb(aabb_a, aabb_b)
			}
		} else {
			// a is a body, b is a triangle index
			should_remove |= !get_body(sim_state, contact.a).alive

			if !should_remove {
				level_geom_handle := Level_Geom_Handle(contact.b)
				level_geom := get_level_geom(sim_state, level_geom_handle)

				should_remove |= !level_geom.alive

				if !should_remove {
					tri_idx := int(contact.local_tri_idx)
					_, vertices, indices := get_level_geom_data(
						sim_state,
						sim_cache,
						level_geom_handle,
					)
					should_remove |= tri_idx * 3 >= len(indices)

					if !should_remove {
						aabb_a := dyn_tlas.body_aabbs[body_handle_to_index(contact.a)]
						tri := get_transformed_triangle(
							vertices,
							indices,
							tri_idx,
							level_geom.x,
							level_geom.q,
						)
						aabb_b := get_triangle_aabb(tri)

						should_remove |= !bvh.test_aabb_vs_aabb(aabb_a, aabb_b)
					}
				}
			}
		}

		pair_from_contact :: proc(contact: Contact) -> (pair: Contact_Pair) {
			if contact.type == .Body_vs_Body {
				pair = make_contact_pair_bodies(
					body_handle_to_index(contact.a),
					body_handle_to_index(Body_Handle(contact.b)),
					contact.shape_a,
					contact.shape_b,
				)
			} else {
				pair = make_contact_pair_body_level(
					Body_Handle(contact.a),
					Level_Geom_Handle(contact.b),
					contact.shape_a,
					contact.local_tri_idx,
				)
			}
			return
		}

		if should_remove {
			removed_pair := pair_from_contact(contact)
			delete_key(&sim_state.contact_container.lookup, removed_pair)

			unordered_remove_soa(&sim_state.contact_container.contacts, i)

			if i < len(sim_state.contact_container.contacts) {
				moved_contact := &sim_state.contact_container.contacts[i]
				moved_pair := pair_from_contact(moved_contact^)
				sim_state.contact_container.lookup[moved_pair] = i32(i)
			}
		} else {
			i += 1
		}
	}
}

// TODO: free intermediate temp allocs
find_new_contacts :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	static_tlas: ^Static_TLAS,
	dyn_tlas: ^Dynamic_TLAS,
	config: Solver_Config,
) {
	tracy.Zone()

	for i in 0 ..< len(sim_state.bodies_slice) {
		assert(i <= int(max(u16)))
		body_idx := u16(i)
		body := &sim_state.bodies_slice[i]

		if body.alive && body.inv_mass != 0 {
			tracy.ZoneN(name.to_string(body.name))

			body_aabb := dyn_tlas.body_aabbs[i]
			{
				tracy.ZoneN("find_new_contacts::dynamic_vs_dynamic")
				it := bvh.iterator_intersect_leaf_aabb(&dyn_tlas.bvh_tree, body_aabb)

				for leaf_node in bvh.iterator_intersect_leaf_next(&it) {
					for j in 0 ..< leaf_node.prim_len {
						other_body_idx :=
							dyn_tlas.bvh_tree.primitives[leaf_node.child_or_prim_start + j]
						other_body := &sim_state.bodies_slice[other_body_idx]
						prim_aabb := dyn_tlas.body_aabbs[other_body_idx]

						if body_idx != other_body_idx &&
						   (bvh.test_aabb_vs_aabb(body_aabb, prim_aabb)) {
							shapes_a_it := shapes_iterator(sim_state, body.shape)
							for shape_a in shapes_iterator_next(&shapes_a_it) {
								shape_a_idx := shapes_a_it.counter - 1

								shape_a_aabb := shape_get_aabb(shape_a^)
								shape_a_aabb = body_transform_shape_aabb(body, shape_a_aabb)
								EXPAND_K :: 2
								expand := lg.abs(EXPAND_K * config.timestep * body.v) + 0.1
								shape_a_aabb.extent += expand * 0.5

								shapes_b_it := shapes_iterator(sim_state, other_body.shape)

								for shape_b in shapes_iterator_next(&shapes_b_it) {
									shape_b_idx := shapes_b_it.counter - 1

									shape_b_aabb := shape_get_aabb(shape_b^)
									shape_b_aabb = body_transform_shape_aabb(
										other_body,
										shape_b_aabb,
									)

									pair := make_contact_pair_bodies(
										int(body_idx),
										int(other_body_idx),
										shape_a_idx,
										shape_b_idx,
									)

									bvh_aabb_a := bvh.AABB {
										min = shape_a_aabb.center - shape_a_aabb.extent,
										max = shape_a_aabb.center + shape_a_aabb.extent,
									}
									bvh_aabb_b := bvh.AABB {
										min = shape_b_aabb.center - shape_b_aabb.extent,
										max = shape_b_aabb.center + shape_b_aabb.extent,
									}

									if bvh.test_aabb_vs_aabb(bvh_aabb_a, bvh_aabb_b) &&
									   !(pair in sim_state.contact_container.lookup) {
										new_contact_idx := len(
											sim_state.contact_container.contacts,
										)
										resize_soa(
											&sim_state.contact_container.contacts,
											new_contact_idx + 1,
										)
										contact := &sim_state.contact_container.contacts[new_contact_idx]

										contact^ = Contact {
											type    = .Body_vs_Body,
											a       = Body_Handle(i + 1),
											b       = i32(other_body_idx + 1),
											shape_a = shape_a_idx,
											shape_b = shape_b_idx,
										}

										sim_state.contact_container.lookup[pair] = i32(
											new_contact_idx,
										)
									}
								}
							}
						}
					}
				}
			}

			{
				tracy_ctx := tracy.ZoneN("find_new_contacts::dynamic_vs_static")

				it := bvh.iterator_intersect_leaf_aabb(&static_tlas.bvh_tree, body_aabb)

				num_contacts_found := 0
				for leaf_node in bvh.iterator_intersect_leaf_next(&it) {
					tracy.ZoneN("intersect_bvh_node")
					for j in 0 ..< leaf_node.prim_len {
						level_geom_handle := index_to_level_geom(
							int(
								static_tlas.bvh_tree.primitives[leaf_node.child_or_prim_start + j],
							),
						)
						level_geom := get_level_geom(sim_state, level_geom_handle)
						if level_geom.alive {
							blas, vertices, indices := get_level_geom_data(
								sim_state,
								sim_cache,
								level_geom_handle,
							)
							body_aabb_in_Level_geom_space := aabb_inv_transform(
								body_aabb,
								level_geom.x,
								level_geom.q,
							)
							blas_it := bvh.iterator_intersect_leaf_aabb(
								&blas,
								body_aabb_in_Level_geom_space,
							)
							for blas_leaf_node in bvh.iterator_intersect_leaf_next(&blas_it) {
								for k in 0 ..< blas_leaf_node.prim_len {
									tri_idx := int(
										blas.primitives[blas_leaf_node.child_or_prim_start + k],
									)
									tri := get_triangle(vertices, indices, tri_idx)
									prim_aabb := get_triangle_aabb(tri)
									prim_aabb.min -= 0.1
									prim_aabb.max += 0.1

									if bvh.test_aabb_vs_aabb(
										body_aabb_in_Level_geom_space,
										prim_aabb,
									) {
										// tracy.ZoneN("body_vs_triangle", false)

										shapes_it := shapes_iterator(sim_state, body.shape)
										for shape in shapes_iterator_next(&shapes_it) {
											// tracy.ZoneN("body_shape_vs_triangle")

											shape_idx := shapes_it.counter - 1
											shape_aabb := shape_get_aabb(shape^)
											shape_aabb = body_transform_shape_aabb(
												body,
												shape_aabb,
											)
											bvh_shape_aabb := bvh.AABB {
												min = shape_aabb.center - shape_aabb.extent,
												max = shape_aabb.center + shape_aabb.extent,
											}
											bvh_shape_aabb = aabb_inv_transform(
												bvh_shape_aabb,
												level_geom.x,
												level_geom.q,
											)

											pair := make_contact_pair_body_level(
												index_to_body_handle(int(body_idx)),
												level_geom_handle,
												shape_idx,
												i32(tri_idx),
											)

											if bvh.test_aabb_vs_aabb(prim_aabb, bvh_shape_aabb) &&
											   !(pair in sim_state.contact_container.lookup) {
												num_contacts_found += 1
												new_contact_idx := len(
													sim_state.contact_container.contacts,
												)
												resize_soa(
													&sim_state.contact_container.contacts,
													new_contact_idx + 1,
												)
												contact := &sim_state.contact_container.contacts[new_contact_idx]

												contact^ = Contact {
													type          = .Body_vs_Level,
													a             = index_to_body_handle(i),
													shape_a       = shape_idx,
													b             = i32(level_geom_handle),
													local_tri_idx = i32(tri_idx),
												}

												sim_state.contact_container.lookup[pair] = i32(
													new_contact_idx,
												)
											}
										}
									}
								}
							}
						}
					}
				}

				tracy.ZoneValue(tracy_ctx, u64(num_contacts_found))
			}
		}
	}
}

// Outer simulation loop for fixed timestepping
simulate :: proc(
	scene: ^Scene,
	config: Solver_Config,
	dt: f32,
	commit := true, // commit = false is a special mode for debugging physics stepping to allow rerunning the same step each frame
	step_mode := Step_Mode.Accumulated_Time,
) -> (
	accumulated_dt: f32,
) {
	tracy.Zone()
	assert(config.timestep > 0)

	state := &scene.solver_state

	prune_immediate(scene)

	sim_state_interp_cache_clear(&scene.interpolation_cache)

	did_copy := false
	sim_state := get_next_sim_state(scene)

	sim_cache: Sim_Cache
	sim_cache.level_geom_asset_bvh = make_map(
		map[Level_Geom_Handle]assets.Loaded_BVH,
		context.temp_allocator,
	)

	simulated_dt := f32(0)

	num_steps := 0
	switch step_mode {
	case .Accumulated_Time:
		state.accumulated_time += dt

		for state.accumulated_time >= config.timestep {
			if !did_copy {
				did_copy = true
				copy_sim_state(get_next_sim_state(scene), get_sim_state(scene))
			}
			num_steps += 1
			state.accumulated_time -= config.timestep
			simulated_dt += config.timestep

			if num_steps < MAX_STEPS {
				simulate_step(scene, sim_state, &sim_cache, config)
			}
		}
	case .Single:
		copy_sim_state(get_next_sim_state(scene), get_sim_state(scene))
		simulate_step(scene, get_next_sim_state(scene), &sim_cache, config)
		num_steps += 1
	}

	if num_steps > 0 && commit {
		flip_sim_state(scene)
	}

	state.simulation_frame += 1
	state.immediate_bodies.num_items = 0
	state.immediate_suspension_constraints.num_items = 0
	state.immediate_engines.num_items = 0
	state.immediate_level_geoms.num_items = 0

	return simulated_dt
}

GLOBAL_PLANE :: collision.Plane {
	normal = Vec3{0, 1, 0},
	dist   = 0,
}

Contact_Type :: enum {
	Body_vs_Body,
	Body_vs_Level,
}

Contact :: struct {
	type:                      Contact_Type,
	a:                         Body_Handle,
	// Body_vs_Body - Body_Handle, Body_vs_Level - Level_Geom_Handle
	b:                         i32,
	// shape index in each body
	shape_a, shape_b:          i32,
	local_tri_idx:             i32,
	manifold:                  collision.Contact_Manifold,
	applied_corrections:       int,

	// For PGS
	total_normal_impulse:      [4]f32,
	// xy - tangent and bitangent (linear friction), z - twist friction around normal (rotational impulse only)
	total_friction_impulse:    [4]Vec2,

	// XPBD stuff
	lambda_normal:             [4]f32,
	lambda_tangent:            f32,
	applied_static_friction:   bool,
	applied_normal_correction: [4]f32,
}

update_contacts :: proc(sim_state: ^Sim_State, sim_cache: ^Sim_Cache, static_tlas: ^Static_TLAS) {
	tracy.Zone()

	temp := runtime.default_temp_allocator_temp_begin()
	defer runtime.default_temp_allocator_temp_end(temp)

	graph_color_bitmask: [4]bit_array.Bit_Array
	for i in 0 ..< len(graph_color_bitmask) {
		bit_array.init(
			&graph_color_bitmask[i],
			len(sim_state.bodies_slice),
			0,
			context.temp_allocator,
		)
	}

	Body_And_Shape_Index :: struct {
		body_index:  i32,
		shape_index: i32,
	}

	body_mesh_cache := make(map[Body_And_Shape_Index]collision.Convex, context.temp_allocator)

	get_body_mesh :: #force_inline proc(
		body_mesh_cache: ^map[Body_And_Shape_Index]collision.Convex,
		sim_state: ^Sim_State,
		body_idx: int,
		shape: i32,
	) -> (
		convex: collision.Convex,
	) {
		key := Body_And_Shape_Index {
			body_index  = i32(body_idx),
			shape_index = shape,
		}
		ok: bool
		convex, ok = body_mesh_cache[key]
		if !ok {
			convex = body_get_convex_shape_by_linear_index_world(
				sim_state,
				&sim_state.bodies_slice[body_idx],
				shape,
			)
			body_mesh_cache[key] = convex
		}
		return convex
	}

	for contact_idx in 0 ..< len(sim_state.contact_container.contacts) {
		contact := &sim_state.contact_container.contacts[contact_idx]

		b_handle := Body_Handle(contact.b) if contact.type == .Body_vs_Body else INVALID_BODY
		body, body2 := get_body(sim_state, contact.a), get_body(sim_state, b_handle)

		assert(body.alive)

		old_manifold := contact.manifold
		old_total_normal_impulse := contact.total_normal_impulse
		old_total_friction_impulse := contact.total_friction_impulse

		contact.manifold = {}
		contact.total_normal_impulse = 0
		contact.total_friction_impulse = 0
		contact.lambda_normal = 0
		contact.lambda_tangent = 0
		contact.applied_static_friction = false
		contact.applied_normal_correction = 0

		// aabb1 := body_get_aabb(body)
		// aabb2: AABB
		// switch contact.type {
		// case .Body_vs_Body:
		// 	aabb2 = body_get_aabb(body2)
		// case .Body_vs_Level:
		// 	level_geom_handle := Level_Geom_Handle(contact.b)
		// 	level_geom := get_level_geom(sim_state, level_geom_handle)
		// 	aabb2 = level_geom.aabb
		// }

		// // TODO: extract common math functions into a sane place
		// if !collision.test_aabb_vs_aabb(
		// 	{min = aabb1.center - aabb1.extent, max = aabb1.center + aabb1.extent},
		// 	{min = aabb2.center - aabb2.extent, max = aabb2.center + aabb2.extent},
		// ) {
		// 	continue
		// }

		m1 := get_body_mesh(
			&body_mesh_cache,
			sim_state,
			body_handle_to_index(contact.a),
			contact.shape_a,
		)
		m2: collision.Convex

		switch contact.type {
		case .Body_vs_Body:
			m2 = get_body_mesh(
				&body_mesh_cache,
				sim_state,
				body_handle_to_index(b_handle),
				contact.shape_b,
			)
		case .Body_vs_Level:
			level_geom_handle := Level_Geom_Handle(contact.b)
			level_geom := get_level_geom(sim_state, level_geom_handle)
			_, vertices, indices := get_level_geom_data(sim_state, sim_cache, level_geom_handle)
			tri := get_transformed_triangle(
				vertices,
				indices,
				int(contact.local_tri_idx),
				level_geom.x,
				level_geom.q,
			)

			m2 = collision.double_sided_triangle_to_convex(tri, context.temp_allocator)
		}

		// Raw manifold has contact points in world space
		raw_manifold, collision := collision.convex_vs_convex_sat(m1, m2)

		if collision {
			assert(raw_manifold.points_len > 0)

			manifold := &contact.manifold
			manifold^ = raw_manifold

			// Convert manifold contact from world to local space
			for point_idx in 0 ..< manifold.points_len {
				manifold.points_a[point_idx] = body_world_to_local(
					body,
					manifold.points_a[point_idx],
				)
				manifold.points_b[point_idx] = body_world_to_local(
					body2,
					manifold.points_b[point_idx],
				)
			}

			preserve_impulse :=
				old_manifold.points_len == raw_manifold.points_len &&
				old_manifold.type == raw_manifold.type

			if preserve_impulse {
				contact.total_normal_impulse = old_total_normal_impulse

				for point_idx in 0 ..< manifold.points_len {
					contact.total_friction_impulse[point_idx].x = lg.dot(
						old_total_friction_impulse[point_idx].x * old_manifold.tangent,
						manifold.tangent,
					)
					contact.total_friction_impulse[point_idx].y = lg.dot(
						old_total_friction_impulse[point_idx].y * old_manifold.bitangent,
						manifold.bitangent,
					)
				}
			}
		}
	}
}

calculate_soft_constraint_params :: proc(
	natural_freq, damping_ratio, dt: f64,
) -> (
	bias_rate: f32,
	mass_coef: f32,
	impulse_coef: f32,
) {
	omega := 2.0 * math.PI * natural_freq // angular frequency
	a1 := 2.0 * damping_ratio + omega * dt
	a2 := dt * omega * a1
	a3 := 1.0 / (1.0 + a2)
	bias_rate = f32(omega / a1)
	mass_coef = f32(a2 * a3)
	impulse_coef = f32(a3)

	return
}

pgs_solve_contacts :: proc(
	sim_state: ^Sim_State,
	config: Solver_Config,
	dt: f32,
	inv_dt: f32,
	apply_bias: bool,
) {
	bias_rate, mass_coef, impulse_coef := calculate_soft_constraint_params(30, 8, f64(dt))
	if !apply_bias {
		mass_coef = 1
		bias_rate = 0
		impulse_coef = 0
	}

	random_order := make([]i32, len(sim_state.contact_container.contacts), context.temp_allocator)
	{
		for i in 0 ..< len(random_order) {
			random_order[i] = i32(i)
		}

		for i in 0 ..< len(random_order) - 1 {
			j := rand.int_max(len(random_order))
			slice.swap(random_order, i, j)
		}
	}

	for i in 0 ..< len(sim_state.contact_container.contacts) {
		contact := &sim_state.contact_container.contacts[random_order[i]]
		manifold := &contact.manifold

		b_handle := Body_Handle(contact.b) if contact.type == .Body_vs_Body else INVALID_BODY
		body1, body2 := get_body(sim_state, contact.a), get_body(sim_state, b_handle)

		for point_idx in 0 ..< manifold.points_len {
			p1, p2 :=
				body_local_to_world(body1, manifold.points_a[point_idx]),
				body_local_to_world(body2, manifold.points_b[point_idx])

			p_diff_normal := lg.dot(p2 - p1, manifold.normal)
			separation := p_diff_normal + 0.01

			w1 := get_body_inverse_mass(body1, manifold.normal, p1)
			w2 := get_body_inverse_mass(body2, manifold.normal, p2)

			w := w1 + w2

			if w == 0 {
				continue
			}

			inv_w := 1.0 / w


			{
				// r1, r2 := p1 - body1.x, p2 - body2.x
				v1 := body_velocity_at_point(body1, p1)
				v2 := body_velocity_at_point(body2, p2)

				delta_v := v2 - v1
				{
					delta_v_norm := lg.dot(delta_v, manifold.normal)

					bias := f32(0.0)

					MAX_BAUMGARTE_VELOCITY :: 1000
					if separation > 0 {
						bias = separation * inv_dt
					} else if apply_bias {
						bias = lg.max(bias_rate * separation, -MAX_BAUMGARTE_VELOCITY)
					}

					incremental_impulse :=
						-inv_w * mass_coef * (delta_v_norm + bias) -
						impulse_coef * contact.total_normal_impulse[point_idx]

					new_total_impulse := max(
						contact.total_normal_impulse[point_idx] + incremental_impulse,
						0,
					)
					applied_impulse := new_total_impulse - contact.total_normal_impulse[point_idx]
					contact.total_normal_impulse[point_idx] = new_total_impulse

					applied_impulse_vec := applied_impulse * manifold.normal

					apply_velocity_correction(body1, -applied_impulse_vec, p1)
					apply_velocity_correction(body2, applied_impulse_vec, p2)
				}
			}

			{
				// r1, r2 := p1 - body1.x, p2 - body2.x
				v1 := body_velocity_at_point(body1, p1)
				v2 := body_velocity_at_point(body2, p2)

				delta_v := v2 - v1

				DYNAMIC_FRICTION :: 0.6
				STATIC_FRICTION :: 0.8
				STATIC_FRICTION_VELOCITY_THRESHOLD :: 0.01

				delta_v_tang := Vec2 {
					lg.dot(delta_v, manifold.tangent),
					lg.dot(delta_v, manifold.bitangent),
				}

				use_static_friction :=
					lg.dot(delta_v_tang, delta_v_tang) <
					STATIC_FRICTION_VELOCITY_THRESHOLD * STATIC_FRICTION_VELOCITY_THRESHOLD
				friction: f32 = use_static_friction ? STATIC_FRICTION : DYNAMIC_FRICTION
				friction_clamp := contact.total_normal_impulse[point_idx] * friction

				incremental_impulse := -inv_w * delta_v_tang

				new_total_impulse: Vec2 = lg.clamp(
					contact.total_friction_impulse[point_idx] + incremental_impulse,
					Vec2(-friction_clamp),
					Vec2(friction_clamp),
				)

				applied_impulse := new_total_impulse - contact.total_friction_impulse[point_idx]
				contact.total_friction_impulse[point_idx] = new_total_impulse

				applied_impulse_vec :=
					applied_impulse.x * manifold.tangent + applied_impulse.y * manifold.bitangent

				// rl.DrawSphereWires(p1, 0.05, 8, 8, rl.RED)
				// rl.DrawLine3D(p1, p1 + v1, rl.RED)
				// rl.DrawSphereWires(p2, 0.05, 8, 8, rl.BLUE)
				// rl.DrawLine3D(p2, p2 + v2, rl.BLUE)

				apply_velocity_correction(body1, -applied_impulse_vec, p1)
				apply_velocity_correction(body2, applied_impulse_vec, p2)
			}
		}
	}
}

// Returns index into the gear ratios array
// -1 (revers) mapped to 0
// 1..N mapped to 0..N-1
lookup_gear_ratio :: #force_inline proc(gear_ratios: []f32, gear: i32) -> (ratio: f32) {
	assert(len(gear_ratios) > 1)
	if gear == 0 {
		return 0
	} else {
		index := int(gear + 1)
		if index > 0 {
			index -= 1
		}
		index = clamp(index, 0, len(gear_ratios) - 1)
		return gear_ratios[index] * (index == 0 ? -1 : 1)
	}
}

pgs_solve_engines :: proc(sim_state: ^Sim_State, config: Solver_Config, dt: f32, inv_dt: f32) {
	for &engine in sim_state.engines {
		if engine.alive {
			rpm_torque_curve := get_engine_curve(sim_state, engine.rpm_torque_curve)
			gear_ratios := get_gear_ratios(sim_state, engine.gear_ratios)

			rpm := angular_velocity_to_rpm(engine.w)
			throttle := engine.throttle
			clutch := f32(0)

			// Blending range in rpm of auto clutch activation
			AUTO_CLUTCH_RPM_RANGE :: f32(100)
			AUTO_BLIP_RPM_RANGE :: f32(100)

			// Prevent stalling
			{
				if rpm < engine.lowest_rpm {
					clutch =
						min(engine.lowest_rpm - rpm, AUTO_CLUTCH_RPM_RANGE) / AUTO_CLUTCH_RPM_RANGE

					throttle = max(
						throttle,
						min((engine.lowest_rpm - AUTO_BLIP_RPM_RANGE) - rpm, AUTO_BLIP_RPM_RANGE) /
						AUTO_BLIP_RPM_RANGE,
					)
				}
			}

			engine.clutch = clutch

			if engine.rev_limit_time < 0.0 {
				engine.rev_limit_time += dt
				throttle = 0
			} else if (rpm >= engine.rev_limit_rpm) {
				engine.rev_limit_time = -engine.rev_limit_interval
				throttle = 0
			}

			// Throttle
			{

				{
					torque: f32

					idx, _ := slice.binary_search_by(
						rpm_torque_curve,
						rpm,
						proc(a: [2]f32, k: f32) -> slice.Ordering {
							return slice.cmp(a[0], k)
						},
					)

					if idx > 0 && idx < len(rpm_torque_curve) - 1 {
						cur_point := rpm_torque_curve[idx]
						next_point := rpm_torque_curve[idx + 1]
						rpm_diff := next_point[0] - cur_point[0]
						alpha := (rpm - cur_point[0]) / rpm_diff

						torque = math.lerp(cur_point[1], next_point[1], alpha)
					} else {
						torque = rpm_torque_curve[math.clamp(idx, 0, len(rpm_torque_curve) - 1)][1]
					}

					// log.debugf("torque: %v Nm", torque)
					torque *= throttle

					engine.w += (torque / engine.inertia) * dt
				}
			}

			// Internal Friction
			{
				delta_omega := -engine.w

				inv_w := engine.inertia

				friction :=
					math.pow(max(engine.w - rpm_to_angular_velocity(engine.lowest_rpm), 0), 2) *
						0.0001 *
						engine.internal_friction +
					engine.internal_friction

				// Not physically based, but we assume when throttle is open there is no friction
				friction *= (1.0 - throttle)

				incremental_impulse := inv_w * delta_omega
				new_total_impulse := math.clamp(
					engine.friction_impulse + incremental_impulse,
					-friction,
					friction,
				)
				applied_impulse := new_total_impulse - engine.friction_impulse
				engine.friction_impulse = new_total_impulse

				engine.w += applied_impulse / engine.inertia
			}

			// Transmission
			{
				gear_ratio := lookup_gear_ratio(gear_ratios, engine.gear)

				if engine.gear != 0 {
					ratio := gear_ratio * engine.axle.final_drive_ratio
					inv_ratio := f32(1.0 / (ratio))

					drive_wheel1 := &engine.axle.wheels[0]
					wheel1 := get_suspension_constraint(sim_state, drive_wheel1.wheel)

					drive_wheel2 := &engine.axle.wheels[1]
					wheel2 := get_suspension_constraint(sim_state, drive_wheel2.wheel)

					w1 := f32(1.0 / (engine.inertia * ratio * ratio))
					w2 := wheel1.inv_inertia
					w3 := wheel2.inv_inertia

					w := w1 + w2 + w3
					inv_w := f32(1.0 / w)

					avg_wheel_vel := (wheel1.w + wheel2.w) * 0.5

					delta_omega := avg_wheel_vel * ratio - (-engine.w)

					incremental_impulse := -inv_w * delta_omega
					engine.axle.engine_impulse += incremental_impulse

					incremental_impulse *= (1.0 - clutch)

					engine.w += incremental_impulse * w1
					wheel1.w += incremental_impulse * w2 * inv_ratio
					wheel2.w += incremental_impulse * w3 * inv_ratio
				}
			}

			// Diff
			{
				switch engine.axle.diff_type {
				case .Open:
				case .Fixed:
					drive_wheel1 := &engine.axle.wheels[0]
					wheel1 := get_suspension_constraint(sim_state, drive_wheel1.wheel)

					drive_wheel2 := &engine.axle.wheels[1]
					wheel2 := get_suspension_constraint(sim_state, drive_wheel2.wheel)

					w1 := wheel1.inv_inertia
					w2 := wheel2.inv_inertia

					w := w1 + w2
					inv_w := f32(1.0 / w)

					delta_omega := wheel2.w - wheel1.w

					incremental_impulse := -inv_w * delta_omega
					engine.axle.diff_impulse += incremental_impulse

					wheel1.w += -incremental_impulse * wheel1.inv_inertia
					wheel2.w += incremental_impulse * wheel2.inv_inertia
				}
			}
		}
	}
}


calculate_ground_vel :: proc(
	body: Body_Ptr,
	hit_body: Body_Ptr,
	p: Vec3,
	right, forward: Vec3,
) -> (
	body_vel_at_contact_patch: Vec3,
	ground_vel: Vec2,
) {
	body_vel_at_contact_patch =
		body_velocity_at_point(body, p) - body_velocity_at_point(hit_body, p)
	ground_vel.x = lg.dot(body_vel_at_contact_patch, right)
	ground_vel.y = lg.dot(body_vel_at_contact_patch, forward)
	return
}

combine_inv_mass :: proc(w1, w2: f32) -> f32 {
	if w1 > 0 && w2 > 0 {
		return (w1 * w2) / (w1 + w2)
	}
	if w1 > 0 {
		return w1
	}

	return w2
}

pgs_solve_suspension :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	static_tlas: ^Static_TLAS,
	dyn_tlas: ^Dynamic_TLAS,
	config: Solver_Config,
	dt: f32,
	inv_dt: f32,
	apply_bias: bool,
) {
	// Solve suspension velocity
	for _, i in sim_state.suspension_constraints {
		v := &sim_state.suspension_constraints_slice[i]
		if v.alive {
			body := get_body(sim_state, v.body)

			if body.alive {
				if v.turn_wheel {
					v.turn_angle = math.clamp(v.turn_angle, -v.steer_lock, v.steer_lock)

					body_forward_vec := body_local_to_world_vec(body, {0, 0, 1})
					body_right_vec := body_local_to_world_vec(body, {1, 0, 0})
					lateral_to_longitudinal_relation :=
						lg.dot(body_right_vec, body.v) / (lg.dot(body_forward_vec, body.v) + 1.0)
					drift_amount :=
						math.smoothstep(
							f32(0.0),
							f32(1.0),
							abs(lateral_to_longitudinal_relation),
						) *
						math.sign(lateral_to_longitudinal_relation)

					v.turn_assist =
						math.smoothstep(f32(0), f32(1), lg.dot(body_forward_vec, body.v)) *
						drift_amount *
						math.PI *
						0.15

					// Clamp the assist to not go outside steer lock
					v.turn_assist =
						math.clamp(v.turn_angle + v.turn_assist, -v.steer_lock, v.steer_lock) -
						v.turn_angle
				}


				wheel_world_pos := body_local_to_world(body, v.rel_pos + body.shape_offset)
				dir := body_local_to_world_vec(body, v.rel_dir)
				prev_hit_body := v.hit_body
				v.hit_t, v.hit_normal, v.hit_body, v.hit = raycast(
					sim_state,
					sim_cache,
					static_tlas,
					dyn_tlas,
					wheel_world_pos,
					dir,
					v.rest,
				)
				hit_body := get_body(sim_state, v.hit_body)
				v.hit_point = wheel_world_pos + dir * v.hit_t

				w_normal1 := get_body_inverse_mass(body, v.hit_normal, v.hit_point)
				w_normal2 := get_body_inverse_mass(hit_body, v.hit_normal, v.hit_point)
				w_normal := combine_inv_mass(w_normal1, w_normal2)
				inv_w_normal := 1.0 / w_normal

				// Drive force
				if true {
					total_impulse := -v.drive_impulse

					v.w += v.inv_inertia * total_impulse * dt
				}

				// Brake force
				if true {
					// How strong is brake pad pushing against brake disc
					brake_pad_impulse := v.brake_impulse

					if brake_pad_impulse != 0 {
						// TODO: figure out what's the realistic value
						brake_friction := f32(1.0)
						friction_clamp := brake_pad_impulse * brake_friction

						incremental_impulse := (1.0 / v.inv_inertia) * -v.w
						new_total_impulse := clamp(
							v.brake_friction_impulse + incremental_impulse,
							-friction_clamp,
							friction_clamp,
						)
						applied_impulse := new_total_impulse - v.brake_friction_impulse
						v.brake_friction_impulse = new_total_impulse

						v.w += v.inv_inertia * applied_impulse
					} else {
						v.brake_friction_impulse = 0
					}
				}

				if !v.hit || v.hit_body != prev_hit_body {
					v.forward = 0
					v.right = 0
					v.lateral_impulse = 0
					v.spring_impulse = 0
					v.longitudinal_impulse = 0
				}

				if v.hit {
					forward := wheel_get_forward_vec(body, v)
					right := lg.normalize0(lg.cross(forward, v.hit_normal))
					forward = lg.normalize0(lg.cross(v.hit_normal, right))

					//v.longitudinal_impulse = lg.dot(forward, v.longitudinal_impulse * v.forward)
					//v.lateral_impulse = lg.dot(right, v.lateral_impulse * v.right)
					v.forward = forward
					v.right = right

					// Spring force
					{
						bias_coef, mass_coef, impulse_coef := calculate_soft_constraint_params(
							f64(v.natural_frequency),
							f64(v.damping),
							f64(dt),
						)
						if !apply_bias {
							// mass_coef = 1
							// bias_coef = 0
							// impulse_coef = 0
						}

						vel := lg.dot(
							body_velocity_at_point(hit_body, v.hit_point) -
							body_velocity_at_point(body, v.hit_point),
							dir,
						)
						x := v.hit_t
						separation := x - v.rest

						// spring_dot_normal := abs(lg.dot(dir, v.hit_normal))
						incremental_impulse :=
							-inv_w_normal * mass_coef * (vel + separation * bias_coef) -
							impulse_coef * v.spring_impulse
						new_spring_impulse := max(0, v.spring_impulse + incremental_impulse)
						applied_impulse := new_spring_impulse - v.spring_impulse
						v.spring_impulse = new_spring_impulse

						apply_velocity_correction(
							body,
							applied_impulse * v.hit_normal,
							v.hit_point,
						)
						apply_velocity_correction(
							hit_body,
							applied_impulse * -v.hit_normal,
							v.hit_point,
						)
					}

					// Positive means spinning forward
					wheel_spin_vel := -v.radius * v.w

					body_vel_at_contact_patch, ground_vel := calculate_ground_vel(
						body,
						hit_body,
						v.hit_point,
						right,
						forward,
					)

					slip_ratio :=
						ground_vel.y == 0 ? 0 : clamp(wheel_spin_vel / ground_vel.y - 1, -1, 1) * 100.0

					slip_angle :=
						-lg.angle_between(Vec2{0, 1}, Vec2{ground_vel.x, ground_vel.y}) *
						math.DEG_PER_RAD

					v.slip_ratio = slip_ratio
					v.slip_angle = slip_angle

					MAX_SLIP_LEN :: f32(2)

					slip_vec := Vec2 {
						slip_angle / PACEJKA94_LATERAL_PEAK_X / MAX_SLIP_LEN,
						slip_ratio / PACEJKA94_LONGITUDINAL_PEAK_X / MAX_SLIP_LEN,
					}

					slip_len := lg.length(slip_vec)
					slip_len = slip_len == 0 ? 0 : min(slip_len, 1) / slip_len
					slip_vec *= slip_len

					v.slip_vec = slip_vec

					// log.debugf("slip_vec: %v", slip_vec)

					camber_cos := -lg.dot(wheel_get_right_vec(body, v), v.right)
					camber_angle_rad := camber_cos < 0.999999 ? math.acos(camber_cos) : 0
					v.camber = camber_angle_rad * math.DEG_PER_RAD * math.sign(v.rel_pos.x)

					long_friction :=
						abs(
							pacejka_94_longitudinal(
								v.pacejka_long,
								slip_ratio,
								v.spring_impulse * inv_dt * 0.001,
							),
						) *
						abs(slip_vec.y)
					lat_friction :=
						abs(
							pacejka_94_lateral(
								v.pacejka_lat,
								slip_angle,
								v.spring_impulse * inv_dt * 0.001,
								v.camber,
							),
						) *
						abs(slip_vec.x)

					// Longitudinal friction
					if true {
						body_vel_at_contact_patch, ground_vel = calculate_ground_vel(
							body,
							hit_body,
							v.hit_point,
							right,
							forward,
						)

						// Wheel linear velocity relative to ground
						relative_vel := ground_vel.y - wheel_spin_vel

						friction_clamp := abs(v.spring_impulse) * long_friction

						w_body1 := get_body_inverse_mass(body, forward, v.hit_point)
						w_body2 := get_body_inverse_mass(hit_body, forward, v.hit_point)
						w_body := w_body1 + v.inv_inertia + w_body2

						w_long := w_body
						inv_w_long := 1.0 / w_long

						incremental_impulse := -inv_w_long * relative_vel
						new_total_impulse := clamp(
							v.longitudinal_impulse + incremental_impulse,
							-friction_clamp,
							friction_clamp,
						)
						applied_impulse := new_total_impulse - v.longitudinal_impulse
						v.longitudinal_impulse = new_total_impulse

						apply_velocity_correction(body, applied_impulse * forward, v.hit_point)
						apply_velocity_correction(
							hit_body,
							-applied_impulse * forward,
							v.hit_point,
						)
						v.w += v.inv_inertia * applied_impulse
					}
					// Lateral friction
					if true {
						body_vel_at_contact_patch, ground_vel = calculate_ground_vel(
							body,
							hit_body,
							v.hit_point,
							right,
							forward,
						)
						lateral_vel := ground_vel.x
						friction_clamp := abs(v.spring_impulse) * lat_friction

						w_body1 := get_body_inverse_mass(body, right, v.hit_point)
						w_body2 := get_body_inverse_mass(hit_body, right, v.hit_point)
						w_body := w_body1 + w_body2

						w_lat := w_body
						inv_w_lat := 1.0 / w_lat

						incremental_impulse := -inv_w_lat * lateral_vel
						new_total_impulse := clamp(
							v.lateral_impulse + incremental_impulse,
							-friction_clamp,
							friction_clamp,
						)
						applied_impulse := new_total_impulse - v.lateral_impulse
						v.lateral_impulse = new_total_impulse

						apply_velocity_correction(body, applied_impulse * right, v.hit_point)
						apply_velocity_correction(hit_body, -applied_impulse * right, v.hit_point)
					}
				}
			}
		}
	}
}

pgs_substep :: proc(
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	static_tlas: ^Static_TLAS,
	dyn_tlas: ^Dynamic_TLAS,
	config: Solver_Config,
	dt: f32,
	inv_dt: f32,
) {
	for i in 0 ..< len(sim_state.bodies_slice) {
		body := &sim_state.bodies_slice[i]

		if body.alive {
			body.v += config.gravity * dt * (body.inv_mass == 0 ? 0 : 1) // special case for gravity, TODO
		}
	}

	// Warm start
	if true {
		for i in 0 ..< len(sim_state.contact_container.contacts) {
			contact := &sim_state.contact_container.contacts[i]
			manifold := &contact.manifold

			body1, body2 :=
				get_body(sim_state, contact.a),
				get_body(
					sim_state,
					contact.type == .Body_vs_Body ? Body_Handle(contact.b) : INVALID_BODY,
				)

			for point_idx in 0 ..< manifold.points_len {
				p1, p2 :=
					body_local_to_world(body1, manifold.points_a[point_idx]),
					body_local_to_world(body2, manifold.points_b[point_idx])
				total_normal_impulse := contact.total_normal_impulse[point_idx] * manifold.normal

				apply_velocity_correction(body1, -total_normal_impulse, p1)
				apply_velocity_correction(body2, total_normal_impulse, p2)

				total_tangent_impulse :=
					contact.total_friction_impulse[point_idx].x * manifold.tangent +
					contact.total_friction_impulse[point_idx].y * manifold.bitangent

				apply_velocity_correction(body1, -total_tangent_impulse, p1)
				apply_velocity_correction(body2, total_tangent_impulse, p2)
			}
		}

		for i in 0 ..< len(sim_state.engines) {
			e := &sim_state.engines[i]

			if e.alive {
				gear_ratios := get_gear_ratios(sim_state, e.gear_ratios)
				// e.w += e.unstall_impulse / e.inertia

				e.w += e.friction_impulse / e.inertia

				// Warm start engine torque
				if false {
					drive_wheel1 := &e.axle.wheels[0]
					wheel1 := get_suspension_constraint(sim_state, drive_wheel1.wheel)

					drive_wheel2 := &e.axle.wheels[1]
					wheel2 := get_suspension_constraint(sim_state, drive_wheel2.wheel)

					gear_ratio := lookup_gear_ratio(gear_ratios, e.gear)

					if e.gear != 0 {
						ratio := gear_ratio * e.axle.final_drive_ratio
						inv_ratio := f32(1.0 / (ratio))

						w1 := f32(1.0 / (e.inertia * ratio * ratio))
						w2 := wheel1.inv_inertia
						w3 := wheel2.inv_inertia

						e.w += e.axle.engine_impulse * w1 * (1.0 - e.clutch)
						wheel1.w += e.axle.engine_impulse * w2 * inv_ratio * (1.0 - e.clutch)
						wheel2.w += e.axle.engine_impulse * w3 * inv_ratio * (1.0 - e.clutch)
					}

					// Warmp start diff impulse
					if false {
						switch e.axle.diff_type {
						case .Open:
						case .Fixed:
							wheel1.w += -e.axle.diff_impulse * wheel1.inv_inertia
							wheel2.w += e.axle.diff_impulse * wheel2.inv_inertia
						}
					}
				}
			}
		}

		for i in 0 ..< len(sim_state.suspension_constraints) {
			s := &sim_state.suspension_constraints_slice[i]

			if s.hit {
				body := get_body(sim_state, s.body)
				hit_body := get_body(sim_state, s.hit_body)
				hit_p := body_local_to_world(
					body,
					s.rel_pos + body.shape_offset + s.rel_dir * s.hit_t,
				)

				apply_velocity_correction(body, s.spring_impulse * s.hit_normal, hit_p)
				apply_velocity_correction(hit_body, s.spring_impulse * -s.hit_normal, hit_p)

				apply_velocity_correction(body, s.lateral_impulse * s.right, hit_p)
				apply_velocity_correction(hit_body, -s.lateral_impulse * s.right, hit_p)

				apply_velocity_correction(body, s.longitudinal_impulse * s.forward, hit_p)
				apply_velocity_correction(hit_body, -s.longitudinal_impulse * s.forward, hit_p)
				s.w += s.inv_inertia * s.longitudinal_impulse
			}

			s.w += s.inv_inertia * s.brake_friction_impulse
		}
	}

	apply_bias := true
	pgs_solve_contacts(sim_state, config, dt, inv_dt, apply_bias)
	pgs_solve_engines(sim_state, config, dt, inv_dt)
	pgs_solve_suspension(
		sim_state,
		sim_cache,
		static_tlas,
		dyn_tlas,
		config,
		dt,
		inv_dt,
		apply_bias,
	)

	for i in 0 ..< len(sim_state.bodies_slice) {
		body := &sim_state.bodies_slice[i]

		if body.alive {
			body.x += body.v * dt

			// NOTE: figure out how this works https://fgiesen.wordpress.com/2012/08/24/quaternion-differentiation/
			q := body.q

			w := body.w
			delta_rot := quaternion(x = w.x, y = w.y, z = w.z, w = 0)
			delta_rot = delta_rot * q
			q.x += 0.5 * dt * delta_rot.x
			q.y += 0.5 * dt * delta_rot.y
			q.z += 0.5 * dt * delta_rot.z
			q.w += 0.5 * dt * delta_rot.w
			q = lg.normalize0(q)

			body.q = q
		}
	}

	for i in 0 ..< len(sim_state.engines) {
		e := &sim_state.engines[i]

		e.q = math.mod_f32(e.q + 0.5 * e.w * dt, math.PI * 2)
	}

	for i in 0 ..< len(sim_state.suspension_constraints_slice) {
		s := &sim_state.suspension_constraints_slice[i]

		s.x += s.v * dt
		s.q = math.mod_f32(s.q + 0.5 * s.w * dt, math.PI * 2)
	}

	apply_bias = false
	pgs_solve_contacts(sim_state, config, dt, inv_dt, apply_bias)
	// pgs_solve_suspension(
	// 	sim_state,
	// 	sim_cache,
	// 	static_tlas,
	// 	dyn_tlas,
	// 	config,
	// 	dt,
	// 	inv_dt,
	// 	apply_bias,
	// )
}

simulate_step :: proc(
	scene: ^Scene,
	sim_state: ^Sim_State,
	sim_cache: ^Sim_Cache,
	config: Solver_Config,
) {
	tracy.Zone()

	// Save prev state
	{
		for i in 0 ..< len(sim_state.bodies) {
			body := &sim_state.bodies_slice[i]

			if body.alive {
				body.prev_x = body.x
				body.prev_v = body.v
				body.prev_q = body.q
				body.prev_w = body.w
			}
		}

		for i in 0 ..< len(sim_state.suspension_constraints) {
			wheel := &sim_state.suspension_constraints_slice[i]

			if wheel.alive {
				wheel.prev_q = wheel.q
				wheel.prev_w = wheel.w
				wheel.prev_hit = wheel.hit
				wheel.prev_hit_normal = wheel.hit_normal
				wheel.prev_hit_body = wheel.hit_body
				wheel.prev_hit_point = wheel.hit_point
				wheel.prev_hit_t = wheel.hit_t
				wheel.prev_turn_angle = wheel.turn_angle
				wheel.prev_turn_assist = wheel.turn_assist
			}
		}
	}

	substeps := config.substreps_minus_one + 1

	dt_64 := f64(config.timestep) / f64(substeps)
	inv_dt_64 := f64(1.0) / dt_64

	dt := f32(dt_64)
	inv_dt := f32(inv_dt_64)

	build_static_tlas(sim_state, &sim_state.static_tlas)
	build_dynamic_tlas(sim_state, config, &sim_state.dynamic_tlas)

	remove_invalid_contacts(sim_state, sim_cache, sim_state.static_tlas, sim_state.dynamic_tlas)
	find_new_contacts(
		sim_state,
		sim_cache,
		&sim_state.static_tlas,
		&sim_state.dynamic_tlas,
		config,
	)
	update_contacts(sim_state, sim_cache, &sim_state.static_tlas)

	Solver :: enum {
		XPBD,
		PGS,
	}

	solver := Solver.PGS

	switch solver {
	case .XPBD:
		for _ in 0 ..< substeps {
			xpbd_substep(sim_state, config, dt, inv_dt)
		}
	case .PGS:
		for _ in 0 ..< substeps {
			pgs_substep(
				sim_state,
				sim_cache,
				&sim_state.static_tlas,
				&sim_state.dynamic_tlas,
				config,
				dt,
				inv_dt,
			)
		}
	}

}

body_solve_velocity :: #force_inline proc(
	body: Body_Ptr,
	prev_x: Vec3,
	prev_q: Quat,
	inv_dt: f32,
) {
	body.v = (body.x - prev_x) * inv_dt

	delta_q := body.q * lg.quaternion_inverse(prev_q)
	body.w = Vec3{delta_q.x, delta_q.y, delta_q.z} * 2.0 * inv_dt

	if delta_q.w < 0 {
		body.w = -body.w
	}
}

calculate_constraint_params1 :: proc(
	dt: f32,
	body: Body_Ptr,
	compliance: f32,
	error: f32,
	error_gradient: Vec3,
	pos: Vec3,
	other_combined_inv_mass: f32,
) -> (
	lambda: f32,
	w: f32,
	correction: Vec3,
	ok: bool,
) {
	if error == 0 {
		return
	}

	w = get_body_inverse_mass(body, error_gradient, pos)
	w += other_combined_inv_mass

	if w == 0 {
		return
	}

	ok = true

	alpha := compliance / dt / dt
	lambda = -error / (w + alpha)
	correction = -error_gradient * -lambda

	return
}

calculate_constraint_params2 :: proc(
	dt: f32,
	body1: Body_Ptr,
	body2: Body_Ptr,
	compliance: f32,
	error: f32,
	error_gradient: Vec3,
	pos1: Vec3,
	pos2: Vec3,
) -> (
	lambda: f32,
	correction1: Vec3,
	correction2: Vec3,
	ok: bool,
) {
	if error == 0 {
		return
	}

	w := get_body_inverse_mass(body1, error_gradient, pos1)
	w += get_body_inverse_mass(body2, error_gradient, pos2)

	if w == 0 {
		return
	}

	ok = true

	alpha := compliance / dt / dt
	lambda = -error / (w + alpha)
	correction1 = -error_gradient * -lambda
	correction2 = error_gradient * -lambda

	return
}

apply_constraint_correction_unilateral :: proc(
	dt: f32,
	body: Body_Ptr,
	compliance: f32,
	error: f32,
	error_gradient: Vec3,
	pos: Vec3,
	other_combined_inv_mass: f32 = 0,
) -> (
	lambda: f32,
) {
	w: f32
	correction: Vec3
	ok: bool
	lambda, w, correction, ok = calculate_constraint_params1(
		dt,
		body,
		compliance,
		error,
		error_gradient,
		pos,
		other_combined_inv_mass,
	)

	if ok {
		apply_position_correction(body, correction, pos)
	}

	return
}

multiply_inv_intertia :: proc(body: Body_Ptr, vec: Vec3) -> (result: Vec3) {
	q := body.q
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(q))

	result = lg.quaternion_mul_vector3(inv_q, vec)
	result *= body.inv_inertia_tensor
	result = lg.quaternion_mul_vector3(q, result)

	return result
}

apply_velocity_correction :: #force_inline proc(body: Body_Ptr, impulse: Vec3, pos: Vec3) {
	apply_velocity_correction_linear(body, impulse)

	angular_impulse := lg.cross(pos - body.x, impulse)

	apply_velocity_correction_angular(body, angular_impulse)
}

apply_velocity_correction_linear :: #force_inline proc "contextless" (
	body: Body_Ptr,
	impulse: Vec3,
) {
	body.v += impulse * body.inv_mass
}

apply_velocity_correction_angular :: #force_inline proc "contextless" (
	body: Body_Ptr,
	angular_impulse: Vec3,
) {
	q := body.q
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(q))
	delta_omega := lg.quaternion_mul_vector3(inv_q, angular_impulse)
	delta_omega *= body.inv_inertia_tensor
	delta_omega = lg.quaternion_mul_vector3(q, delta_omega)

	body.w += delta_omega
}

apply_position_correction :: proc(body: Body_Ptr, corr: Vec3, pos: Vec3) {
	body.x += corr * body.inv_mass

	q := body.q
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(q))
	delta_omega := pos - body.x
	delta_omega = lg.cross(delta_omega, corr)
	delta_omega = lg.quaternion_mul_vector3(inv_q, delta_omega)
	delta_omega *= body.inv_inertia_tensor
	delta_omega = lg.quaternion_mul_vector3(q, delta_omega)

	delta_rot := quaternion(x = delta_omega.x, y = delta_omega.y, z = delta_omega.z, w = 0)
	delta_rot *= q
	q.x += 0.5 * delta_rot.x
	q.y += 0.5 * delta_rot.y
	q.z += 0.5 * delta_rot.z
	q.w += 0.5 * delta_rot.w
	q = lg.normalize0(q)

	body.q = q
}

// Total inverse mass (linear + angular)
get_body_inverse_mass :: proc(body: Body_Ptr, normal, pos: Vec3) -> f32 {
	linear, angular := get_body_inverse_mass_separate(body, normal, pos)

	return linear + angular
}

get_body_inverse_mass_separate :: proc(
	body: Body_Ptr,
	normal, pos: Vec3,
) -> (
	linear: f32,
	angular: f32,
) {
	q := body.q
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(q))

	rn := pos - body.x
	rn = lg.cross(rn, normal)
	rn = lg.quaternion_mul_vector3(inv_q, rn)

	linear = body.inv_mass
	angular = lg.dot(rn, rn * body.inv_inertia_tensor)

	return
}

get_body_angular_inverse_mass :: proc(body: Body_Ptr, normal: Vec3) -> f32 {
	q := body.q
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(q))

	n := lg.quaternion_mul_vector3(inv_q, normal)

	return lg.dot(n, n * body.inv_inertia_tensor)
}