extends ARVROrigin

# Provides different methods to calibrate the camera origin in VR setups
#
# HMD cameras (ARVRCamera node in Godot) assume the position and orientation
# of the HMD in the physical world, relative to the ARVROrigin node.
#
# The camera position is very unpredictable, though:
#
# - An HTC Vive player with seated setup will have the camera about 1 meter
#   above the position of the ARVROrigin
#
# - An HTC Vive player with standing setup will either have the camera 2 meters
#   above the ARVROrigin node (if he's actually standing) or 1 meter above
#   and up to 2 meters horizontally (X/Z) away from the center (as it's unlikely
#   that his chair is standing alone in the center of the play area)
#
# - An Oculus DK2 player looking straight ahead will likely have the camera
#   exactly at the ARVROrigin (since the head position gets calibrated as
#   being a the identity transform)
#
# I have no access to an Oculus CV1. It may act like the HTC Vive or the DK2.
#
# The VirtualRealityOrigin class offers helper methods to recenter the camera
# so that whatever position the player choses can have the HMD camera be
# at the identity transform relative to the **parent** of the ARVROrigin node.
#
# Suggested node setup:
#
#   o Observer              (can be anything, empty Spatial node recommended)
#     o PhysicalOrigin      (ARVROrigin + this script)
#       o EyeCamera         (ARVRCamera)
#
# If your game runs without VR, both the PhysicalOrigin and the EyeCamera will
# stay at their identity transforms, thus the whole thing acts like a vanilla camera.

# -------------------------------------------------------------------------------------------------

# Maximum framerate the averaging algorithm expects to encounter
# @remarks
#   If the framerate is higher than this, the sampling period will be shortened
#   by a bit. This will not cause a bad calibration, but the sampling will
#   finish in less than the time requested by the caller.
const MaximumAnticipatedFramerate = 90.0

# How long the automatic calibration on startup will run
const StartupCalibrationPeriod = 0.25

# How long the manual calibration will run
const ManualCalibrationPeriod = 0.5

# Current action of the VR observer
enum Action {
	# Nothing is happening
	Idle,
	# Calibration samples are being collected while fading out
	FadingOutCalibrating,
	# Fading back in after calibration
	FadingIn
}

# -------------------------------------------------------------------------------------------------

# Whether to automatically perform a calibration when the scene has loaded
# @remarks
#   This is recommended (thus, defaulting to true) due to the unpredictability
#   of the calibration center between different VR headsets.
export var CalibrateOnStartup = true

# Transform the HMD camera should have with the player's head neutral
# @remarks
#   Changing this doesn't do anything for the 
export var TargetTransform = Transform(
	Vector3(1.0, 0.0, 0.0),
	Vector3(0.0, 1.0, 0.0),
	Vector3(0.0, 0.0, 1.0),
	Vector3(0.0, 0.0, 0.0)
)

# -------------------------------------------------------------------------------------------------

# The camera through which the user is viewing the scene
var hmdCamera
# Manager class that handles the virtual reality setup etc.
var virtualRealityManager
# Camera controller that will be used to fade out during VR calibration
var cameraController

# Orientation samples captured during the averaging phase
var orientationSamples = []
# Translation samples captured during the averaging phase
var translationSamples = []
# Number of averaging samples that will be captured
var averagingSampleCount = 0

# Action currently being performed by calibration system
var currentAction = Action.Idle
# Seconds that have elapsed in the current calibration action
var elapsedActionSeconds = 0.0
# Total seconds the current calibration action will run for
var actionLengthSeconds = 0.0

# How many samples to skip due to scene warm-up effect
# @remarks
#   The first (possilbly first few) camera positions do not contain the proper physical
#   position of the HMD and would throw averaging off. So when CalibrateOnSceneLoad
#   is used, this is the amount of samples that will be skipped.
var warmupSamples = 5

# -------------------------------------------------------------------------------------------------

# Initializes the script component when the node is added to the scene
func _ready():
	InitializeIfNotYetDone()
	
	if self.virtualRealityManager == null:
		set_process(false)
		return
	
	# If we should calibrate on startup, either apply the existing calibration
	# (so that no every scene has to repeat the calibration process - if players
	# move their heads it's most likely during loading screens!) or
	# calibrate now if it's the first scene loaded by the game. 
	if self.CalibrateOnStartup:
		var calibrationTransform = self.virtualRealityManager.GetCalibrationTransform()
		if calibrationTransform == null:
			print(
				"%.3f No VR calibration, recalibrating if/when tracking data arrives" %
				(OS.get_ticks_msec() / 1000.0)
			)
			CalibrateToIdentity(StartupCalibrationPeriod)
		else:
			print("%.3f VR calibration data present" % (OS.get_ticks_msec() / 1000.0))
			applyCalibration(calibrationTransform)
			set_process(false)
	else: # No calibration desired
		ResetCalibration()
		set_process(false)

