package physics

import "core:fmt"
import "core:math"
import "core:math/rand"
import "core:math/linalg"
import rl "vendor:raylib"
import "vendor:raylib/rlgl"

/*
	XPBD algorithm
	- integrate
	- solve constraints
	- correct positions
*/

quat :: quaternion128

RigidBody :: struct {
	pos: [3]f32,
	rot: quat,

	vel: [3]f32,
	omega: [3]f32, // angular velocity

	inv_mass: f32,
	inv_inertia: matrix[3,3]f32,

	static, active: bool,
	col: Collider,

	// previous state
	prev_pos: [3]f32,
	prev_rot: quaternion128,
}

// decent enough approximations for drag without doing real math.
// otherwise we would have to calculate forces and integrate them properly.
// fraction of velocity LOST per second (0.0 = none, higher = more damping)
LINEAR_DAMPING :: 0.2
ANGULAR_DAMPING :: 0.01
FRICTION :: 0.5
GRAVITY :: [3]f32{0, -10, 0}
DT :: f32(1.0/60.0)
SUBSTEP_CNT :: 10

bodies : [dynamic]RigidBody
contacts : [dynamic]Contact

body_create :: proc(col: Collider, pos: [3]f32, mass: f32,
	rot := linalg.Quaternionf32(1), static := false) -> RigidBody {
	static := static

	body := RigidBody {
		pos      = pos,
		rot      = rot,
		prev_pos = pos,
		prev_rot = rot,
		col = col,
	}
	if mass <= 0  && !static {
		fmt.println("warn:" + 
			"non static body with no mass was set to static")
		static = true
	}
	if !static {
		body.inv_mass = 1.0 / mass
		body.inv_inertia = collider_inv_inertia(col, mass)
	}

	body.static = static
	return body
}

inv_inertia_world :: proc(body: ^RigidBody) -> matrix[3,3]f32 {
	rot_mat := linalg.matrix3_from_quaternion(body.rot)
	return rot_mat * body.inv_inertia * linalg.transpose(rot_mat)
}

generalized_inv_mass :: proc(body: ^RigidBody, r: [3]f32, n: [3]f32) -> f32 {
	r_cross_n := linalg.cross(r, n)
	inv_inertia_w := inv_inertia_world(body)
	angular := linalg.dot(r_cross_n, inv_inertia_w * r_cross_n)
	return body.inv_mass + angular
}

apply_impulse :: proc(body: ^RigidBody, r: [3]f32, impulse: [3]f32) {
	// Note the sign convention: callers pass +impulse to one body and 
	// -impulse to the other, so apply_impulse itself never negates
	// Newton's third law is the caller's job.

	// linear: heavier bodies move less
	body.pos += impulse * body.inv_mass

	// angular: torque = r × impulse
	inv_inertia_w := inv_inertia_world(body)
	angular := inv_inertia_w * linalg.cross(r, impulse)
	angular_q := quaternion(x = angular.x, y = angular.y,
				z = angular.z, w = 0)
	spin := angular_q * body.rot
	body.rot += quaternion(x = spin.x * 0.5, y = spin.y * 0.5,
			       z = spin.z * 0.5, w = spin.w * 0.5)
	body.rot = linalg.normalize(body.rot)
}

solve_positional :: proc(a, b: ^RigidBody, r1, r2, correction: [3]f32,
	compliance, dt: f32) -> f32 {
	c := linalg.length(correction)

	if c == 0 {
		return 0
	}

	n := correction / c
	w1 := generalized_inv_mass(a, r1, n)
	w2 := generalized_inv_mass(b, r2, n)
	w_sum := w1 + w2

	if w_sum == 0 {
		return 0
	}

	alpha := compliance / (dt * dt)
	delta_lambda := c / (w_sum + alpha)
	impulse := delta_lambda * n

	apply_impulse(a, r1, impulse)
	apply_impulse(b, r2, -impulse)

	return delta_lambda
}

