from scipy import spatial
import os
import pickle
from six.moves import cPickle
from keras import backend as K
from keras.layers import Input, Conv2D, MaxPooling2D
from keras.models import Model
import cv2
from matplotlib import pyplot as plt
import numpy as np
import operator


from flask import Flask, request, jsonify
import numpy as np
from PIL import Image

app = Flask(__name__)

@app.route('/')
def hello_world():
    return 'Hello World!'

@app.route('/test', methods=["POST"])
def test():
    data = {}
    if request.method == "POST":
        print("in test")

        img = Image.open(request.files['file'])
        img = np.array(img)
        # plt.imshow(img)
        img = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
        # print(img)
        a = ranking_module(img)
        # print(a)
        data["filename_list"] = []
        count = 0
        for i in a:
            data["filename_list"].append(i)
            count+=1
            if count >= 10:
               break

        # data["filename_list"].append('topcit.png')
        # data["filename_list"].append('양동욱_21.jpg')

    return jsonify(data)


class Config:

	def __init__(self):

		# Print the process or not
		self.verbose = True

		# base network --> Todo! ResNet 으로 바꿔야함
		self.network = 'vgg'

		# Setting for data augmentation
		self.use_horizontal_flips = False
		self.use_vertical_flips = False
		self.rot_90 = False

		# Anchor box scales
    	# If im_size is smaller, anchor_box_scales should be scaled
	    # Original anchor_box_scales in the paper is [128, 256, 512]
		self.anchor_box_scales = [64, 128, 256] 

		# Anchor box ratios
		self.anchor_box_ratios = [[1, 1], [1./math.sqrt(2), 2./math.sqrt(2)], [2./math.sqrt(2), 1./math.sqrt(2)]]

		# Size to resize the smallest side of the image
		# Original setting in paper is 600. Set to 300 in here to save training time
		self.im_size = 300

		# image channel-wise mean to subtract
		self.img_channel_mean = [103.939, 116.779, 123.68]
		self.img_scaling_factor = 1.0

		# number of ROIs at once --> Todo! If ResNet, RoI will be 3.
		self.num_rois = 4

		# stride at the RPN (this depends on the network configuration)
		self.rpn_stride = 16

		self.balanced_classes = False

		# scaling the stdev
		self.std_scaling = 4.0
		self.classifier_regr_std = [8.0, 8.0, 4.0, 4.0]

		# overlaps for RPN --> Todo! If ResNet, IoU can be reduced. (what value? : 아마 0.5도 괜찮을듯 논문에서 본 적 있음)
		self.rpn_min_overlap = 0.3
		self.rpn_max_overlap = 0.7

		# overlaps for classifier ROIs
		self.classifier_min_overlap = 0.1
		self.classifier_max_overlap = 0.5

		# placeholder for the class mapping, automatically generated by the parser
		self.class_mapping = None

		self.model_path = None

"""# VGG Model"""

def get_img_output_length(width, height):
    def get_output_length(input_length):
        return input_length//16

    return get_output_length(width), get_output_length(height)    

def nn_base(input_tensor=None, trainable=False):


    input_shape = (None, None, 3)

    if input_tensor is None:
        img_input = Input(shape=input_shape)
    else:
        if not K.is_keras_tensor(input_tensor):
            img_input = Input(tensor=input_tensor, shape=input_shape)
        else:
            img_input = input_tensor

    bn_axis = 3

    # Block 1
    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1')(img_input)
    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv2')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x)

    # Block 2
    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv2')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x)

    # Block 3
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv1')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv2')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x)

    # Block 4
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x)

    # Block 5
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv3')(x)
    # x = MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x)

    return x

"""# Load Config Pickle"""



"""# Image Preprocessing"""

def format_img_size(img, C):
	""" formats the image size based on config """
	img_min_side = float(C.im_size)
	(height,width,_) = img.shape
		
	if width <= height:
		ratio = img_min_side/width
		new_height = int(ratio * height)
		new_width = int(img_min_side)
	else:
		ratio = img_min_side/height
		new_width = int(ratio * width)
		new_height = int(img_min_side)
	img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
	return img, ratio	

def format_img_channels(img, C):
	""" formats the image channels based on config """
	img = img[:, :, (2, 1, 0)]
	img = img.astype(np.float32)
	img[:, :, 0] -= C.img_channel_mean[0]
	img[:, :, 1] -= C.img_channel_mean[1]
	img[:, :, 2] -= C.img_channel_mean[2]
	img /= C.img_scaling_factor
	img = np.transpose(img, (2, 0, 1))
	img = np.expand_dims(img, axis=0)
	return img

