extends SceneTree

## Visual regression capture for the local Titanfall motion lab rigs.
##
## godot --path . --windowed --resolution 800x800 \
##   --script res://debug/titanfall_motion_capture.gd -- <output-dir> [model.glb]

const DEFAULT_MODEL := \
	"res://assets/characters/local_reference/titanfall_motion/mannequin.glb"
const SPEED := 13.0

var _out_dir := "tmp_titanfall_motion_capture"
var _model_path := DEFAULT_MODEL
var _model: SkinnedPlayerModel


func _initialize() -> void:
	var args := OS.get_cmdline_user_args()
	if not args.is_empty():
		_out_dir = args[0]
	if args.size() > 1:
		_model_path = args[1]
	if not _out_dir.is_absolute_path():
		_out_dir = ProjectSettings.globalize_path(_out_dir)
	DirAccess.make_dir_recursive_absolute(_out_dir)
	call_deferred("_run")


func _run() -> void:
	var scene := Node3D.new()
	root.add_child(scene)
	current_scene = scene
	var environment_node := WorldEnvironment.new()
	var environment := Environment.new()
	environment.background_mode = Environment.BG_COLOR
	environment.background_color = Color(0.08, 0.09, 0.12)
	environment.ambient_light_source = Environment.AMBIENT_SOURCE_COLOR
	environment.ambient_light_color = Color(0.78, 0.82, 0.95)
	environment.ambient_light_energy = 1.3
	environment.tonemap_mode = Environment.TONE_MAPPER_FILMIC
	environment_node.environment = environment
	scene.add_child(environment_node)
	var sun := DirectionalLight3D.new()
	sun.rotation_degrees = Vector3(-35.0, 35.0, 0.0)
	sun.light_energy = 2.2
	scene.add_child(sun)
	var camera := Camera3D.new()
	scene.add_child(camera)
	camera.current = true
	camera.global_position = Vector3(0.0, 1.15, 3.0)
	camera.look_at(Vector3(0.0, 0.9, 0.0), Vector3.UP)

	var body := CharacterBody3D.new()
	scene.add_child(body)
	_model = SkinnedPlayerModel.new()
	_model.model_path = _model_path
	_model.first_person_mode = false
	_model.shadows_only = false
	body.add_child(_model)
	for _frame in 60:
		await process_frame
	if not _model.loaded:
		printerr("titanfall_motion_capture: model failed to load")
		quit(1)
		return

	await _drive("ground", 0.0, 90)
	_snap("idle")
	await _wall(-1.0, Vector3.RIGHT)
	_snap("wall_left")
	await _drive("ground", 0.0, 90)
	await _wall(1.0, Vector3.LEFT)
	_snap("wall_right")

	_model.set_wall_side(0.0)
	_model.set_wall_surface(Vector3.ZERO, Vector3.ZERO)
	_model.set_wall_glide_motion(Vector3.ZERO)
	await _drive("ground", 0.0, 1)
	_snap("post_right_01")
	await _drive("ground", 0.0, 19)
	_snap("post_right_20")
	await _drive("ground", 0.0, 70)
	_snap("post_right_idle")

	# Speed-scaled traversal entry: preserve three early frames so a hard cut is
	# visually obvious in review while the final wall pose remains unchanged.
	body.velocity = Vector3(0.0, -3.0, -SPEED)
	await _drive("air", SPEED, 20)
	_model.set_wall_side(1.0)
	_model.set_wall_glide_motion(Vector3(0.0, 0.0, -SPEED))
	await _drive("wall_run", SPEED, 1)
	_snap("wall_transition_01")
	await _drive("wall_run", SPEED, 3)
	_snap("wall_transition_04")
	await _drive("wall_run", SPEED, 6)
	_snap("wall_transition_10")

	# The stopping one-shot is driven while speed physically decays instead of
	# leaving the high-rate Sprint cycle active until an instant zero.
	_model.set_locomotion(0.0, 1.0, 0.0)
	await _drive("ground", 11.0, 45)
	_snap("stop_00_run")
	for sample in [["stop_01", 10.0, 1], ["stop_06", 4.5, 5],
			["stop_14", 1.2, 8]]:
		await _drive("stop", float(sample[1]), int(sample[2]))
		_snap(String(sample[0]))
	await _drive("ground", 0.0, 18)
	_snap("stop_32_idle")

	# Successive hops select source and offline-mirrored source performances.
	body.velocity = Vector3(0.0, 8.5, -8.0)
	await _drive("air", 8.0, 12)
	_snap("jump_normal")
	await _drive("ground", 0.0, 24)
	body.velocity = Vector3(0.0, 8.5, -8.0)
	await _drive("air_alt", 8.0, 12)
	_snap("jump_alternate")
	quit()


