import tensorflow as tf
import numpy as np
import cv2
import time
import pdb
import os




def nms(output, output1, output2, confThreshold=0.5,nmsThreshold=0.3):
	scores = output2
	classes = output1
	print(output1.shape)
	print(output2.shape)

	conf_mask = scores > confThreshold
	valConf = np.nonzero(conf_mask)[0]
	filteredBoxes = []
	if len(valConf) == 0:
		return filteredBoxes
	# print(valConf)
	validBoxes = output[valConf,:]
	classes = output1[valConf]
	validConfidence = output2[valConf]
	print(validBoxes)
	# print(classes)
	# print(validConfidence)
	unique_classes = list(set(classes.reshape(-1)))
	top = np.copy(validBoxes[:,0])
	left = np.copy(validBoxes[:,1])
	down = np.copy(validBoxes[:,2])
	right = np.copy(validBoxes[:,3])
	print(left.shape)
	for i in range(len(left)):
		left[i] = max(min(left[i],299),0)
		top[i] = max(min(top[i],299),0)
		right[i] = max(min(right[i],299),0)
		down[i] = max(min(down[i],299),0)
	# left = max(min(left,299),0)
	# top = max(min(top,299),0)
	# right = max(min(right,299),0)
	# down = max(min(down,299),0)
	validBoxes[:,0] = left
	validBoxes[:,1] = top
	validBoxes[:,2] = right
	validBoxes[:,3] = down
	print(validBoxes)
	boxAttr = validBoxes
	boxAttr[:,2] = boxAttr[:,2]- boxAttr[:,0]
	boxAttr[:,3] = boxAttr[:,3]- boxAttr[:,1]
	print(boxAttr)
	for cl in unique_classes :
		classMask = np.nonzero((classes == cl))[0]
		ClassBoxes = boxAttr[classMask,:]
		ClassConfidence = validConfidence[classMask]
		indices = cv2.dnn.NMSBoxes(ClassBoxes.tolist(),ClassConfidence.tolist(),confThreshold, nmsThreshold)
		for i in indices :
			ClassRes = [cl]
			ClassRes.extend(ClassConfidence[i])
			ClassRes.extend(ClassBoxes[i].tolist()[0])
			filteredBoxes.append(ClassRes)
	return filteredBoxes



def get_face_roi_test_newnms(net,img_batch,orig_h,orig_w,Scaling):
    sess = net['sess']			
    inp = net['inp']
    out = net['out']
    out1 = net['out1']
    out2 = net['out2']
    out3 = net['out3']
    feed_dict = {inp : np.concatenate(img_batch, 0)}

    output = sess.run(out, feed_dict)
    output1 = sess.run(out1, feed_dict)
    output2 = sess.run(out2, feed_dict)
    output3 = sess.run(out3, feed_dict)

    print(output.shape)
    print(output1.shape)
    print(output2.shape)
    print(output3.shape)
    # print(output)
    output = output[0]
    output1 = output1[0]
    output2 = output2[0]
    print(output)
    output = output*300
    # print(output)
    filteredBoxes = nms(output, output1, output2, net['min_prob'],net['min_iou'])
    f_ = np.array(filteredBoxes)
    print(f_)
    f = np.copy(f_[:,2:])
    print(f)
    f[:,0] = f[:,0]*(1280/300)
    f[:,2] = f[:,2]*(1280/300)
    f[:,1] = f[:,1]*(720/300)
    f[:,3] = f[:,3]*(720/300)
    # print(f[:,2]*Scaling)
    print(f)
    


def init_tf(GPU_ID, fd_modelfile, batch_size, model_r, model_c, debug):#used in test functions
    device_name = "GPU:"+str(GPU_ID)
    with tf.device(device_name):
        #init config
        cfg = dict({'allow_soft_placement': True,'log_device_placement': False})
        utility = 0.9
        cfg['gpu_options'] = tf.GPUOptions(per_process_gpu_memory_fraction = utility)
        cfg['allow_soft_placement'] = True
        cfg['device_count'] = {'GPU': 1}
        cfg['use_per_session_threads'] = True
        sess = tf.Session(config = tf.ConfigProto(**cfg))
        #load model
        with tf.gfile.FastGFile(fd_modelfile, "rb") as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            tf.import_graph_def(graph_def,name="")

        inp = tf.get_default_graph().get_tensor_by_name('image_tensor:0')       
        out = tf.get_default_graph().get_tensor_by_name('detection_boxes:0')    
        out1 = tf.get_default_graph().get_tensor_by_name('detection_classes:0')
        out2 = tf.get_default_graph().get_tensor_by_name('detection_scores:0')
        out3 = tf.get_default_graph().get_tensor_by_name('num_detections:0')
        
        c= inp.shape[3].value
        h= inp.shape[1].value
        w= inp.shape[2].value


        sess.run(tf.global_variables_initializer())

        tempimg = np.zeros((model_r, model_c, 3), dtype = "uint8") #for dummy pass
        tempimg = np.expand_dims(tempimg, 0)
        fd_list = []
        for i in range(batch_size):
            fd_list.append(tempimg)
        feed_dict = {inp : np.concatenate(fd_list, 0)}
        output = sess.run(out, feed_dict) 
        
    net = {'sess':sess,\
    'inp':inp,\
    'out':out,\
    'c':c,\
    'h':h,\
    'w':w,\
    'min_prob':0.5, 
    'min_iou':0.3,
    'batsiz':batch_size,\
    'out1':out1,\
    'out2':out2,\
    'out3':out3}
    return net


if __name__ == '__main__':

    model_path = "../models/ssd_inception_v2_coco_2018_01_28/1/model.pb" # (300,300,3)  
    model_r = 300 
    model_c = 300 

    net = init_tf(0, model_path, 1, model_r, model_c, 2) 

    bs = [1]
    # bs = [6,7,8,16,32]
    num_itr = 1
    for bs_i in bs:
        net['batsiz'] = bs_i
        # for nf in range(8):
        fd_list = []
        temp_i = cv2.imread("./sample_720p.jpg")

        h1, w1, _ = temp_i.shape
        maxdim = max(h1,w1)
        temp_img = np.zeros((maxdim, maxdim, 3), dtype = "uint8")
        temp_img[0:h1,0:w1] = temp_i

        Scaling = float(maxdim) / float(model_c)
        # imsz1 = cv2.resize(temp_img, (model_c,model_r)) #maintain_aspect_ratio: 1
        imsz1 = cv2.resize(temp_i, (model_c,model_r)) #maintain_aspect_ratio:0
        imsz1 = imsz1[:,:,(2,1,0)]
        
        this_inp = np.expand_dims(imsz1, 0)
        inp_feed = list()
        for i in range(net['batsiz']):
            inp_feed.append(this_inp)
        t_fullcall = 0.0

        for itr in range(num_itr):
            t0 = time.time()
            get_face_roi_test_newnms(net,inp_feed,0,0,Scaling)
            t1 = time.time()
            tic = t1 - t0
            print(tic)

    print('done')