solve_attachment :: proc(body: ^RigidBody, local_point, world_point: [3]f32,
	compliance, dt: f32) {
	r := linalg.quaternion_mul_vector3(body.rot, local_point)
	correction := world_point - (body.pos + r)

	magnitude := linalg.length(correction)
	if magnitude == 0 {
		return
	}
	n := correction / magnitude
	w := generalized_inv_mass(body, r, n)

	if w == 0 {
		return
	}

	alpha := compliance / (dt * dt)
	delta_lambda := magnitude / (w + alpha)
	apply_impulse(body, r, delta_lambda * n)
}

friction_solve :: proc(contact: ^Contact, dt: f32) {
	a := contact.a
	b := contact.b
	n := contact.normal

	r1 := linalg.quaternion_mul_vector3(a.rot, contact.r1_local)
	r2 := linalg.quaternion_mul_vector3(b.rot, contact.r2_local)

	p1_cur := a.pos + r1
	p2_cur := b.pos + r2
	p1_prev := a.prev_pos + 
		linalg.quaternion_mul_vector3(a.prev_rot, contact.r1_local)
	p2_prev := b.prev_pos +
		linalg.quaternion_mul_vector3(b.prev_rot, contact.r2_local)

	delta := (p1_cur - p1_prev) - (p2_cur - p2_prev)
	delta_t := delta - linalg.dot(delta, n) * n
	slide := linalg.length(delta_t)

	if slide < 1e-8 {
		return
	}
	tangent_dir := delta_t / slide

	correction_mag := min(slide, FRICTION * contact.normal_lambda)
	correction := -correction_mag * tangent_dir

	solve_positional(a, b, r1, r2, correction, 0, dt)

	// fmt.printfln("slide=%v clamp=%v applied=%v normal_lambda=%v",
		// slide, FRICTION * contact.normal_lambda,
		// min(slide, FRICTION * contact.normal_lambda),
		// contact.normal_lambda)
}

integrate :: proc(body: ^RigidBody, dt: f32) {
	body.prev_pos = body.pos
	body.prev_rot = body.rot

	// linear motion
	if body.inv_mass > 0 {
		body.vel += GRAVITY * dt
	}
	body.pos += body.vel * dt

	// angular motion
	omega_q := quaternion(
		x = body.omega.x, 
		y = body.omega.y, 
		z = body.omega.z, 
		w = 0)
	spin := omega_q * body.rot
	scale := 0.5 * dt
	delta := quaternion(
		x = spin.x * scale,
		y = spin.y * scale,
		z = spin.z * scale,
		w = spin.w * scale,
	)
	body.rot += delta
	body.rot = linalg.normalize(body.rot)
}

recover_velocity :: proc(body: ^RigidBody, dt: f32) {
	body.vel = (body.pos - body.prev_pos) / dt

	rdelta := body.rot * linalg.quaternion_inverse(body.prev_rot)
	body.omega = 2 * [3]f32{rdelta.x, rdelta.y, rdelta.z} / dt

	if rdelta.w < 0 {
		body.omega = -body.omega
	}
}

/*
	physics steps:
	- integrate
	- solve collisions/contacts
	- solve constraints
	- recover velocity
*/

physics_simulate :: proc() {
	dt := DT / f32(SUBSTEP_CNT)

	// raised to pow of dt to remove frame rate dependency
	linear_damp_factor  := math.pow(1 - LINEAR_DAMPING, dt)
	angular_damp_factor := math.pow(1 - ANGULAR_DAMPING, dt)

	for substep in 0..<SUBSTEP_CNT {
		for &body in bodies {
			integrate(&body, dt)
		}

		clear(&contacts)
		for i in 0..<len(bodies) {
			for j in i+1..<len(bodies) {
				collide(&bodies[i], &bodies[j], &contacts)
			}
		}

		for &contact in contacts {
			contact_solve(&contact, dt)
			friction_solve(&contact, dt)
		}

		for &body in bodies {
			recover_velocity(&body, dt)
			body.vel *= linear_damp_factor
			body.omega *= angular_damp_factor
		}
	}
}