# -------------------------------------------------------------------------------------------------

# Called once per visual frame to update the visual state of the node
# @param  delta  Time that has passed since the last frame
func _process(delta):
	if self.warmupSamples > 0:
		self.warmupSamples -= 1
		return
	
	# Nothing going on? Stop the process callback
	if self.currentAction == Action.Idle:
		set_process(false)
		return
	
	# Fading back in after calibration? 
	if self.currentAction == Action.FadingIn:
		self.elapsedActionSeconds += delta
		if self.elapsedActionSeconds >= self.actionLengthSeconds:
			self.cameraController.OverrideFadeLevel = 1.0
			self.currentAction = Action.Idle
			set_process(false)
		else:
			self.cameraController.OverrideFadeLevel = (
				self.elapsedActionSeconds / self.actionLengthSeconds
			)
		return
	
	# If time is up, finish up the calibration
	self.elapsedActionSeconds += delta
	if self.elapsedActionSeconds >= self.actionLengthSeconds:
		finishCalibrationSampleCollection()
		self.elapsedActionSeconds = 0.0
		self.currentAction = Action.FadingIn
		return
	
	self.cameraController.OverrideFadeLevel = (
		1.0 - (self.elapsedActionSeconds / self.actionLengthSeconds)
	)
	
	# Collect an additional sample if there's still space in the sample list
	# (if the sample list is full, we don't finish the calibration yet because
	# we want to wait for the screen to be fully faded out before adjusting
	# the camera position to avoid disorienting the player).
	var maximumSampleCount = self.translationSamples.size()
	if self.averagingSampleCount < (maximumSampleCount - 1):
		self.translationSamples[self.averagingSampleCount] = (
			self.hmdCamera.translation # Local space translation
		)
		self.orientationSamples[self.averagingSampleCount] = (
			Quat(self.hmdCamera.transform.basis) # Local space orientation as Quaternion
		)
		self.averagingSampleCount += 1

# -------------------------------------------------------------------------------------------------

# Called to process any player input
# @param  event  Event describing the input performed by the player
func _input(var event):
	if event.is_action_pressed("vr_recalibrate"):
		if self.currentAction == Action.Idle:
			print("%.3f Player requested VR recalibration" % (OS.get_ticks_msec() / 1000.0))
			CalibrateToIdentity(ManualCalibrationPeriod)

# -------------------------------------------------------------------------------------------------

# Initializes the scene loader if this hasn't already happened
# @remarks
#   With this method, other components can avoid having to depend on an initialization
#   order that needs to be set up behind the scenes (and that would have no documentation
#   in an obvious place to indicate that)
func InitializeIfNotYetDone():
	if self.hmdCamera == null:
		self.hmdCamera = findHmdCamera()
		assert(self.hmdCamera != null)
	
	if self.virtualRealityManager == null:
		self.virtualRealityManager = get_node("/root/Systems/VirtualRealityManager")
		if self.virtualRealityManager != null:
			self.virtualRealityManager.InitializeIfNotYetDone()
	
	if self.cameraController == null:
		self.cameraController = get_node("..")
		assert(self.cameraController != null)

# -------------------------------------------------------------------------------------------------

# Resets the ARVR origin to its identity transform
# @remarks
#   It is unpredictable where the HMD camera will end up in relation to the ARVROrigin
#   node, so unless you are specifically aiming at a 'room scale setup', calling this
#   could have your camera end up 1-2 meters above the ARVROrigin and quite possibly
#   a meter on the X/Z plane, too.
func ResetCalibration():
	self.transform = self.TargetTransform

# -------------------------------------------------------------------------------------------------

# Recenters the ARVR origin so the current HMD camera position places the camera
# at the identity transform relative to the ARVROrigin node's parent.
# @param  averagingSeconds  Number of seconds to average HMD position over
# @remarks
#   It recommended to warn the user before calling this. If the user just looked
#   to the side or reached for his drink, that head position+orientation would
#   become the new neutral.
func CalibrateToIdentity(var averagingSeconds = 0.0):
	self.TargetTransform.basis = Basis(
		Vector3(1.0, 0.0, 0.0),
		Vector3(0.0, 1.0, 0.0),
		Vector3(0.0, 0.0, 1.0)
	)
	self.TargetTransform.origin = Vector3(0.0, 0.0, 0.0)
	
	if !self.virtualRealityManager.IsVirtualRealityActive():
		return
	
	if averagingSeconds < (1.0 / MaximumAnticipatedFramerate):
		var counterTransform = self.hmdCamera.transform.inverse()
		self.virtualRealityManager.StoreCalibrationTransform(counterTransform)
		applyCalibration(counterTransform)
	else:
		startCalibrationSampleCollection(averagingSeconds)

