import numpy as np
from numpy import expand_dims
from keras.models import load_model
from keras.preprocessing.image import load_img
from keras.preprocessing.image import img_to_array
from keras.backend import resize_images
from matplotlib import pyplot
from matplotlib.patches import Rectangle
import cv2

# load and prepare an image
def load_image_pixels(filename, shape):
    # load the image to get its shape
    image = load_img(filename)
    width, height = image.size
    # load the image with the required size
    image = load_img(filename, target_size=shape)
    # convert to numpy array
    image = img_to_array(image)
    # scale pixel values to [0, 1]
    image = image.astype('float32')
    image /= 255.0
    # add a dimension so that we have one sample
    image = expand_dims(image, 0)
    return image, width, height

def load_image_stream(frame, shape):
    width, height = frame.shape[1], frame.shape[0]
    frame = cv2.resize(frame, (shape[0], shape[1]), interpolation = cv2.INTER_LINEAR) 
    # scale pixel values to [0, 1]
    frame = frame.astype('float32')
    frame /= 255.0
    # add a dimension so that we have one sample
    frame = expand_dims(frame, 0)
    return frame, width, height

class BoundBox:
	def __init__(self, xmin, ymin, xmax, ymax, objness = None, classes = None):
		self.xmin = xmin
		self.ymin = ymin
		self.xmax = xmax
		self.ymax = ymax
		self.objness = objness
		self.classes = classes
		self.label = -1
		self.score = -1
 
	def get_label(self):
		if self.label == -1:
			self.label = np.argmax(self.classes)
 
		return self.label
 
	def get_score(self):
		if self.score == -1:
			self.score = self.classes[self.get_label()]
 
		return self.score

def _sigmoid(x):
	return 1. / (1. + np.exp(-x))

def decode_netout(netout, anchors, obj_thresh, net_h, net_w):
	grid_h, grid_w = netout.shape[:2]
	nb_box = 3
	netout = netout.reshape((grid_h, grid_w, nb_box, -1))
	nb_class = netout.shape[-1] - 5
	boxes = []
	netout[..., :2]  = _sigmoid(netout[..., :2])
	netout[..., 4:]  = _sigmoid(netout[..., 4:])
	netout[..., 5:]  = netout[..., 4][..., np.newaxis] * netout[..., 5:]
	netout[..., 5:] *= netout[..., 5:] > obj_thresh
 
	for i in range(grid_h*grid_w):
		row = i / grid_w
		col = i % grid_w
		for b in range(nb_box):
			# 4th element is objectness score
			objectness = netout[int(row)][int(col)][b][4]
			if(objectness.all() <= obj_thresh): continue
			# first 4 elements are x, y, w, and h
			x, y, w, h = netout[int(row)][int(col)][b][:4]
			x = (col + x) / grid_w # center position, unit: image width
			y = (row + y) / grid_h # center position, unit: image height
			w = anchors[2 * b + 0] * np.exp(w) / net_w # unit: image width
			h = anchors[2 * b + 1] * np.exp(h) / net_h # unit: image height
			# last elements are class probabilities
			classes = netout[int(row)][col][b][5:]
			box = BoundBox(x-w/2, y-h/2, x+w/2, y+h/2, objectness, classes)
			boxes.append(box)
	return boxes

def correct_yolo_boxes(boxes, image_h, image_w, net_h, net_w):
	new_w, new_h = net_w, net_h
	for i in range(len(boxes)):
		x_offset, x_scale = (net_w - new_w)/2./net_w, float(new_w)/net_w
		y_offset, y_scale = (net_h - new_h)/2./net_h, float(new_h)/net_h
		boxes[i].xmin = int((boxes[i].xmin - x_offset) / x_scale * image_w)
		boxes[i].xmax = int((boxes[i].xmax - x_offset) / x_scale * image_w)
		boxes[i].ymin = int((boxes[i].ymin - y_offset) / y_scale * image_h)
		boxes[i].ymax = int((boxes[i].ymax - y_offset) / y_scale * image_h)

def _interval_overlap(interval_a, interval_b):
	x1, x2 = interval_a
	x3, x4 = interval_b
	if x3 < x1:
		if x4 < x1:
			return 0
		else:
			return min(x2,x4) - x1
	else:
		if x2 < x3:
			 return 0
		else:
			return min(x2,x4) - x3