// BOX DEBUG HELPERS
box_corners :: proc(body: ^RigidBody) -> [8][3]f32 {
	half := body.col.(Box).sz / 2
	axes := box_axes(body)
	corners: [8][3]f32
	idx := 0
	for sx in ([2]f32{-1, 1}) {
	for sy in ([2]f32{-1, 1}) {
	for sz in ([2]f32{-1, 1}) {
		local := axes[0]*half.x*sx + axes[1]*half.y*sy + axes[2]*half.z*sz
		corners[idx] = body.pos + local
		idx += 1
	}}}
	return corners
}

// returns the axis index (0,1,2) and sign (+1/-1) of the face
// whose outward normal best matches `dir`
box_best_face :: proc(axes: [3][3]f32, dir: [3]f32) -> (axis: int, sign: f32) {
	best_dot := f32(-1e30)
	for i in 0..<3 {
		d := linalg.dot(axes[i], dir)
		if d > best_dot {
			best_dot = d
			axis = i
			sign = 1
		}
		if -d > best_dot {
			best_dot = -d
			axis = i
			sign = -1
		}
	}
	return
}

// the 4 world-space corners of the face identified by (axis, sign)
box_face_corners :: proc(body: ^RigidBody, axes: [3][3]f32,
	axis: int, sign: f32) -> [4][3]f32 {
	half := body.col.(Box).sz / 2
	// the two axes spanning the face (the ones that aren't `axis`)
	u := (axis + 1) % 3
	v := (axis + 2) % 3

	// center of the face: box center pushed out to the face along `axis`
	face_center := body.pos + axes[axis] * (half[axis] * sign)

	hu := half[u]
	hv := half[v]
	return {
		face_center + axes[u]*hu + axes[v]*hv,
		face_center + axes[u]*hu - axes[v]*hv,
		face_center - axes[u]*hu - axes[v]*hv,
		face_center - axes[u]*hu + axes[v]*hv,
	}
}

// for a confirmed box-box contact with normal `n` (points a->b),
// returns the reference and incident face corner sets.
// also returns which body owns the reference (for later).
box_contact_faces :: proc(a, b: ^RigidBody, n: [3]f32) ->
	(ref_corners, inc_corners: [4][3]f32, ref_is_a: bool) {

	axes_a := box_axes(a)
	axes_b := box_axes(b)

	// how aligned is each box's best face with the normal?
	// reference face = most aligned with the contact normal.
	// for body a, the separating direction pushing toward b is +n.
	// for body b, it's -n. we compare alignment.
	a_axis, a_sign := box_best_face(axes_a, n)
	b_axis, b_sign := box_best_face(axes_b, -n)

	a_align := abs(linalg.dot(axes_a[a_axis] * a_sign, n))
	b_align := abs(linalg.dot(axes_b[b_axis] * b_sign, -n))

	if a_align >= b_align {
		// a owns the reference face; b owns the incident
		ref_corners = box_face_corners(a, axes_a, a_axis, a_sign)
		// incident face on b: most anti-aligned with n (points back toward a)
		i_axis, i_sign := box_best_face(axes_b, -n)
		inc_corners = box_face_corners(b, axes_b, i_axis, i_sign)
		ref_is_a = true
	} else {
		ref_corners = box_face_corners(b, axes_b, b_axis, b_sign)
		i_axis, i_sign := box_best_face(axes_a, n)
		inc_corners = box_face_corners(a, axes_a, i_axis, i_sign)
		ref_is_a = false
	}
	return
}

// END BOX DEBUG HELPERS

