package physics

import "collision"
import "core:math"
import lg "core:math/linalg"
import rl "vendor:raylib"

_ :: math

Solver_Config :: struct {
	// Will automatically do fixed timestep
	timestep: f32,
	gravity:  rl.Vector3,
}

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

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

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

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

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

	dt := config.timestep
	inv_dt := 1.0 / dt

	// Integrate positions and rotations
	for &body, i in scene.bodies {
		if body.alive {
			body_states[i].prev_x = body.x
			body.v += dt * config.gravity
			body.x += dt * body.v

			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
			delta_rot := quaternion(x = body.w.x, y = body.w.y, z = body.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 &v in scene.suspension_constraints {
		if v.alive {
			body := get_body(scene, v.body)

			q := body.q
			pos := body_local_to_world(body, v.rel_pos)
			dir := lg.quaternion_mul_vector3(q, 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,
				)
			}
		}
	}

	// Compute new linear and angular velocities
	for &body, i in scene.bodies {
		if body.alive {
			body.v = (body.x - body_states[i].prev_x) * inv_dt

			delta_q := body.q * lg.quaternion_inverse(body_states[i].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_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_inverse(q)

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

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

	return w
}