simple_skeleton_ik_3d.gd

@tool
extends Node
class_name SimpleSkeletonIK3D

var selected_bone_indices: PackedInt32Array
var bones := []
var bone_chain_length: float
var bone_lengths: PackedFloat32Array
var bone_origins: PackedVector3Array
var bone_directions: PackedVector3Array
var bone_rotations := [] # quaternions

var ik_target_origin: Vector3
var ik_pole_target_origin: Vector3

var root_bone: String
var tail_bone: String
@export var error_tolerance := 0.001
@export var incremental_solver := false
@export var pre_rotate := true  # Perform initial rotation of IK chain to point towards target
@export var target_node: Node
@export var pole_target_node: Node
@export var debug := false : set = _set_debug
@export var rotate_tail := true

var available_bone_names := []

var debug_geom: MeshInstance3D

func _get_property_list() -> Array:
  var available_bone_names_hint := ",".join(available_bone_names)
  return [
    {
      name = "root_bone",
      type = TYPE_STRING,
      hint = PROPERTY_HINT_ENUM,
      hint_string = available_bone_names_hint,
    },
    {
      name = "tail_bone",
      type = TYPE_STRING,
      hint = PROPERTY_HINT_ENUM,
      hint_string = available_bone_names_hint,
    },
  ]

func _notification(what: int) -> void:
  match what:
    NOTIFICATION_UNPARENTED:
      update_bones()
    NOTIFICATION_PARENTED:
      update_bones()

func get_parent_skeleton() -> Skeleton3D:
  var parent_node := get_parent()
  if parent_node and parent_node is Skeleton3D:
    return parent_node
  return null

func update_bones() -> void:
  var skeleton := get_parent_skeleton()
  if skeleton:
    if Engine.is_editor_hint():
      update_available_bone_names(skeleton)
    update_selected_bone_indices(skeleton)
    initialize_bone_chain(skeleton)

# Repopulate available_bone_names with the bone's from skeleton
func update_available_bone_names(skeleton: Skeleton3D) -> void:
  available_bone_names.clear()
  for index in range(skeleton.get_bone_count()):
    var bone_name: String = skeleton.get_bone_name(index)
    available_bone_names.append(bone_name)
  notify_property_list_changed()

# Repopulate bone_indices with the bone indices from the root
# bone to the tail bone of the selected IK chain.
func update_selected_bone_indices(skeleton: Skeleton3D) -> void:
  if not root_bone or not tail_bone:
    return
  var root_index := skeleton.find_bone(root_bone)
  var tail_index := skeleton.find_bone(tail_bone)
  if tail_index < 0 or root_index < 0:
    push_warning("Unable to find both end bones selected bone chain")
    return
  selected_bone_indices.clear()
  while tail_index != root_index:
    selected_bone_indices.push_back(tail_index)
    tail_index = skeleton.get_bone_parent(tail_index)
  selected_bone_indices.push_back(tail_index)
  selected_bone_indices.reverse()

func initialize_bone_chain(skeleton: Skeleton3D) -> void:
  bones.clear()
  bone_lengths.clear()
  bone_directions.clear()
  bone_rotations.clear()
  for b in selected_bone_indices:
    bones.append(skeleton.get_bone_global_rest(b))
  bones.append(Transform3D(bones[-1].basis, Vector3.ZERO)) # dummy tail bone for length calculation
  bone_chain_length = bones[0].origin.distance_to(bones[1].origin)
  for i in len(bones)-1:
    var bone_length := (bones[i].origin as Vector3).distance_to(bones[i+1].origin)
    bone_chain_length += bone_length
    bone_lengths.append(bone_length)
    bone_directions.append((bones[i+1].origin - bones[i].origin).normalized())
    bone_rotations.append(bones[i].basis.get_rotation_quaternion())

func to_vec2(v: Vector3) -> Vector2:
  return Vector2(v.z, v.y)

func to_vec3(v: Vector2) -> Vector3:
  return Vector3(0.0, v.y, v.x)

