import heapq

class DynamicOpticalSystem(object):
	'''A dynamically varying optical system.

	This can be used as a starting point for an implementation of an adaptive optics system. The optics
	are changed as specific moments in time, which can be scheduled by submitting a callback. The time
	points are bridged using the integrate function, which is written by the derived class. This
	function should integrate on each detector without reading out the detector. Detector readouts or
	DM changes should be implemented as callbacks.
	'''
	def __init__(self):
		self.callbacks = []
		self.t = 0
		self.callback_counter = 0

	def evolve_until(self, t):
		'''Evolve the optical system until time `t`.

		Parameters
		----------
		t : scalar
			The point in time to which to simulate this optical system. Callbacks and integrations
			are handled along the way.
		'''
		if t < self.t:
			raise ValueError('Backwards evolution is not allowed.')

		end = False
		while not end:
			t_next, _, callback = self.callbacks[0]
			if t_next < t:
				integration_time = t_next - self.t
				heapq.heappop(self.callbacks)
			else:
				integration_time = t - self.t
				end = True

			# Avoid multiple expensive integrations if we have multiple callbacks
			# at the same time.
			if integration_time > 1e-6:
				self.integrate(integration_time)
				self.t += integration_time

			if not end:
				callback()

	def integrate(self, integration_time):
		'''Integrate the current optical system for a certain integration time.

		This function should be implemented by a user in a derived class.

		Parameters
		----------
		integration_time : scalar
			The integration time with which to integrate the optical system.
		'''
		pass

	def add_callback(self, t, callback):
		'''Add a callback to the list of callbacks.

		This function can even be called during handling of a callback. This is
		especially useful when implementing periodic callbacks: a callback can
		reinsert itself at a later time at the end of handling of the callback.

		Parameters
		----------
		t : scalar
			The time at which to call the callback.
		callback : function
			The function to call at time `t`. This function should have no arguments.
		'''
		# Callback counter is to avoid comparison of callback functions
		# if the times are equal, as comparison of functions is not allowed.
		heapq.heappush(self.callbacks, (t, self.callback_counter, callback))
		self.callback_counter += 1