func _wall(side: float, normal: Vector3) -> void:
	_model.set_wall_side(side)
	_model.set_wall_surface(normal, Vector3.ZERO)
	_model.set_wall_glide_motion(Vector3(0.0, 0.0, -SPEED))
	await _drive("wall_run", SPEED, 90)


func _drive(state: String, speed: float, frames: int) -> void:
	for _frame in frames:
		_model.update_state(state, speed, false)
		await process_frame


func _snap(name: String) -> void:
	var image := root.get_viewport().get_texture().get_image()
	var path := _out_dir.path_join(name + ".png")
	image.save_png(path)
	print("titanfall_motion_capture: saved %s metrics=%s" % [
		path, _pose_metrics()])


func _pose_metrics() -> Dictionary:
	var roles: Dictionary = _model._rig_info.get("roles", {})
	var chain := [_model._role_bone("hips", ["hips", "Hips", "pelvis"])]
	var neck_name := String(roles.get("neck", ""))
	var head_name := String(roles.get("head", ""))
	for role in roles.get("spine", []):
		var name := String(role)
		if name == neck_name or name == head_name:
			continue
		chain.append(RigRoles.find_imported_bone(_model.skeleton, name))
	var max_torso_joint_angle := 0.0
	var max_torso_joint := ""
	for index in range(1, chain.size()):
		if chain[index - 1] < 0 or chain[index] < 0:
			continue
		var parent_pose := _model.skeleton.get_bone_global_pose(chain[index - 1])
		var child_pose := _model.skeleton.get_bone_global_pose(chain[index])
		var parent_rest := _model.skeleton.get_bone_global_rest(chain[index - 1])
		var child_rest := _model.skeleton.get_bone_global_rest(chain[index])
		var posed_relative := parent_pose.basis.inverse() * child_pose.basis
		var rest_relative := parent_rest.basis.inverse() * child_rest.basis
		var joint_angle := posed_relative.get_rotation_quaternion().angle_to(
			rest_relative.get_rotation_quaternion())
		if joint_angle > max_torso_joint_angle:
			max_torso_joint_angle = joint_angle
			max_torso_joint = _model.skeleton.get_bone_name(chain[index])
	var max_leg_horizontal_ratio := 0.0
	var knee_rise := {"L": 0.0, "R": 0.0}
	for side in ["L", "R"]:
		var thigh := _model._role_bone("thigh." + side,
			["thigh." + side, "Left thigh" if side == "L" else "Right thigh"])
		var shin := _model._role_bone("shin." + side,
			["shin." + side, "Left shin" if side == "L" else "Right shin"])
		if thigh < 0 or shin < 0:
			continue
		var segment := _model.skeleton.get_bone_global_pose(shin).origin \
			- _model.skeleton.get_bone_global_pose(thigh).origin
		segment = _model.skeleton.global_transform.basis * segment
		knee_rise[side] = segment.y
		if segment.length() > 0.0001:
			max_leg_horizontal_ratio = maxf(max_leg_horizontal_ratio,
				Vector2(segment.x, segment.z).length() / segment.length())
	return {
		"torso_joint_angle": max_torso_joint_angle,
		"torso_joint": max_torso_joint,
		"leg_horizontal_ratio": max_leg_horizontal_ratio,
		"left_knee_segment_y": knee_rise["L"],
		"right_knee_segment_y": knee_rise["R"],
		"knee_lead": float(knee_rise["R"]) - float(knee_rise["L"]),
		"clip": _model.current_clip_debug(),
		"crossfade": _model._state_trans.xfade_time,
	}