package physics

import "collision"
import "core:log"
import "core:math"
import lg "core:math/linalg"
import "game:halfedge"
import "game:name"
import "libs:tracy"

_ :: log

inertia_tensor_sphere :: proc(radius: f32) -> (tensor: Matrix3) {
	tensor = radius * radius * (2.0 / 3.0)

	return
}

inertia_tensor_box :: proc(size: Vec3) -> (tensor: Matrix3) {
	CONSTANT :: f32(1.0 / 12.0)

	tensor[0][0] = size.z * size.z + size.y * size.y
	tensor[1][1] = size.x * size.x + size.z * size.z
	tensor[2][2] = size.x * size.x + size.y * size.y

	tensor = tensor * Matrix3(CONSTANT)

	return
}

inertia_tensor_collision_shape :: proc(shape: Input_Shape_Internal) -> (tensor: Matrix3) {
	switch s in shape {
	case Shape_Box:
		tensor = inertia_tensor_box(s.size)
	case Shape_Convex:
		// TODO: assuming precomputed
		tensor = s.inertia_tensor
	}

	return
}

center_of_mass_shape :: proc(shape: Input_Shape_Internal) -> (com: Vec3) {
	switch s in shape {
	case Shape_Box:
		com = 0
	case Shape_Convex:
		com = s.center_of_mass
	}
	return
}

volume_shape :: proc(shape: Input_Shape_Internal) -> (vol: f32) {
	switch s in shape {
	case Shape_Box:
		vol = s.size.x * s.size.y * s.size.z
	case Shape_Convex:
		vol = s.total_volume
	}
	return
}

combined_center_of_mass :: proc(shapes: []Input_Shape) -> (com: Vec3) {
	total_weight := f32(0)
	for shape in shapes {
		inner_com := center_of_mass_shape(shape.inner_shape) + shape.rel_x
		volume := volume_shape(shape.inner_shape) * (shape.density_minus_one + 1)
		com += inner_com * volume
		total_weight += volume
	}
	com = total_weight > 0 ? com / total_weight : 0
	return
}

// Returns normalized inertia
// https://pybullet.org/Bullet/phpBB3/viewtopic.php?t=246
combined_inertia_tensor :: proc(shapes: []Input_Shape) -> (inertia: Matrix3) {
	combined_com := combined_center_of_mass(shapes)

	total_volume := f32(0)
	for shape in shapes {
		com := center_of_mass_shape(shape.inner_shape) + shape.rel_x
		r := com - combined_com
		local_inertia := inertia_tensor_collision_shape(shape.inner_shape)
		rotation_mat: Matrix3 = auto_cast lg.matrix3_from_quaternion(shape.rel_q)
		inertia_common_align := lg.transpose(rotation_mat) * local_inertia * rotation_mat
		volume := volume_shape(shape.inner_shape) * (shape.density_minus_one + 1)
		dot_r := lg.dot(r, r)
		// outer_r: Matrix3 = auto_cast lg.outer_product(r, r)
		#unroll for i in 0 ..< 3 {
			#unroll for j in 0 ..< 3 {
				kron: f32 = i == j ? 1 : 0
				inertia[i][j] += inertia_common_align[i][j] + volume * (dot_r * kron - r[i] * r[j])
			}
		}
		// inertia = inertia + inertia_common_align * (Matrix3(dot_r) - outer_r) * Matrix3(volume)
		total_volume += volume
	}
	inertia *= Matrix3(1.0 / total_volume)
	return
}

body_local_to_world :: #force_inline proc(body: Body_Ptr, pos: Vec3) -> Vec3 {
	return body.x + lg.quaternion_mul_vector3(body.q, pos)
}

body_local_to_world_vec :: #force_inline proc(body: Body_Ptr, vec: Vec3) -> Vec3 {
	return lg.normalize0(lg.quaternion_mul_vector3(body.q, vec))
}

body_world_to_local :: #force_inline proc(body: Body_Ptr, pos: Vec3) -> Vec3 {
	inv_q := lg.quaternion_normalize0(lg.quaternion_inverse(body.q))
	return lg.quaternion_mul_vector3(inv_q, pos - body.x)
}

body_world_to_local_vec :: #force_inline proc(body: Body_Ptr, vec: Vec3) -> Vec3 {
	inv_q := lg.quaternion_inverse(body.q)
	return lg.normalize0(lg.quaternion_mul_vector3(inv_q, vec))
}

body_velocity_at_point :: #force_inline proc(body: Body_Ptr, pos: Vec3) -> Vec3 {
	return body.v + lg.cross(body.w, pos - body.x)
}