# func _process(delta: float) -> void:
#   _physics_process(delta)

func _process(_delta: float) -> void:
  var skeleton := get_parent_skeleton()
  if not skeleton:
    push_warning("no skeleton set")
    return
  if selected_bone_indices.size() < 3:
    push_warning("not enough bones in chain (minimum: 3): {}" % selected_bone_indices.size())
    return
  if not target_node:
    push_warning("no target node set")
    return
  fabrik_step(skeleton)


func debug_line(a: Vector3, b: Vector3, color: Color) -> void:
  var dbg := get_debug_mesh()
  if not dbg:
    return
  dbg.surface_begin(Mesh.PRIMITIVE_LINES)
  dbg.surface_set_color(color)
  dbg.surface_add_vertex(a)
  dbg.surface_add_vertex(b)
  dbg.surface_end()

func debug_bone_x_axis(color := Color.RED) -> void:
  var dbg := get_debug_mesh()
  if not dbg:
    return
  dbg.surface_begin(Mesh.PRIMITIVE_LINES)
  dbg.surface_set_color(color)
  for i in len(bones)-1:
    dbg.surface_add_vertex(bones[i].origin)
    dbg.surface_add_vertex(bones[i].origin + Vector3.LEFT * Transform3D(Basis(bone_rotations[i]), Vector3.ZERO) * 1.0)
  dbg.surface_end()

func fabrik_step(skeleton: Skeleton3D) -> void:
  var dbg := get_debug_mesh()
  if dbg:
    dbg.clear_surfaces()
  var num_joints := len(selected_bone_indices)
  # Collect current bone global pose transforms
  bones.resize(num_joints)
  bone_origins.resize(num_joints)
  for i in num_joints:
    if incremental_solver:
      bones[i] = skeleton.get_bone_global_pose(selected_bone_indices[i])
    else:
      bones[i] = skeleton.get_bone_global_rest(selected_bone_indices[i])
    bone_origins[i] = bones[i].origin
  # Collect current IK target positions
  ik_target_origin = skeleton.to_local(target_node.global_transform.origin)
  if pole_target_node:
    ik_pole_target_origin = skeleton.to_local(pole_target_node.global_transform.origin)

  # Perform an initial rotation of the entire chain toward the IK target
  # direction if enabled
  var pre_transform: Transform3D
  if pre_rotate:
    var chain_direction := (bone_origins[-1] - bone_origins[0]).normalized() as Vector3
    var target_direction := (ik_target_origin - bone_origins[0]).normalized() as Vector3
    var diff_rotation := Quaternion(target_direction, chain_direction)
    pre_transform = Transform3D(Basis.IDENTITY, -bone_origins[0])
    # if pole_target_node:
    #   pre_transform = Transform3D(Basis(Quaternion(rget_direction, Vector3.LEFT)), Vector3.ZERO) * pre_transform
    pre_transform = Transform3D(Basis(diff_rotation), bone_origins[0]) * pre_transform
    # pre_transform = Transform3D(Basis(diff_rotation), bone_origins[0]) * Transform3D(Basis.IDENTITY, -bone_origins[0])
    # if pole_target_node:
    #   pre_transform = pre_transform * Transform3D(Basis(Quaternion(target_direction, deg_to_rad(1.0))), Vector3.ZERO)
    for i in num_joints:
      bone_origins[i] = bone_origins[i] * pre_transform
      debug_line(bones[i].origin, bones[i].origin + Vector3.RIGHT * diff_rotation * 0.5, Color.DARK_RED)

  var root_distance_to_target := bone_origins[0].distance_to(ik_target_origin)
  if root_distance_to_target > bone_chain_length:
    var direction := (ik_target_origin - bone_origins[0]).normalized()
    for i in range(1, num_joints):
      bone_origins[i] = bone_origins[i-1] + direction * bone_lengths[i-1]
  else:
    var root_origin := bone_origins[0]
    var max_iterations := 5
    var iteration := 0
    while iteration < max_iterations:
      iteration += 1
      # Forward reaching phase
      bone_origins[-1] = ik_target_origin
      for i in range(num_joints-2, -1, -1):
        bone_origins[i] = bone_origins[i+1] \
            + (bone_lengths[i] / bone_origins[i+1].distance_to(bone_origins[i])) \
            * (bone_origins[i] - bone_origins[i+1])
      # Backward reaching phase
      bone_origins[0] = root_origin
      for i in range(num_joints-1):
        bone_origins[i+1] = bone_origins[i] \
            + (bone_lengths[i] / bone_origins[i+1].distance_to(bone_origins[i])) \
            * (bone_origins[i+1] - bone_origins[i])
      # Check if tail bone is sufficiently close to the target position
      if bone_origins[-1].distance_to(ik_target_origin) < error_tolerance:
        break
  if rotate_tail:
    var last_bone_index := num_joints-1
    bones[last_bone_index].origin = bone_origins[last_bone_index]
    bones[last_bone_index].basis = skeleton.global_transform.basis.inverse() * target_node.global_transform.basis * Basis(bone_rotations[last_bone_index])
    skeleton.set_bone_global_pose_override(selected_bone_indices[last_bone_index], bones[last_bone_index], 1.0, false)
  # Apply transformations to skeleton
  for i in num_joints-1:
    bones[i].origin = bone_origins[i]
    var bone_direction := (bone_origins[i+1] - bone_origins[i]).normalized()
    var target_rotation := Quaternion(bone_directions[i], bone_direction)
    if pre_transform:
      var bone_x_axis_direction: Quaternion = Quaternion(pre_transform.basis.inverse()) * bone_rotations[i]
      var roll_correction := Quaternion(target_rotation * Vector3.LEFT, bone_x_axis_direction * Vector3.LEFT)
      if i < num_joints-1:
        bones[i].basis = Basis(roll_correction * target_rotation * bone_rotations[i])
    else:
      bones[i].basis = Basis(target_rotation * bone_rotations[i])

    skeleton.set_bone_global_pose_override(selected_bone_indices[i], bones[i], 1.0, false)

  # Draw debug geometry
  # debug_bone_x_axis(Color.ORANGE)
  debug_line(bones[0].origin, ik_target_origin, Color.GREEN)
  if pole_target_node:
    debug_line(bones[0].origin, ik_pole_target_origin, Color.RED)

