import tensorflow as tf
import numpy as np
import time
import cv2
from hitnet.utils_hitnet import *


drivingStereo_config = CameraConfig(0.546, 1000)

class HitNet():

	def __init__(self, model_path, model_type=ModelType.eth3d, camera_config=drivingStereo_config):

		self.fps = 0
		self.timeLastPrediction = time.time()
		self.frameCounter = 0
		self.camera_config = camera_config

		# Initialize model
		self.model = self.initialize_model(model_path, model_type)

	def __call__(self, left_img, right_img):

		return self.estimate_disparity(left_img, right_img)

	def initialize_model(self, model_path, model_type):

		self.model_type = model_type

		with tf.io.gfile.GFile(model_path, "rb") as f:
			graph_def = tf.compat.v1.GraphDef()
			loaded = graph_def.ParseFromString(f.read())

		# Wrap frozen graph to ConcreteFunctions
		if self.model_type == ModelType.flyingthings:
			model = wrap_frozen_graph(graph_def=graph_def,
										inputs="input:0",
										outputs=["reference_output_disparity:0","secondary_output_disparity:0"])

		else:
			model = wrap_frozen_graph(graph_def=graph_def,
										inputs="input:0",
										outputs="reference_output_disparity:0")

		return model

	def estimate_disparity(self, left_img, right_img):

		input_tensor = self.prepare_input(left_img, right_img)

		# Perform inference on the image
		if self.model_type == ModelType.flyingthings:
			left_disparity, right_disparity = self.inference(input_tensor)
			self.disparity_map = left_disparity
		else:
			self.disparity_map = self.inference(input_tensor)

		return self.disparity_map

	def get_depth(self):
		return self.camera_config.f*self.camera_config.baseline/self.disparity_map

	def prepare_input(self, left_img, right_img):

		if (self.model_type == ModelType.eth3d):

			# Shape (1, None, None, 2)
			left_img = cv2.cvtColor(left_img, cv2.COLOR_BGR2GRAY)
			right_img = cv2.cvtColor(right_img, cv2.COLOR_BGR2GRAY)

			left_img = np.expand_dims(left_img,2)
			right_img = np.expand_dims(right_img,2)
			combined_img = np.concatenate((left_img, right_img), axis=-1) / 255.0
		else:
			# Shape (1, None, None, 6)
			left_img = cv2.cvtColor(left_img, cv2.COLOR_BGR2RGB)
			right_img = cv2.cvtColor(right_img, cv2.COLOR_BGR2RGB)
			combined_img = np.concatenate((left_img, right_img), axis=-1) / 255.0


		return tf.convert_to_tensor(np.expand_dims(combined_img, 0), dtype=tf.float32)
		
	def inference(self, input_tensor):
		output = self.model(input_tensor)

		return np.squeeze(output)



	







[((15, 28, 15, 39), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((86, 9, 86, 27), 'numpy.squeeze', 'np.squeeze', ({(86, 20, 86, 26): 'output'}, {}), '(output)', True, 'import numpy as np\n'), ((30, 7, 30, 42), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', ({(30, 25, 30, 35): 'model_path', (30, 37, 30, 41): '"""rb"""'}, {}), "(model_path, 'rb')", True, 'import tensorflow as tf\n'), ((31, 15, 31, 38), 'tensorflow.compat.v1.GraphDef', 'tf.compat.v1.GraphDef', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((68, 14, 68, 56), 'cv2.cvtColor', 'cv2.cvtColor', ({(68, 27, 68, 35): 'left_img', (68, 37, 68, 55): 'cv2.COLOR_BGR2GRAY'}, {}), '(left_img, cv2.COLOR_BGR2GRAY)', False, 'import cv2\n'), ((69, 15, 69, 58), 'cv2.cvtColor', 'cv2.cvtColor', ({(69, 28, 69, 37): 'right_img', (69, 39, 69, 57): 'cv2.COLOR_BGR2GRAY'}, {}), '(right_img, cv2.COLOR_BGR2GRAY)', False, 'import cv2\n'), ((71, 14, 71, 40), 'numpy.expand_dims', 'np.expand_dims', ({(71, 29, 71, 37): 'left_img', (71, 38, 71, 39): '2'}, {}), '(left_img, 2)', True, 'import numpy as np\n'), ((72, 15, 72, 42), 'numpy.expand_dims', 'np.expand_dims', ({(72, 30, 72, 39): 'right_img', (72, 40, 72, 41): '2'}, {}), '(right_img, 2)', True, 'import numpy as np\n'), ((76, 14, 76, 55), 'cv2.cvtColor', 'cv2.cvtColor', ({(76, 27, 76, 35): 'left_img', (76, 37, 76, 54): 'cv2.COLOR_BGR2RGB'}, {}), '(left_img, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((77, 15, 77, 57), 'cv2.cvtColor', 'cv2.cvtColor', ({(77, 28, 77, 37): 'right_img', (77, 39, 77, 56): 'cv2.COLOR_BGR2RGB'}, {}), '(right_img, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((81, 30, 81, 61), 'numpy.expand_dims', 'np.expand_dims', ({(81, 45, 81, 57): 'combined_img', (81, 59, 81, 60): '(0)'}, {}), '(combined_img, 0)', True, 'import numpy as np\n'), ((73, 18, 73, 64), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((78, 18, 78, 64), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n')]