# -------------------------------------------------------------------------------------------------

# Recenters the ARVR origin so the current HMD camera position places the camera
# at the specified orientation and location relative to the ARVROrigin node's parent.
# @param  orientation       Orientation (Basis) the camera should end up in
# @param  translation       Translation (Vector3) the camera should end up at
# @param  averagingSeconds  Number of seconds to average HMD position over
# @remarks
#   It recommended to warn the user before calling this. If the user just looked
#   to the side or reached for his drink, that head position+orientation would
#   become the new neutral.
func CalibrateTo(var orientation, var translation, var averagingSeconds = 0.0):
	self.TargetTransform.basis = orientation
	self.TargetTransform.origin = translation
	
	if !self.virtualRealityManager.IsVirtualRealityActive():
		return
	
	if averagingSeconds < (1.0 / MaximumAnticipatedFramerate):
		var counterTransform = self.hmdCamera.transform.inverse()
		self.virtualRealityManager.StoreCalibrationTransform(counterTransform)
		applyCalibration(counterTransform)
	else:
		startCalibrationSampleCollection(averagingSeconds)

# -------------------------------------------------------------------------------------------------

# Returns whether the HMD is being calibrated right now
# @returns True if the HMD is currently being calibrated, false otherwise
func IsCalibrating():
	if self.virtualRealityManager.IsVirtualRealityActive():
		return (self.averagingSeconds > 0.0)
	else:
		return false

# -------------------------------------------------------------------------------------------------

# Looks for the HMD camera in the children of the ARVROrigin node
# @returns The HMD camera or null if not found
func findHmdCamera():
	for childNode in get_children():
		if childNode is ARVRCamera:
			return childNode
	
	return null

# -------------------------------------------------------------------------------------------------

# Sets up the calibration sample collection
# @param  averagingSeconds  Duration over which averaging samples will be collection
func startCalibrationSampleCollection(var averagingSeconds):
	var maximumSampleCount = averagingSeconds * MaximumAnticipatedFramerate
	
	self.currentAction = Action.FadingOutCalibrating
	self.elapsedActionSeconds = 0.0
	self.actionLengthSeconds = averagingSeconds
	self.orientationSamples.resize(maximumSampleCount)
	self.translationSamples.resize(maximumSampleCount)
	self.averagingSampleCount = 0
	
	print("%.3f VR calibration started" % (OS.get_ticks_msec() / 1000.0))
	set_process(true)

# -------------------------------------------------------------------------------------------------

# Applies the calibration after all averaging samples have been collected
func finishCalibrationSampleCollection():
	var averageOrientation = averageOrientations(
		self.orientationSamples, self.averagingSampleCount
	)
	var averageTranslation = averageTranslations(
		self.translationSamples, self.averagingSampleCount
	)
	
	# Build a transform that stores the average orientation and translation
	# of the HMD camera during the sampling duration
	var averageHmdTransform = Transform(averageOrientation)
	averageHmdTransform.origin = averageTranslation
	
	# Now recalibrate to the averaged head transform + desired neutral transform
	var counterTransform = averageHmdTransform.inverse()
	self.virtualRealityManager.StoreCalibrationTransform(counterTransform)
	applyCalibration(counterTransform)
	
	print("%.3f VR calibration completed" % (OS.get_ticks_msec() / 1000.0))

# -------------------------------------------------------------------------------------------------

# Applies a new HMD calibration transform
# @param  counterTransform  Transform that the physics HMDs transform in its neutral position
func applyCalibration(var counterTransform):
	self.transform = self.TargetTransform * counterTransform

# -------------------------------------------------------------------------------------------------

# Averages a set of quaterions
# @param  quaternionArray  Array of quaternions that will be averaged
# @param  count            Number of quaternions in the array
# @returns The average of all orientations in the quaternion array
static func averageOrientations(var quaternionArray, var count):
	var summedOrientations = Quat(0.0, 0.0, 0.0, 0.0)
	
	while count > 0:
		count -= 1
		
		summedOrientations.x += quaternionArray[count].x
		summedOrientations.y += quaternionArray[count].y
		summedOrientations.z += quaternionArray[count].z
		summedOrientations.w += quaternionArray[count].w
	
	return summedOrientations.normalized()

# -------------------------------------------------------------------------------------------------

# Averages a set of translations
# @param  vectorArray  Array of vectors that will be averaged
# @param  count        Number of vectors in the array
# @returns The average of all translations in the vector array
static func averageTranslations(var vectorArray, var count):
	var summedTranslations = Vector3(0.0, 0.0, 0.0)
	var totalTranslationCount = count
	
	while count > 0:
		count -= 1
		summedTranslations += vectorArray[count]
	
	return (summedTranslations / totalTranslationCount)