wheel_get_rel_wheel_pos :: #force_inline proc(
	body: Body_Ptr,
	wheel: Suspension_Constraint_Ptr,
) -> Vec3 {
	t := wheel.hit_t > 0 ? wheel.hit_t : wheel.rest
	return wheel.rel_pos + wheel.rel_dir * (t - wheel.radius)
}

wheel_get_right_vec :: #force_inline proc(
	body: Body_Ptr,
	wheel: Suspension_Constraint_Ptr,
) -> Vec3 {
	local_right := lg.quaternion_mul_vector3(
		lg.quaternion_angle_axis(wheel.turn_angle, Vec3{0, 1, 0}),
		Vec3{1, 0, 0},
	)
	return body_local_to_world_vec(body, local_right)
}

wheel_get_forward_vec :: #force_inline proc(
	body: Body_Ptr,
	wheel: Suspension_Constraint_Ptr,
) -> Vec3 {
	local_right := lg.quaternion_mul_vector3(
		lg.quaternion_angle_axis(wheel.turn_angle, Vec3{0, 1, 0}),
		Vec3{0, 0, 1},
	)
	return body_local_to_world_vec(body, local_right)
}

body_get_shape_offset_local :: proc(body: Body_Ptr) -> (offset: Vec3) {
	return body.shape_offset
}

// Returns the shape's world position
// Shape can be offset from COM
body_get_shape_pos :: proc(body: Body_Ptr) -> Vec3 {
	offset := body_get_shape_offset_local(body)
	return body_local_to_world(body, offset)
}

shape_get_convex_local :: proc(
	sim_state: ^Sim_State,
	shape: Collision_Shape_Internal,
	allocator := context.temp_allocator,
) -> (
	mesh: collision.Convex,
) {
	switch s in shape {
	case Shape_Box:
		mesh = collision.box_to_convex(collision.Box{rad = s.size * 0.5}, allocator)
	case Internal_Shape_Convex:
		mesh = convex_container_get_mesh(&sim_state.convex_container, s.mesh)
		// TODO: make sure this works as intendent
		mesh = halfedge.copy_mesh(mesh, allocator)
	}

	return
}

Shapes_Iterator :: struct {
	sim_state: ^Sim_State,
	first_idx: i32,
	cur_idx:   i32,
	counter:   i32,
}

shapes_iterator :: proc(sim_state: ^Sim_State, first_shape: i32) -> (it: Shapes_Iterator) {
	it.sim_state = sim_state
	it.first_idx = first_shape
	it.cur_idx = it.first_idx
	return
}

shapes_iterator_next :: proc(
	it: ^Shapes_Iterator,
) -> (
	shape: ^Collision_Shape,
	idx: i32,
	ok: bool,
) {
	if it.cur_idx == it.first_idx && it.counter != 0 {
		return nil, it.counter, false
	}

	shape = &it.sim_state.shapes[it.cur_idx]
	idx = it.cur_idx
	ok = true

	it.cur_idx = shape.next_index
	it.counter += 1
	return
}

get_shape_by_index :: proc(
	sim_state: ^Sim_State,
	first_shape: i32,
	index: i32,
) -> (
	shape_idx: i32,
) {
	tracy.Zone()
	shape_idx = -1
	i := i32(0)
	it := shapes_iterator(sim_state, first_shape)
	for _, idx in shapes_iterator_next(&it) {
		if i == index {
			shape_idx = idx
			break
		}
		i += 1
	}
	return
}

body_get_convex_shape_world :: proc(
	sim_state: ^Sim_State,
	body: Body_Ptr,
	shape: Collision_Shape,
	allocator := context.temp_allocator,
) -> (
	mesh: collision.Convex,
) {
	tracy.Zone()
	body_transform := lg.matrix4_from_trs_f32(body_get_shape_pos(body), body.q, 1)
	shape_transform := lg.matrix4_translate(shape.rel_x) * lg.matrix4_from_quaternion(shape.rel_q)

	transform := body_transform * shape_transform
	mesh = shape_get_convex_local(sim_state, shape.inner_shape, allocator)
	halfedge.transform_mesh(&mesh, transform)
	return
}

body_get_convex_shape_by_linear_index_world :: proc(
	sim_state: ^Sim_State,
	body: Body_Ptr,
	shape_linear_idx: i32,
	allocator := context.temp_allocator,
) -> (
	mesh: collision.Convex,
) {
	tracy.Zone()
	if shape_linear_idx == 1 {
		log.debugf("linear idx: %v, body %v", shape_linear_idx, name.to_string(body.name))
	}
	shape_idx := get_shape_by_index(sim_state, body.shape, shape_linear_idx)
	return body_get_convex_shape_world(sim_state, body, sim_state.shapes[shape_idx], allocator)
}

