package physics

import "collision"
import "core:fmt"
import "core:math"
import lg "core:math/linalg"
import "core:slice"
import "libs:tracy"
import rl "vendor:raylib"

_ :: math
_ :: fmt

Solver_Config :: struct {
	// Will automatically do fixed timestep
	timestep:            f32,
	gravity:             rl.Vector3,
	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

// Outer simulation loop for fixed timestepping
simulate :: proc(scene: ^Scene, state: ^Solver_State, config: Solver_Config, dt: f32) {
	assert(config.timestep > 0)

	prune_immediate(scene, state)

	state.accumulated_time += dt

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

		if num_steps < MAX_STEPS {
			simulate_step(scene, config)
		}
	}

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

Body_Sim_State :: struct {
	prev_x: rl.Vector3,
	prev_v: rl.Vector3,
	prev_w: rl.Vector3,
	prev_q: rl.Quaternion,
}

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

Contact_Pair :: struct {
	a, b:                      Body_Handle,
	prev_x_a, prev_x_b:        rl.Vector3,
	prev_q_a, prev_q_b:        rl.Quaternion,
	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,
}

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

	body_states := make([]Body_Sim_State, len(scene.bodies), context.temp_allocator)

	scene.contact_pairs_len = 0

	substeps := config.substreps_minus_one + 1

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

	for _ in 0 ..< substeps {
		// Integrate positions and rotations
		for &body, i in scene.bodies {
			if body.alive {
				body_states[i].prev_x = body.x
				body_states[i].prev_v = body.v
				body_states[i].prev_w = body.w
				body_states[i].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
			}
		}


		Body_Pair :: struct {
			a, b: int,
		}
		handled_pairs := make_map(map[Body_Pair]bool, context.temp_allocator)

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

			for _, i in scene.bodies {
				body := &scene.bodies_slice[i]
				if body.alive {
					for _, j in scene.bodies {
						body2 := &scene.bodies_slice[j]

						if i != j &&
						   body2.alive &&
						   !handled_pairs[{a = min(i, j), b = max(i, j)}] {
							m1, m2 :=
								body_get_convex_shape_world(body),
								body_get_convex_shape_world(body2)

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

							if collision {
								contact_pair := &scene.contact_pairs[scene.contact_pairs_len]
								contact_pair^ = Contact_Pair {
									a        = Body_Handle(i + 1),
									b        = Body_Handle(j + 1),
									prev_x_a = body.x,
									prev_x_b = body2.x,
									prev_q_a = body.q,
									prev_q_b = body2.q,
									manifold = raw_manifold,
								}
								scene.contact_pairs_len += 1
								manifold := &contact_pair.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],
									)
								}
								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)

									handled_pairs[{a = min(i, j), b = max(i, j)}] = true

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

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

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

			if false {
				context = context
				context.user_ptr = scene
				slice.sort_by(
					scene.contact_pairs[:scene.contact_pairs_len],
					proc(c1, c2: Contact_Pair) -> bool {
						scene := cast(^Scene)context.user_ptr

						find_min_contact_y :: proc(
							scene: ^Scene,
							c: Contact_Pair,
						) -> (
							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(scene, c1)
						min_y2 := find_min_contact_y(scene, c2)

						return min_y1 > min_y2
					},
				)
			}

			for &contact_pair in scene.contact_pairs[:scene.contact_pairs_len] {
				manifold := contact_pair.manifold
				body, body2 := get_body(scene, contact_pair.a), get_body(scene, contact_pair.b)
				i, j := int(contact_pair.a) - 1, int(contact_pair.b) - 1

				for point_idx in 0 ..< manifold.points_len {
					lambda_norm := contact_pair.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_states[i].prev_x +
							lg.quaternion_mul_vector3(
								body_states[i].prev_q,
								manifold.points_a[point_idx],
							)
						prev_p2 :=
							body_states[j].prev_x +
							lg.quaternion_mul_vector3(
								body_states[j].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 / max(f32(contact_pair.applied_corrections) * 0.5, 1),
									-tangent_diff_normalized,
									p1,
									p2,
								)

							STATIC_FRICTION :: 0.6
							if ok_tangent && delta_lambda_tangent < STATIC_FRICTION * lambda_norm {
								contact_pair.applied_static_friction[point_idx] = true
								contact_pair.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 scene.suspension_constraints {
				if v.alive {
					body := get_body(scene, 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,
						)
					}
				}
			}
		}

		solve_velocities(scene, body_states, inv_dt)

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

			for &pair in scene.contact_pairs[:scene.contact_pairs_len] {
				i, j := int(pair.a) - 1, int(pair.b) - 1
				manifold := &pair.manifold

				body, body2 := get_body(scene, pair.a), get_body(scene, pair.b)
				s1, s2 := body_states[i], body_states[j]
				prev_q1, prev_q2 := s1.prev_q, s2.prev_q

				for point_idx in 0..<manifold.points_len {
					if pair.lambda_normal == 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 := (s1.prev_v + lg.cross(s1.prev_w, prev_r1)) - (s2.prev_v + lg.cross(s2.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.5

					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 + body.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 &pair in scene.contact_pairs[:scene.contact_pairs_len] {
				manifold := &pair.manifold
				body1 := get_body(scene, pair.a)
				body2 := get_body(scene, pair.b)

				for point_idx in 0 ..< pair.manifold.points_len {
					if pair.applied_static_friction[point_idx] {
						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 :: 0.4
					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(pair.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 scene.suspension_constraints {
			v := &scene.suspension_constraints_slice[i]
			if v.alive {
				body_idx := int(v.body) - 1
				body := get_body(scene, v.body)

				if body.alive && v.hit {
					prev_x, prev_q := body_states[body_idx].prev_x, body_states[body_idx].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)
					body_state := body_states[i32(v.body) - 1]

					// 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_state, inv_dt)
					}

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

						_ = apply_constraint_correction_unilateral(
							dt,
							body,
							0,
							total_impulse * dt * dt,
							-forward,
							wheel_world_pos,
						)
						body_solve_velocity(body, body_state, 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_state, inv_dt)
					}
				}
			}
		}
	}
}

solve_velocities :: proc(scene: ^Scene, body_states: []Body_Sim_State, inv_dt: f32) {
	// Compute new linear and angular velocities
	for _, i in scene.bodies_slice {
		body := &scene.bodies_slice[i]
		if body.alive {
			body_solve_velocity(body, body_states[i], inv_dt)
		}
	}
}

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

	delta_q := body.q * lg.quaternion_inverse(state.prev_q)
	body.w = rl.Vector3{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: rl.Vector3,
	pos: rl.Vector3,
	other_combined_inv_mass: f32,
) -> (
	lambda: f32,
	w: f32,
	correction: rl.Vector3,
	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: rl.Vector3,
	pos1: rl.Vector3,
	pos2: rl.Vector3,
) -> (
	lambda: f32,
	correction1: rl.Vector3,
	correction2: rl.Vector3,
	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: rl.Vector3,
	pos: rl.Vector3,
	other_combined_inv_mass: f32 = 0,
) -> (
	lambda: f32,
) {
	w: f32
	correction: rl.Vector3
	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: rl.Vector3) -> (result: rl.Vector3) {
	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: rl.Vector3, pos: rl.Vector3) {
	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: rl.Vector3) -> 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
}