package physics

import "collision"
import "core:fmt"
import "core:math"
import lg "core:math/linalg"
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,
}

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

		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_q: rl.Quaternion,
}

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

simulate_step :: proc(scene: ^Scene, config: Solver_Config) {
	body_states := make([]Body_Sim_State, len(scene.bodies), context.temp_allocator)

	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.v += config.gravity * dt
				body.x += body.v * dt

				body_states[i].prev_q = body.q

				// TODO: Probably can do it using built in quaternion math but I have no idea how it works
				// 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 scene.bodies {
			body := &scene.bodies_slice[i]
			if body.alive {
				normal := GLOBAL_PLANE.normal
				penetration: f32
				hit_point: rl.Vector3
				hit := false

				switch s in body.shape {
				case Shape_Box:
					extent := s.size * 0.5
					local_plane := collision.Plane {
						normal = body_world_to_local_vec(body, GLOBAL_PLANE.normal),
						dist   = -lg.dot(GLOBAL_PLANE.normal, body.x) - GLOBAL_PLANE.dist,
					}

					penetration, hit = collision.test_box_vs_halfspace(
						collision.Box{pos = 0, rad = extent},
						local_plane,
					)
					if hit {
						local_hit_point: rl.Vector3
						min_distance := f32(math.F32_MAX)
						for corner in collision.BOX_CORNERS_NORM {
							local_pos := extent * corner
							dist := collision.signed_distance_plane(local_pos, local_plane)
							if dist < min_distance {
								min_distance = dist
								local_hit_point = local_pos
							}
						}
						hit_point = body_local_to_world(body, local_hit_point)
						penetration = -min(min_distance, 0)
					}
				case Shape_Sphere:
					penetration, hit = collision.test_sphere_vs_halfspace(
						collision.Sphere{pos = body.x, rad = s.radius},
						GLOBAL_PLANE,
					)
				}

				if hit {
					apply_constraint_correction_unilateral(
						dt,
						body,
						0,
						penetration,
						normal,
						hit_point,
						0,
					)
				}
			}
		}

		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 := v.hit_point - pos
					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)

		// 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 {
						local_contact_pos := v.hit_point - body.x
						vel_contact := body_velocity_at_local_point(body, local_contact_pos)
						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
	}
}

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,
) {
	if error == 0 {
		return
	}

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

	if w == 0 {
		return
	}

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

	delta_pos := error_gradient * -lambda

	apply_correction(body, delta_pos, pos)
}

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 :=
		rn.x * rn.x * body.inv_inertia_tensor.x +
		rn.y * rn.y * body.inv_inertia_tensor.y +
		rn.z * rn.z * body.inv_inertia_tensor.z

	w += body.inv_mass

	return w
}