package physics

import "bvh"
import "collision"
import "core:container/bit_array"
import "core:fmt"
import "core:math"
import lg "core:math/linalg"
import "core:math/rand"
import "core:slice"
import "libs:tracy"

_ :: rand
_ :: math
_ :: fmt
_ :: slice

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
	num_referenced_bodies:                 i32,
	num_referenced_suspension_constraints: i32,
	immedate_bodies:                       map[u32]Immedate_State(Body_Handle),
	immediate_suspension_constraints:      map[u32]Immedate_State(Suspension_Constraint_Handle),
}

destroy_solver_state :: proc(state: ^Solver_State) {
	delete(state.immedate_bodies)
	delete(state.immediate_suspension_constraints)
}

Immedate_State :: struct($T: typeid) {
	handle:   T,
	// When was this referenced last time (frame number)
	last_ref: u32,
}

MAX_STEPS :: 10

// TODO: move into scene.odin
// Copy current state to next
sim_state_copy :: proc(dst: ^Sim_State, src: ^Sim_State) {
	tracy.Zone()
	convex_container_reconcile(&src.convex_container)

	dst.num_bodies = src.num_bodies
	dst.first_free_body_plus_one = src.first_free_body_plus_one
	dst.first_free_suspension_constraint_plus_one = src.first_free_suspension_constraint_plus_one

	resize(&dst.bodies, len(src.bodies))
	resize(&dst.suspension_constraints, len(src.suspension_constraints))

	dst.bodies_slice = dst.bodies[:]
	dst.suspension_constraints_slice = dst.suspension_constraints[:]

	for i in 0 ..< len(dst.bodies) {
		dst.bodies[i] = src.bodies[i]
	}
	for i in 0 ..< len(dst.suspension_constraints) {
		dst.suspension_constraints[i] = src.suspension_constraints[i]
	}

	contact_container_copy(&dst.contact_container, src.contact_container)
	convex_container_copy(&dst.convex_container, src.convex_container)
}

Step_Mode :: enum {
	Accumulated_Time,
	Single,
}

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

// TODO: free intermediate temp allocs
// Creates TLAS using temp allocator
build_tlas :: proc(sim_state: ^Sim_State, config: Solver_Config) -> TLAS {
	tracy.Zone()

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

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

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

				phys_aabb := body_get_aabb(body)

				EXPAND_K :: 2
				expand := lg.abs(EXPAND_K * config.timestep * body.v)
				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, context.temp_allocator)

	return TLAS{bvh_tree = sim_state_bvh, body_aabbs = body_aabbs}
}

// TODO: free intermediate temp allocs
find_new_contacts :: proc(sim_state: ^Sim_State, tlas: ^TLAS) {
	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_aabb := tlas.body_aabbs[i]
			it := bvh.iterator_intersect_leaf(&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 := tlas.bvh_tree.primitives[leaf_node.child_or_prim_start + j]
					prim_aabb := tlas.body_aabbs[other_body_idx]

					pair := make_contact_pair(i32(body_idx), i32(other_body_idx))
					if body_idx != other_body_idx &&
					   bvh.test_aabb_vs_aabb(body_aabb, prim_aabb) &&
					   !(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 {
							a = Body_Handle(i + 1),
							b = Body_Handle(other_body_idx + 1),
						}

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

// Outer simulation loop for fixed timestepping
simulate :: proc(
	scene: ^Scene,
	state: ^Solver_State,
	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,
) {
	tracy.Zone()
	assert(config.timestep > 0)

	prune_immediate(scene, state)

	sim_state_copy(get_next_sim_state(scene), get_sim_state(scene))

	sim_state := get_next_sim_state(scene)

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

		for state.accumulated_time >= config.timestep {
			num_steps += 1
			state.accumulated_time -= config.timestep

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

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

	state.simulation_frame += 1
	state.num_referenced_bodies = 0
	state.num_referenced_suspension_constraints = 0
}

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

Contact :: struct {
	a, b:                      Body_Handle,
	prev_x_a, prev_x_b:        Vec3,
	prev_q_a, prev_q_b:        Quat,
	manifold:                  collision.Contact_Manifold,
	applied_corrections:       int,
	lambda_normal:             [4]f32,
	lambda_tangent:            [4]f32,
	applied_static_friction:   [4]bool,
	applied_normal_correction: [4]f32,
}

update_contacts :: proc(sim_state: ^Sim_State) {
	tracy.Zone()

	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,
		)
	}

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

		i, j := i32(contact.a) - 1, i32(contact.b) - 1

		body, body2 := &sim_state.bodies_slice[i], &sim_state.bodies_slice[j]

		assert(body.alive)
		assert(body2.alive)

		contact.prev_x_a = body.x
		contact.prev_x_b = body2.x
		contact.prev_q_a = body.q
		contact.prev_q_b = body2.q
		contact.manifold = {}
		contact.lambda_normal = 0
		contact.lambda_tangent = 0
		contact.applied_static_friction = false
		contact.applied_normal_correction = 0

		aabb1, aabb2 := body_get_aabb(body), body_get_aabb(body2)

		// 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, m2 :=
			body_get_convex_shape_world(sim_state, body),
			body_get_convex_shape_world(sim_state, body2)

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

		if collision {
			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],
				)
			}
		}
	}
}