body_get_convex_shapes_world :: proc(
	sim_state: ^Sim_State,
	body: Body_Ptr,
	allocator := context.temp_allocator,
) -> (
	meshes: []collision.Convex,
) {
	tracy.Zone()
	body_transform :=
		lg.matrix4_translate_f32(body_get_shape_pos(body)) * lg.matrix4_from_quaternion_f32(body.q)

	meshes_arr := make([dynamic]collision.Convex, context.temp_allocator)
	defer delete(meshes_arr)

	it := shapes_iterator(sim_state, body.shape)
	for shape in shapes_iterator_next(&it) {
		shape_transform :=
			lg.matrix4_translate(shape.rel_x) * lg.matrix4_from_quaternion(shape.rel_q)

		transform := body_transform * shape_transform
		mesh := shape_get_convex_local(sim_state, shape.inner_shape, allocator)
		halfedge.transform_mesh(&mesh, transform)
		append(&meshes_arr, mesh)
	}

	meshes = make([]collision.Convex, len(meshes_arr), allocator)
	copy(meshes, meshes_arr[:])
	return
}

level_geom_get_convex_shape :: proc(
	sim_state: ^Sim_State,
	level_geom: Level_Geom_Ptr,
	tri_idx: int,
	allocator := context.temp_allocator,
) -> (
	mesh: collision.Convex,
) {
	vertices, indices := get_level_geom_data(sim_state, level_geom.geometry)
	return collision.double_sided_triangle_to_convex(
		get_transformed_triangle(vertices, indices, tri_idx, level_geom.x, level_geom.q),
		allocator,
	)
}

input_shape_internal_get_aabb :: proc(shape: Input_Shape_Internal) -> (aabb: AABB) {
	switch s in shape {
	case Shape_Box:
		aabb.center = 0
		aabb.extent = s.size * 0.5
	case Shape_Convex:
		aabb.center = s.mesh.center
		aabb.extent = s.mesh.extent
	}

	return aabb
}

internal_shape_get_aabb :: proc(shape: Collision_Shape_Internal) -> (aabb: AABB) {
	switch s in shape {
	case Shape_Box:
		aabb.center = 0
		aabb.extent = s.size * 0.5
	case Internal_Shape_Convex:
		aabb.center = s.center
		aabb.extent = s.extent
	}
	return
}

shape_get_aabb :: proc(shape: Collision_Shape) -> (aabb: AABB) {
	local_aabb := internal_shape_get_aabb(shape.inner_shape)
	aabb = aabb_transform(local_aabb, shape.rel_x, shape.rel_q)
	return
}

rotate_extent :: proc(extent: Vec3, q: Quat) -> Vec3 {
	rotation_mat := lg.matrix3_from_quaternion(q)
	result := lg.abs(extent.xxx * rotation_mat[0])
	result += lg.abs(extent.yyy * rotation_mat[1])
	result += lg.abs(extent.zzz * rotation_mat[2])
	return result
}

aabb_transform :: proc(local_aabb: AABB, x: Vec3, q: Quat) -> (aabb: AABB) {
	aabb.center = lg.quaternion_mul_vector3(q, local_aabb.center) + x

	aabb.extent = rotate_extent(local_aabb.extent, q)

	return
}

input_shape_get_aabb :: proc(shapes: []Input_Shape) -> (aabb: AABB) {
	min, max: Vec3 = max(f32), min(f32)

	for shape in shapes {
		local_aabb := input_shape_internal_get_aabb(shape.inner_shape)
		aabb = aabb_transform(local_aabb, shape.rel_x, shape.rel_q)
		shape_min := aabb.center - aabb.extent
		shape_max := aabb.center + aabb.extent

		min, max = lg.min(shape_min, min), lg.max(shape_max, max)
	}

	aabb.center = (max + min) * 0.5
	aabb.extent = (max - min) * 0.5
	return
}

body_transform_shape_aabb :: proc(body: Body_Ptr, local_aabb: AABB) -> (aabb: AABB) {
	offset_world := body_local_to_world(body, body.shape_offset)
	aabb = aabb_transform(local_aabb, offset_world, body.q)

	return
}

body_get_aabb :: proc(body: Body_Ptr) -> (aabb: AABB) {
	local_aabb := body.shape_aabb

	return body_transform_shape_aabb(body, local_aabb)
}

rpm_to_angular_velocity :: proc(rpm: f32) -> f32 {
	return 2.0 * math.PI * (rpm / 60.0)
}

angular_velocity_to_rpm :: proc(w: f32) -> f32 {
	return (w / (2.0 * math.PI)) * 60.0
}