func pole_constraint() -> void:
  if pole_target_node:
    var limb_axis := (bone_origins[2] - bone_origins[0]).normalized()
    var pole_direction := (ik_pole_target_origin - bone_origins[0]).normalized()
    var bone_direction := (bone_origins[1] - bone_origins[0]).normalized()

    pole_direction = (pole_direction - limb_axis * pole_direction.dot(limb_axis)).normalized()
    bone_direction = (bone_direction - limb_axis * bone_direction.dot(limb_axis)).normalized()
    var angle := Quaternion(bone_direction, pole_direction).normalized()
    bone_origins[1] = angle * (bone_origins[1] - bone_origins[0]) + bone_origins[0]


func _set_debug(value: bool) -> void:
  debug = value
  if not Engine.is_editor_hint():
    return
  var skeleton := get_parent_skeleton()
  if not skeleton:
    return
  if value:
    if not debug_geom:
      debug_geom = MeshInstance3D.new()
      debug_geom.mesh = ImmediateMesh.new()
      var material := StandardMaterial3D.new()
      material['shading_mode'] = BaseMaterial3D.ShadingMode.SHADING_MODE_UNSHADED
      material['vertex_color_use_as_albedo'] = true
      material['cull_mode'] = BaseMaterial3D.CULL_DISABLED
      material['no_depth_test'] = true
      debug_geom['material_override'] = material
      skeleton.add_child(debug_geom)
  else:
    if debug_geom:
      skeleton.remove_child(debug_geom)
      debug_geom = null

func get_debug_mesh() -> Mesh:
  return debug_geom.mesh if debug_geom else null