package physics

import lg "core:math/linalg"
import "libs:tracy"

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: f32 = lg.dot(p2 - p1, manifold.normal)
					separation: f32 = 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_position_correction(body, corr1, p1)
						apply_position_correction(body2, corr2, p2)
					}
				}

				if false && contact.lambda_normal[point_idx] != 0 {
					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)

					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: Vec3 = (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.00002,
								-tangent_diff_len,
								-tangent_diff_normalized,
								p1,
								p2,
							)

						STATIC_FRICTION :: 1000
						if ok_tangent &&
						   delta_lambda_tangent >
							   STATIC_FRICTION * contact.lambda_normal[point_idx] {
							contact.applied_static_friction = true
							contact.lambda_tangent = delta_lambda_tangent

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

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

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

			prev_p1, prev_p2: Vec3
			p1, p2: Vec3
			total_lambda_normal := f32(0)
			friciton_points_len := 0

			for point_idx in 0 ..< contact.manifold.points_len {
				if contact.lambda_normal == 0 {
					continue
				}

				total_lambda_normal += contact.lambda_normal[point_idx]
				friciton_points_len += 1

				point_p1, point_p2 :=
					body_local_to_world(body1, manifold.points_a[point_idx]),
					body_local_to_world(body2, manifold.points_b[point_idx])

				p1 += point_p1
				p2 += point_p2

				prev_point_p1 :=
					body1.prev_x +
					lg.quaternion_mul_vector3(body1.prev_q, manifold.points_a[point_idx])
				prev_point_p2 :=
					body2.prev_x +
					lg.quaternion_mul_vector3(body2.prev_q, manifold.points_b[point_idx])

				prev_p1 += prev_point_p1
				prev_p2 += prev_point_p2
			}

			if friciton_points_len > 0 {
				p1 /= f32(friciton_points_len)
				p2 /= f32(friciton_points_len)

				prev_p1 /= f32(friciton_points_len)
				prev_p2 /= f32(friciton_points_len)

				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,
							body1,
							body2,
							0,
							-tangent_diff_len,
							tangent_diff_normalized,
							p1,
							p2,
						)

					STATIC_FRICTION :: 1.0
					if ok_tangent {
						contact.applied_static_friction = true
						contact.lambda_tangent = delta_lambda_tangent

						apply_position_correction(body1, corr1_tangent, p1)
						apply_position_correction(body2, corr2_tangent, p2)
					}
				}
			}
		}
	}

	{
		// 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 {
			if contact.manifold.points_len == 0 {
				continue
			}

			manifold := &contact.manifold
			body1 := get_body(sim_state, contact.a)
			body2 := get_body(sim_state, contact.b)

			friction_p1, friction_p2: Vec3
			total_lambda_normal := f32(0)
			friciton_points_len := 0

			for point_idx in 0 ..< contact.manifold.points_len {
				if contact.applied_static_friction || contact.lambda_normal == 0 {
					continue
				}

				total_lambda_normal += contact.lambda_normal[point_idx]
				friciton_points_len += 1

				p1, p2 :=
					body_local_to_world(body1, manifold.points_a[point_idx]),
					body_local_to_world(body2, manifold.points_b[point_idx])

				friction_p1 += p1
				friction_p2 += p2
			}

			if friciton_points_len > 0 {
				friction_p1 /= f32(friciton_points_len)
				friction_p2 /= f32(friciton_points_len)
				v1 := body_velocity_at_point(body1, friction_p1)
				v2 := body_velocity_at_point(body2, friction_p2)
				r1, r2 := friction_p1 - body1.x, friction_p2 - body2.x

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

				DYNAMIC_FRICTION :: 0.5
				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, friction_p1),
						get_body_inverse_mass(body2, v_tangent_norm, friction_p2)

					w := w1 + w2

					if w != 0 {
						delta_v :=
							-v_tangent_norm *
							min(
								dt * DYNAMIC_FRICTION * abs(total_lambda_normal / (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))
					}
				}

				{
					angular_vel_error :=
						lg.dot(body1.w, manifold.normal) - lg.dot(body2.w, manifold.normal)

					w1, w2 :=
						get_body_angular_inverse_mass(body1, manifold.normal),
						get_body_angular_inverse_mass(body2, manifold.normal)

					w := w1 + w2

					if w != 0 {
						angular_impulse := manifold.normal * -angular_vel_error / (w1 + w2)

						apply_angular_velocity_correction :: proc(
							body: Body_Ptr,
							angular_impulse: Vec3,
						) {
							q := body.q
							inv_q := lg.quaternion_inverse(q)

							delta_omega := angular_impulse
							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)

							body.w += delta_omega
						}

						apply_angular_velocity_correction(body1, angular_impulse)
						apply_angular_velocity_correction(body2, -angular_impulse)
					}
				}
			}
		}
	}


	// 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_position_correction(body, corr, v.hit_point)
					body_solve_velocity(body, body.prev_x, body.prev_q, inv_dt)
				}
			}
		}
	}
}