xpbd_predict_positions :: proc(sim_state: ^Sim_State, config: Solver_Config, dt: f32) {
	// Integrate positions and rotations
	for &body in sim_state.bodies {
		if body.alive {
			body.prev_x = body.x
			body.prev_v = body.v
			body.prev_w = body.w
			body.prev_q = body.q

			body.v += config.gravity * dt * (body.inv_mass == 0 ? 0 : 1) // special case for gravity, TODO
			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
		}
	}
}

xpbd_substep :: proc(sim_state: ^Sim_State, config: Solver_Config, dt: f32, inv_dt: f32) {
	xpbd_predict_positions(sim_state, config, dt)

	Body_Pair :: struct {
		a, b: int,
	}

	{
		tracy.ZoneN("simulate_step::solve_collisions")
		for i in 0 ..< len(sim_state.contact_container.contacts) {
			contact := &sim_state.contact_container.contacts[i]
			body, body2 := get_body(sim_state, contact.a), get_body(sim_state, contact.b)

			contact^ = Contact {
				a        = contact.a,
				b        = contact.b,
				prev_x_a = body.x,
				prev_x_b = body2.x,
				prev_q_a = body.q,
				prev_q_b = body2.q,
				manifold = contact.manifold,
			}

			manifold := &contact.manifold

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

				p_diff_normal := lg.dot(p2 - p1, manifold.normal)
				separation := min(p_diff_normal, 0)

				lambda_norm, corr1, corr2, ok := calculate_constraint_params2(
					dt,
					body,
					body2,
					0.00002,
					separation,
					-manifold.normal,
					p1,
					p2,
				)
				if ok {
					contact.applied_normal_correction[point_idx] = -separation
					contact.applied_corrections += 1
					contact.lambda_normal[point_idx] = lambda_norm

					apply_correction(body, corr1, p1)
					apply_correction(body2, corr2, p2)
				}
			}
		}
	}

	if false {
		tracy.ZoneN("simulate_step::static_friction")

		when false {
			context = context
			context.user_ptr = sim_state
			slice.sort_by(
				sim_state.contact_pairs[:sim_state.contact_pairs_len],
				proc(c1, c2: Contact) -> bool {
					sim_state := cast(^Sim_State)context.user_ptr

					find_min_contact_y :: proc(
						scene: ^Sim_State,
						c: Contact,
					) -> (
						min_contact_y: f32,
					) {
						min_contact_y = max(f32)
						body_a, body_b := get_body(scene, c.a), get_body(scene, c.b)
						for i in 0 ..< c.manifold.points_len {
							min_contact_y = min(
								body_local_to_world(body_a, c.manifold.points_a[i]).y,
								body_local_to_world(body_b, c.manifold.points_b[i]).y,
							)
						}
						return
					}

					min_y1 := find_min_contact_y(sim_state, c1)
					min_y2 := find_min_contact_y(sim_state, c2)

					return min_y1 > min_y2
				},
			)
		}

		for &contact in sim_state.contact_container.contacts {
			manifold := contact.manifold
			body, body2 := get_body(sim_state, contact.a), get_body(sim_state, contact.b)

			for point_idx in 0 ..< manifold.points_len {
				lambda_norm := contact.lambda_normal[point_idx]
				if lambda_norm != 0 {
					p1 := body_local_to_world(body, manifold.points_a[point_idx])
					p2 := body_local_to_world(body2, manifold.points_b[point_idx])
					prev_p1 :=
						body.prev_x +
						lg.quaternion_mul_vector3(body.prev_q, manifold.points_a[point_idx])
					prev_p2 :=
						body2.prev_x +
						lg.quaternion_mul_vector3(body2.prev_q, manifold.points_b[point_idx])

					p_diff_tangent := (p1 - prev_p1) - (p2 - prev_p2)
					p_diff_tangent -= lg.dot(p_diff_tangent, manifold.normal) * manifold.normal

					tangent_diff_len := lg.length(p_diff_tangent)

					if tangent_diff_len > 0 {
						tangent_diff_normalized := p_diff_tangent / tangent_diff_len

						delta_lambda_tangent, corr1_tangent, corr2_tangent, ok_tangent :=
							calculate_constraint_params2(
								dt,
								body,
								body2,
								0,
								-tangent_diff_len,
								-tangent_diff_normalized,
								p1,
								p2,
							)

						STATIC_FRICTION :: 0
						if ok_tangent && delta_lambda_tangent < STATIC_FRICTION * lambda_norm {
							contact.applied_static_friction[point_idx] = true
							contact.lambda_tangent[point_idx] = delta_lambda_tangent

							apply_correction(body, corr1_tangent, p1)
							apply_correction(body2, corr2_tangent, p2)
						}
					}
				}
			}
		}
	}

	{
		tracy.ZoneN("simulate_step::suspension_constraints")

		for &v in sim_state.suspension_constraints {
			if v.alive {
				body := get_body(sim_state, v.body)

				pos := body_local_to_world(body, v.rel_pos)
				dir := body_local_to_world_vec(body, v.rel_dir)
				pos2 := pos + dir * v.rest
				v.hit_t, v.hit_point, v.hit = collision.intersect_segment_plane(
					{pos, pos2},
					collision.plane_from_point_normal({}, collision.Vec3{0, 1, 0}),
				)
				if v.hit {
					corr := pos - v.hit_point
					distance := lg.length(corr)
					corr = corr / distance if distance > 0 else 0

					_ = apply_constraint_correction_unilateral(
						dt,
						body,
						v.compliance,
						error = distance - v.rest,
						error_gradient = corr,
						pos = pos,
						other_combined_inv_mass = 0,
					)
				}
			}
		}
	}

	{
		// Compute new linear and angular velocities
		for _, i in sim_state.bodies_slice {
			body := &sim_state.bodies_slice[i]
			if body.alive {
				body_solve_velocity(body, body.prev_x, body.prev_q, inv_dt)
			}
		}
	}

	// Restituion
	if true {
		tracy.ZoneN("simulate_step::restitution")

		for &contact in sim_state.contact_container.contacts {
			manifold := &contact.manifold


			body, body2 := get_body(sim_state, contact.a), get_body(sim_state, contact.b)
			prev_q1, prev_q2 := body.prev_q, body2.prev_q

			for point_idx in 0 ..< manifold.points_len {
				if contact.lambda_normal[point_idx] == 0 {
					continue
				}
				prev_r1 := lg.quaternion_mul_vector3(prev_q1, manifold.points_a[point_idx])
				prev_r2 := lg.quaternion_mul_vector3(prev_q2, manifold.points_b[point_idx])

				r1 := lg.quaternion_mul_vector3(body.q, manifold.points_a[point_idx])
				r2 := lg.quaternion_mul_vector3(body2.q, manifold.points_b[point_idx])

				prev_v :=
					(body.prev_v + lg.cross(body.prev_w, prev_r1)) -
					(body2.prev_v + lg.cross(body2.prev_w, prev_r2))
				v := (body.v + lg.cross(body.w, r1)) - (body2.v + lg.cross(body2.w, r2))

				prev_v_normal := lg.dot(prev_v, manifold.normal)
				v_normal := lg.dot(v, manifold.normal)

				RESTITUTION :: 0

				restitution := f32(RESTITUTION)

				if abs(v_normal) <= 2 * abs(lg.dot(manifold.normal, -config.gravity) * dt) {
					restitution = 0
				}

				delta_v := manifold.normal * (-v_normal + min(-RESTITUTION * prev_v_normal, 0))

				w1 := get_body_inverse_mass(body, manifold.normal, r1 + body.x)
				w2 := get_body_inverse_mass(body2, manifold.normal, r2 + body2.x)
				w := w1 + w2

				if w != 0 {
					p := delta_v / w

					body.v += p * body.inv_mass
					body2.v -= p * body2.inv_mass

					body.w += multiply_inv_intertia(body, lg.cross(r1, p))
					body2.w -= multiply_inv_intertia(body2, lg.cross(r2, p))
				}
			}
		}
	}

	if true {
		tracy.ZoneN("simulate_step::dynamic_friction")

		for &contact in sim_state.contact_container.contacts {
			manifold := &contact.manifold
			body1 := get_body(sim_state, contact.a)
			body2 := get_body(sim_state, contact.b)

			for point_idx in 0 ..< contact.manifold.points_len {
				if contact.applied_static_friction[point_idx] || contact.lambda_normal == 0 {
					continue
				}
				p1, p2 :=
					body_local_to_world(body1, manifold.points_a[point_idx]),
					body_local_to_world(body2, manifold.points_b[point_idx])
				r1, r2 := p1 - body1.x, p2 - body2.x
				v1 := body_velocity_at_point(body1, p1)
				v2 := body_velocity_at_point(body2, p2)

				v := v1 - v2
				v_normal := lg.dot(manifold.normal, v) * manifold.normal
				v_tangent := v - v_normal

				DYNAMIC_FRICTION :: 1
				v_tangent_len := lg.length(v_tangent)
				if v_tangent_len > 0 {
					v_tangent_norm := v_tangent / v_tangent_len

					w1, w2 :=
						get_body_inverse_mass(body1, v_tangent_norm, p1),
						get_body_inverse_mass(body2, v_tangent_norm, p2)

					w := w1 + w2

					if w > 0 {
						delta_v :=
							-v_tangent_norm *
							min(
								dt *
								DYNAMIC_FRICTION *
								abs(contact.lambda_normal[point_idx] / (dt * dt)),
								v_tangent_len / w,
							)

						// delta_v_norm := lg.normalize0(delta_v)

						p := delta_v

						body1.v += p * body1.inv_mass
						body2.v -= p * body2.inv_mass

						body1.w += multiply_inv_intertia(body1, lg.cross(r1, p))
						body2.w -= multiply_inv_intertia(body2, lg.cross(r2, p))
					}
				}
			}
		}
	}


	// 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 && v.hit {
				prev_x, prev_q := body.prev_x, body.prev_q

				wheel_world_pos := body_local_to_world(body, v.rel_pos)
				prev_wheel_world_pos := prev_x + lg.quaternion_mul_vector3(prev_q, v.rel_pos)

				vel_3d := (wheel_world_pos - prev_wheel_world_pos) * inv_dt
				dir := body_local_to_world_vec(body, v.rel_dir)

				// Spring damping
				if true {
					vel := lg.dot(vel_3d, dir)
					damping := -vel * min(v.damping * dt, 1)

					_ = apply_constraint_correction_unilateral(
						dt,
						body,
						0,
						error = damping,
						error_gradient = -dir,
						pos = wheel_world_pos,
					)

					body_solve_velocity(body, body.prev_x, body.prev_q, inv_dt)
				}

				// Drive forces
				if true {
					total_impulse := v.drive_impulse - v.brake_impulse
					forward := body_local_to_world_vec(body, Vec3{0, 0, 1})

					_ = apply_constraint_correction_unilateral(
						dt,
						body,
						0,
						total_impulse * dt * dt,
						-forward,
						wheel_world_pos,
					)
					body_solve_velocity(body, body.prev_x, body.prev_q, inv_dt)
				}

				// Lateral friction
				if true {
					vel_contact := body_velocity_at_point(body, v.hit_point)
					right := wheel_get_right_vec(body, v)

					lateral_vel := lg.dot(right, vel_contact)

					friction := f32(0.5)
					impulse := -lateral_vel * friction
					corr := right * impulse * dt
					v.applied_impulse.x = impulse

					apply_correction(body, corr, v.hit_point)
					body_solve_velocity(body, body.prev_x, body.prev_q, inv_dt)
				}
			}
		}
	}
}