def format_img(img, C):
	""" formats an image for model prediction based on config """
	img, ratio = format_img_size(img, C)
	img = format_img_channels(img, C)
	return img, ratio





#print(len(load_pickles.keys()))
#print(load_pickles.items())

"""# Demo
 - Tested ranking module by query image
"""

#print(load_pickles.keys())

"""# Version 2 : Input is pixel data & not exists in database"""

def ranking_module(query) :
    global count
    global C
    global model_backbone
    global load_pickles
    count += 1

    if count == 1:
        base_path = '/ENGN/'
        #test_path = 'drive/My Drive/capstone/retrieval_test/retrieval_annotations.txt'  # Test data (annotation file)
        test_base_path = '/ENGN/retrieval_test'  # Directory to save the test images
        config_output_filename = os.path.join(base_path, 'model_vgg_config.pickle')
        pickles = 'backbone_pickles'
        with open(config_output_filename, 'rb') as f_in:
            C = pickle.load(f_in)
        # turn off any data augmentation at test time
        C.use_horizontal_flips = False
        C.use_vertical_flips = False
        C.rot_90 = False
        """# Construct and Load the Model"""
        num_features = 512
        input_shape_img = (None, None, 3)
        img_input = Input(shape=input_shape_img)
        # define the base network *Todo : (VGG here, can be Resnet50 ...)*
        shared_layers = nn_base(img_input, trainable=True)
        # model backbone
        model_backbone = Model(img_input, shared_layers)
        model_backbone.load_weights('/ENGN/model_frcnn_vgg.hdf5', by_name=True)
        # model_backbone.load_weights(C.model_path, by_name=True)

        # model backbone
        '''
        Load Pickles 
         - predict 512 values using backbone model
         - saved these values as pickle file
        '''
        load_pickles = cPickle.load(open(os.path.join(test_base_path, pickles), "rb", True))

    query_class=''
    X, ratio = format_img(query, C)
    X = np.transpose(X, (0, 2, 3, 1))
    # 변형부분
    result=model_backbone.predict(X)
    query_pooling = np.max(result[0], axis=(0, 1))
    result={}
    
    for filepath, value in load_pickles.items() :
        cosine=spatial.distance.cosine(query_pooling, value['avg_pooling'])
        result[filepath] = [cosine, value['id']]
        
    cls_count={'rose' : 0, 'tiger' : 0, 'dragon' : 0, 'geometric' : 0}
    temp_result=sorted(result.items(), key=lambda x: x[1][0])
    
    for i in temp_result[:20] :
        cls_count[load_pickles[i[0]]['cls']]+=1
    # max_key = max(cls_count.items(), key=lambda k: k[1])
    # print(cls_count)
    query_class=sorted(cls_count.items(), key=lambda x: x[1], reverse=True)
    # print(query_class)
    for filename, value in result.items() :
        alpha = 1 if cls_count[load_pickles[filename]['cls']]==0 else cls_count[load_pickles[filename]['cls']]
        result[filename][0] +=(20/alpha) * 0.08
        
        # if load_pickles[filename]['cls'] == query_class[0][0] :

        #     continue
        # elif load_pickles[filename]['cls'] == query_class[1][0] :
        #     result[filename][0] *= 2
        #     continue
        # else : 
        #     result[filename][0] *= 2.5
        #     continue
    
    result=sorted(result.items(), key=lambda x: x[1])
    #print(query_class)
    # print(result)
    
    img=query
    # plt.figure(figsize=(24,20))
    # plt.imshow(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
    # plt.subplot(653),plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)),plt.title('query image')
    # frame   = plt.gca()
    # frame.axes.get_xaxis().set_visible(False)
    # frame.axes.get_yaxis().set_visible(False)
    # k=0
    # for i in result :
    #     print(i[1][1])
    #     # print(i[1][0])
    # for i in result[:20] :
        
    #     img = cv2.imread(i[0])
        
    #     plt.subplot(6,5,k+6),plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)), plt.title("rank " + str(k+1) + " : " + str(i[1][0]))
    #     k+=1
    #     frame   = plt.gca()
    #     frame.axes.get_xaxis().set_visible(False)
    #     frame.axes.get_yaxis().set_visible(False)
    id_output=[x[1][1] for x in result]

    return id_output



count = 0
C = None
model_backbone = None
load_pickles = None

if __name__ == '__main__':
    app.run(host='0.0.0.0')

