from pykinect_azure.k4abt import _k4abt
from pykinect_azure.k4abt.frame import Frame
from pykinect_azure.k4abt._k4abtTypes import k4abt_tracker_default_configuration
from pykinect_azure.k4a.device import Device
from pykinect_azure.k4a._k4atypes import K4A_WAIT_INFINITE
from pykinect_azure.utils import get_k4abt_lite_model_path

class Tracker:
	def __init__(self, calibration, model_type):

		self.calibration = calibration
		self._handle = self.create(model_type)
		self.frame = None

	def __del__(self):

		self.destroy()

	def is_valid(self):
		return self._handle

	def is_frame_initialized(self):
		return self.frame

	def handle(self):
		return self._handle

	def destroy(self):
		if self.is_valid():
			_k4abt.k4abt_tracker_destroy(self._handle)
			self._handle = None

	def update(self, capture=None, timeout_in_ms=K4A_WAIT_INFINITE):
		# Add capture to the body tracker processing queue
		if capture:
			self.enqueue_capture(capture.handle(), timeout_in_ms)
		else:
			self.enqueue_capture(Device.capture.handle(), timeout_in_ms)

		return self.pop_result(timeout_in_ms)

	def enqueue_capture(self, capture_handle, timeout_in_ms=K4A_WAIT_INFINITE):
		_k4abt.VERIFY(_k4abt.k4abt_tracker_enqueue_capture(self._handle, capture_handle, timeout_in_ms), "Body tracker capture enqueue failed!")

	def pop_result(self, timeout_in_ms=K4A_WAIT_INFINITE):

		if self.is_frame_initialized():
			self.frame.release()
			_k4abt.VERIFY(_k4abt.k4abt_tracker_pop_result(self._handle, self.frame.handle(), timeout_in_ms), "Body tracker get body frame failed!")
		else:
			frame_handle = _k4abt.k4abt_frame_t()
			_k4abt.VERIFY(_k4abt.k4abt_tracker_pop_result(self._handle, frame_handle, timeout_in_ms), "Body tracker get body frame failed!")
			self.frame = Frame(frame_handle, self.calibration)

		return self.frame

	def set_temporal_smoothing(self, smoothing_factor):
		_k4abt.k4abt_tracker_set_temporal_smoothing(self._handle, smoothing_factor)

	def shutdown(self):
		_k4abt.k4abt_tracker_shutdown(self._handle)

	def create(self, model_type):

		tracker_config = self.get_tracker_configuration(model_type)

		tracker_handle = _k4abt.k4abt_tracker_t()
		_k4abt.VERIFY(_k4abt.k4abt_tracker_create(self.calibration.handle(), tracker_config, tracker_handle), "Body tracker initialization failed!")
		
		return tracker_handle

	def get_tracker_configuration(self, model_type):
		tracker_config = k4abt_tracker_default_configuration

		if model_type == _k4abt.K4ABT_LITE_MODEL:
			tracker_config.model_path = get_k4abt_lite_model_path()

		return tracker_config