pgs_substep :: proc(sim_state: ^Sim_State, 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
		}
	}

	// TODO: apply impulses
	// for i in 0 ..< len(sim_state.contact_pairs) {
	// 	contact_pair := &sim_state.contact_pairs[i]

	//   
	// }

	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
		}
	}
}

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

	substeps := config.substreps_minus_one + 1

	dt := config.timestep / f32(substeps)
	inv_dt := 1.0 / dt


	tlas := build_tlas(sim_state, config)
	find_new_contacts(sim_state, &tlas)

	{
		tracy.ZoneN("simulate_step::remove_invalid_contacts")
		i := 0
		for i < len(sim_state.contact_container.contacts) {
			contact := sim_state.contact_container.contacts[i]

			aabb_a := tlas.body_aabbs[int(contact.a) - 1]
			aabb_b := tlas.body_aabbs[int(contact.b) - 1]

			if !bvh.test_aabb_vs_aabb(aabb_a, aabb_b) {
				removed_pair := make_contact_pair(i32(contact.a) - 1, i32(contact.b) - 1)
				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 := make_contact_pair(
						i32(moved_contact.a) - 1,
						i32(moved_contact.b) - 1,
					)
					sim_state.contact_container.lookup[moved_pair] = i32(i)
				}
			} else {
				i += 1
			}
		}
	}

	update_contacts(sim_state)

	Solver :: enum {
		XPBD,
		PGS,
	}

	solver := Solver.XPBD

	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, 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_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_correction :: proc(body: Body_Ptr, corr: Vec3, pos: Vec3) {
	// rl.DrawSphereWires(pos, 0.5, 4, 4, rl.BLUE)
	// rl.DrawLine3D(pos, pos + corr, rl.BLUE)

	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
}

get_body_inverse_mass :: proc(body: Body_Ptr, normal, pos: Vec3) -> 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)

	w := lg.dot(rn, rn * body.inv_inertia_tensor)
	w += body.inv_mass

	return w
}