def bbox_iou(box1, box2):
	intersect_w = _interval_overlap([box1.xmin, box1.xmax], [box2.xmin, box2.xmax])
	intersect_h = _interval_overlap([box1.ymin, box1.ymax], [box2.ymin, box2.ymax])
	intersect = intersect_w * intersect_h
	w1, h1 = box1.xmax-box1.xmin, box1.ymax-box1.ymin
	w2, h2 = box2.xmax-box2.xmin, box2.ymax-box2.ymin
	union = w1*h1 + w2*h2 - intersect
	return float(intersect) / union
 
def do_nms(boxes, nms_thresh):
	if len(boxes) > 0:
		nb_class = len(boxes[0].classes)
	else:
		return
	for c in range(nb_class):
		sorted_indices = np.argsort([-box.classes[c] for box in boxes])
		for i in range(len(sorted_indices)):
			index_i = sorted_indices[i]
			if boxes[index_i].classes[c] == 0: continue
			for j in range(i+1, len(sorted_indices)):
				index_j = sorted_indices[j]
				if bbox_iou(boxes[index_i], boxes[index_j]) >= nms_thresh:
					boxes[index_j].classes[c] = 0

# get all of the results above a threshold
def get_boxes(boxes, labels, thresh):
	v_boxes, v_labels, v_scores, x_mins = list(), list(), list(), list()
	# enumerate all boxes
	for box in boxes:
		# enumerate all possible labels
		for i in range(len(labels)):
			# check if the threshold for this label is high enough
			if box.classes[i] > thresh:
				v_boxes.append(box)
				v_labels.append(labels[i])
				v_scores.append(box.classes[i]*100)
				x_mins.append(box.xmin)
				# don't break, many labels may trigger for one box
	unique = list(np.unique(x_mins))
	chosen_indexes = []
	for i in range(len(unique)):
		indexes = [j for j, e in enumerate(x_mins) if e == unique[i]]
		scores = [v_scores[k] for k in indexes]
		chosen_indexes.append(indexes[scores.index(max(scores))])
	return [v_boxes[l] for l in chosen_indexes], [v_labels[l] for l in chosen_indexes], [v_scores[l] for l in chosen_indexes]

# draw all results
def draw_boxes(filename, v_boxes, v_labels, v_scores):
	# load the image
	data = pyplot.imread(filename)
	# plot the image
	pyplot.imshow(data)
	pyplot.axis("off")
	# get the context for drawing boxes
	ax = pyplot.gca()
	# plot each box
	for i in range(len(v_boxes)):
		box = v_boxes[i]
		# get coordinates
		y1, x1, y2, x2 = box.ymin, box.xmin, box.ymax, box.xmax
		# calculate width and height of the box
		width, height = x2 - x1, y2 - y1
		if v_labels[i] == "with_mask":
			rect = Rectangle((x1, y1), width, height, fill=False, color = "green")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "green")
		elif v_labels[i] == "without_mask":
			rect = Rectangle((x1, y1), width, height, fill=False, color = "red")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "red")
		else:
			rect = Rectangle((x1, y1), width, height, fill=False, color = "orange")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "orange")
	# show the plot
	pyplot.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
	pyplot.margins(0,0)
	pyplot.show()

def draw_boxes_stream(frame, v_boxes, v_labels, v_scores):
	# plot the image
	pyplot.imshow(frame)
	pyplot.axis("off")
	# get the context for drawing boxes
	ax = pyplot.gca()
	# plot each box
	for i in range(len(v_boxes)):
		box = v_boxes[i]
		# get coordinates
		y1, x1, y2, x2 = box.ymin, box.xmin, box.ymax, box.xmax
		# calculate width and height of the box
		width, height = x2 - x1, y2 - y1
		if v_labels[i] == "with_mask":
			rect = Rectangle((x1, y1), width, height, fill=False, color = "green")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "green")
		elif v_labels[i] == "without_mask":
			rect = Rectangle((x1, y1), width, height, fill=False, color = "red")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "red")
		else:
			rect = Rectangle((x1, y1), width, height, fill=False, color = "orange")
			ax.add_patch(rect)
			label = "%s (%.3f)" % (v_labels[i], v_scores[i])
			pyplot.text(x1, y1, label, color = "orange")
	# show the plot
	pyplot.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
	pyplot.margins(0,0)
	pyplot.savefig("frame_det.png");pyplot.clf()