body_draw :: proc(body: ^RigidBody) {
	angle, axis := linalg.angle_axis_from_quaternion(body.rot)
	rlgl.PushMatrix()
	rlgl.Translatef(body.pos.x, body.pos.y, body.pos.z)
	rlgl.Rotatef(angle * 180 / math.PI, axis.x, axis.y, axis.z)

	switch shape in body.col {
	case Sphere:
		rl.DrawSphere({0,0,0}, shape.radius, rl.Fade(rl.BLUE, 0.5))
		rl.DrawSphereWires({0,0,0}, shape.radius + 0.01, 8, 16, rl.BLACK)

	case Box:
		rl.DrawCubeV({0,0,0}, shape.sz, rl.GRAY)
		rl.DrawCubeWiresV({0,0,0}, shape.sz + 0.005, rl.BLACK)
	}

	rlgl.PopMatrix()
}

start_sim :: proc() {
	clear(&bodies)

	// one cube near the edge, given a shove toward it
	// cube := body_create(Box{sz={1,1,1}}, {4, 2, 0}, 10)
	// cube.vel = {3, 0, 0}   // push it toward the +x edge
	// append(&bodies, cube)

	// floor
	append(&bodies, body_create(Box{sz={50,1,50}}, {0,-0.55,0}, 0, static = true))

	// 4 walls
	// append(&bodies, body_create(Box{sz={50,2,1}}, {0,1.05,24.5}, 0, static = true))
	// append(&bodies, body_create(Box{sz={50,2,1}}, {0,1.05,-24.5}, 0, static = true))
	// append(&bodies, body_create(Box{sz={1,2,47.95}}, {24.5,1.05,0}, 0, static = true))
	// append(&bodies, body_create(Box{sz={1,2,47.95}}, {-24.5,1.05,0}, 0, static = true))

	// larger box floating in the air
	// append(&bodies, body_create(Box{sz={2,2,2}}, {0, 1, 0}, 0, static = true))

	// N := 3
	// for x in 0..<N {
	// for y in 0..<N {
	// for z in 0..<N {
	// 	spacing := f32(2.5)
	// 	hw := (spacing) * f32(N) / f32(2)
	// 	pos := [3]f32 {
	// 		spacing * f32(x) - hw,
	// 		spacing * f32(y) + 10,
	// 		spacing * f32(z) - hw,
	// 	}
	// 	// rot := linalg.quaternion_from_pitch_yaw_roll_f32(10, 50, 23)
	// 	rot := quaternion128(1)
	// 	// append(&bodies, body_create(Sphere{radius = 0.5}, pos, 10, rot))
	// 	append(&bodies, body_create(Box{sz={1,1,1}}, pos, 10, rot))
	// }}}

	append(&bodies, body_create(Box{sz={1,1,1}}, {0, 1, 0}, 1))
	append(&bodies, body_create(Box{sz={1,1,1}}, {0, 2.2, 0}, 1))
}

main :: proc() {
	rl.SetTraceLogLevel(.WARNING)
	rl.SetConfigFlags({.MSAA_4X_HINT})
	rl.InitWindow(1920, 1080, "Physics")
	rl.SetWindowState({.WINDOW_RESIZABLE, .FULLSCREEN_MODE})
	rl.SetTargetFPS(60)
	rl.SetExitKey(.KEY_NULL)
	rl.DisableCursor()

	camera := rl.Camera3D {
		position = {5, 5, 5},
		target = {0, 1, 0},
		up = {0, 1, 0},
		fovy = 60,
		projection = .PERSPECTIVE,
	}

	start_sim()

	cursor_locked := true
	for !rl.WindowShouldClose() {
		if rl.IsKeyPressed(.ESCAPE) {
			cursor_locked = !cursor_locked
			if cursor_locked {
				rl.DisableCursor()
			} else {
				rl.EnableCursor()
			}
		}

		if rl.IsKeyPressed(.R) {
			start_sim()
		}

		if cursor_locked {
			rl.UpdateCamera(&camera, .THIRD_PERSON)
		}

		physics_simulate()

		rl.BeginDrawing()
		rl.ClearBackground(rl.RAYWHITE)

		rl.BeginMode3D(camera)
		// rl.DrawGrid(50, 1)

		for &body in bodies {
			body_draw(&body)
		}

		rl.EndMode3D()

		rl.DrawFPS(10, 10)
		rl.EndDrawing()
	}
}