from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='3'

############################################################
                    #CONFIG SETTING
############################################################

class Config:

    def __init__(self):

        # Print the process or not
        self.verbose = True

        # Name of base network
        self.network = 'vgg'

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

        # Anchor box scales
    # Note that 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 = [100, 175, 350]
        #self.anchor_box_scales = [64, 128, 256, 512, 1024]

        # 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)]]
        self.anchor_box_ratios = [[1, 1], [1, 2], [2, 1]]

        # 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 = 720

        # 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
        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
        self.rpn_num_train_examples = 150 
        self.rpn_min_overlap = 0.3
        self.rpn_max_overlap = 0.65

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

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

        self.model_path = None

import argparse

parser = argparse.ArgumentParser(description='Argument')
parser.add_argument('--device', type=str,
                    help='device visible for keras (CPU: -1)')
parser.add_argument('--num_rois', type=int,
			help='num_rois')
parser.add_argument('--graph', type=str,
                help='path to graph')

parser.add_argument('--images_folder', type=str,
                    help='path to images')
parser.add_argument('--image', type=str,
                    help='path to images')

parser.add_argument('--confidence', type=float, default=0.5,
                    help='Confidence Treshold')
parser.add_argument('--config', type=str,
                help='configuration file')
parser.add_argument('--nms', type=float, default=0.5,
			help='non max supression threshold')
parser.add_argument('-p','--only_pig', action='store_true',
			help='show only pig')
parser.add_argument('-r','--rpn', action='store_true',
			help='show result rpn')
parser.add_argument('-s','--save_output', action='store_true',
			help='show result rpn')
parser.add_argument('-iiou','--iiou', action='store_true',
			help='show result rpn')

args = parser.parse_args()

base_path = os.getcwd()
args = parser.parse_args()
save_output = args.save_output
device = args.device
use_iiou = args.iiou

if device:
    os.environ["CUDA_VISIBLE_DEVICES"]=device

if args.num_rois:
    num_rois = args.num_rois
else:
    num_rois = 4

imgs_path = []

# Process folder
images_folder = args.images_folder
if images_folder:
    imgs_path =  os.listdir(images_folder)
    imgs_path = [images_folder + "/" + s for s in imgs_path]

# Process one image
image = args.image
if image:
    imgs_path = [image]

graph_path = args.graph

# Load config file
import pickle
config_path = args.config
with open(config_path, 'rb') as f_in:
	C = pickle.load(f_in)

show_only_pig = args.only_pig
show_rpn = args.rpn
confidence_threshold = args.confidence
nms_threshold = args.nms
num_rois = args.num_rois

import random
import pprint
import sys
import time
import numpy as np
from optparse import OptionParser
import math
import cv2
import copy
from matplotlib import pyplot as plt
import tensorflow as tf
import pandas as pd
import keyboard

from sklearn.metrics import average_precision_score

from keras import backend as K
from keras.optimizers import Adam, SGD, RMSprop
from keras.layers import Flatten, Dense, Input, Conv2D, MaxPooling2D, Dropout
from keras.layers import GlobalAveragePooling2D, GlobalMaxPooling2D, TimeDistributed
from keras.engine.topology import get_source_inputs
from keras.utils import layer_utils
from keras.utils.data_utils import get_file
from keras.objectives import categorical_crossentropy

from keras.models import Model
from keras.utils import generic_utils
from keras.engine import Layer, InputSpec
from keras import initializers, regularizers

from new_utils import *

############################################################
                        #Start testing
############################################################


from tensorflow.python.client import device_lib

print("List of available devices on tensorflow : ")
print(device_lib.list_local_devices())

print("List of available devices on Keras : ")
print(K.tensorflow_backend._get_available_gpus())
print("#############")


# turn off any data augmentation at test time
C.use_horizontal_flips = False
C.use_vertical_flips = False
C.rot_90 = False
C.num_rois = num_rois

num_features = 512

input_shape_img = (None, None, 3)
input_shape_features = (None, None, num_features)

img_input = Input(shape=input_shape_img)
roi_input = Input(shape=(C.num_rois, 4))
feature_map_input = Input(shape=input_shape_features)

# define the base network (VGG here, can be Resnet50, Inception, etc)
shared_layers = nn_base(img_input, trainable=True)

# define the RPN, built on the base layers
num_anchors = len(C.anchor_box_scales) * len(C.anchor_box_ratios)
rpn_layers = rpn_layer(shared_layers, num_anchors)

classifier = classifier_layer(feature_map_input, roi_input, C.num_rois, nb_classes=len(C.class_mapping))

model_rpn = Model(img_input, rpn_layers)
model_classifier_only = Model([feature_map_input, roi_input], classifier)

model_classifier = Model([feature_map_input, roi_input], classifier)

print('Loading weights from {}'.format(graph_path))
model_rpn.load_weights(graph_path, by_name=True)
model_classifier.load_weights(graph_path, by_name=True)

model_rpn.compile(optimizer='sgd', loss='mse')
model_classifier.compile(optimizer='sgd', loss='mse')

# Switch key value for class mapping
class_mapping = C.class_mapping
class_mapping = {v: k for k, v in class_mapping.items()}
print("class_mapping",class_mapping)

class_to_color = {class_mapping[v]: np.random.randint(0, 255, 3) for v in class_mapping}

classes = {}

# If the box classification value is less than this, we ignore this box


for idx, img_name in enumerate(imgs_path):
    if not img_name.lower().endswith(('.bmp', '.jpeg', '.jpg', '.png', '.tif', '.tiff')):
        continue
    img = cv2.imread(img_name)
    st = time.time()
    X, ratio = format_img(img, C)

    X = np.transpose(X, (0, 2, 3, 1))

    # get output layer Y1, Y2 from the RPN and the feature maps F
    # Y1: y_rpn_cls
    # Y2: y_rpn_regr
    [Y1, Y2, F] = model_rpn.predict(X)

    # Get bboxes by applying NMS 
    # R.shape = (300, 4)
    R,probs = rpn_to_roi(Y1, Y2, C, K.image_dim_ordering(), max_boxes=300, overlap_thresh=0.7)
    if show_rpn:
        img_rpn = img.copy()
        for i in range (R.shape[0]):
            if probs[i] > confidence_threshold:
                x1 = R[i][0]*16
                y1 = R[i][1]*16
                x2 = R[i][2]*16
                y2 = R[i][3]*16

                textOrg = (x1, y1+20)
                textLabel = '{}'.format(int(100*probs[i]))
                (retval,baseLine) = cv2.getTextSize(textLabel,cv2.FONT_HERSHEY_COMPLEX,1,1)
                cv2.rectangle(img_rpn, (x1, y1), (x2, y2), (random.randint(0,255),random.randint(0,255),random.randint(0,255)), 2)
                cv2.rectangle(img_rpn, (textOrg[0] - 5, textOrg[1]+baseLine - 5), (textOrg[0]+retval[0] + 5, textOrg[1]-retval[1] - 5), (0, 0, 0), 1)
                cv2.rectangle(img_rpn, (textOrg[0] - 5,textOrg[1]+baseLine - 5), (textOrg[0]+retval[0] + 5, textOrg[1]-retval[1] - 5), (255, 255, 255), -1)

                cv2.putText(img_rpn, textLabel, textOrg, cv2.FONT_HERSHEY_DUPLEX, 1, (0, 0, 0), 1)



    # convert from (x1,y1,x2,y2) to (x,y,w,h)
    R[:, 2] -= R[:, 0]
    R[:, 3] -= R[:, 1]

    # apply the spatial pyramid pooling to the proposed regions
    bboxes = {}
    probs = {}

    for jk in range(R.shape[0]//C.num_rois + 1):
        ROIs = np.expand_dims(R[C.num_rois*jk:C.num_rois*(jk+1), :], axis=0)
        if ROIs.shape[1] == 0:
            break

        if jk == R.shape[0]//C.num_rois:
            #pad R
            curr_shape = ROIs.shape
            target_shape = (curr_shape[0],C.num_rois,curr_shape[2])
            ROIs_padded = np.zeros(target_shape).astype(ROIs.dtype)
            ROIs_padded[:, :curr_shape[1], :] = ROIs
            ROIs_padded[0, curr_shape[1]:, :] = ROIs[0, 0, :]
            ROIs = ROIs_padded

        [P_cls, P_regr] = model_classifier_only.predict([F, ROIs])

        # Calculate bboxes coordinates on resized image
        for ii in range(P_cls.shape[1]):
            # Ignore 'bg' class
            if np.max(P_cls[0, ii, :]) < confidence_threshold or np.argmax(P_cls[0, ii, :]) == (P_cls.shape[2] - 1):
                continue

            cls_name = class_mapping[np.argmax(P_cls[0, ii, :])]

            if cls_name not in bboxes:
                bboxes[cls_name] = []
                probs[cls_name] = []

            (x, y, w, h) = ROIs[0, ii, :]

            cls_num = np.argmax(P_cls[0, ii, :])
            try:
                (tx, ty, tw, th) = P_regr[0, ii, 4*cls_num:4*(cls_num+1)]
                tx /= C.classifier_regr_std[0]
                ty /= C.classifier_regr_std[1]
                tw /= C.classifier_regr_std[2]
                th /= C.classifier_regr_std[3]
                x, y, w, h = apply_regr(x, y, w, h, tx, ty, tw, th)
            except:
                pass
            bboxes[cls_name].append([C.rpn_stride*x, C.rpn_stride*y, C.rpn_stride*(x+w), C.rpn_stride*(y+h)])
            probs[cls_name].append(np.max(P_cls[0, ii, :]))

    all_dets = [] 

    for key in bboxes:
        if ((show_only_pig is False) or (show_only_pig is True and key == 'pig')):
            bbox = np.array(bboxes[key])
            if use_iiou:
                new_boxes, new_probs, valid  = non_max_suppression_fast_iiou(bbox, np.array(probs[key]), overlap_thresh=nms_threshold)
            else:
                new_boxes, new_probs, valid  = non_max_suppression_fast(bbox, np.array(probs[key]), overlap_thresh=nms_threshold)
            for jk in range(new_boxes.shape[0]):
                (x1, y1, x2, y2) = new_boxes[jk,:]

                # Calculate real coordinates on original image
                (real_x1, real_y1, real_x2, real_y2) = get_real_coordinates(ratio, x1, y1, x2, y2)

                color = (random.randint(0,255),random.randint(0,255),random.randint(0,255))
                
                cv2.rectangle(img,(real_x1, real_y1), (real_x2, real_y2), color,4)

                textLabel = '{}: {}'.format(key,int(100*new_probs[jk]))
                all_dets.append((key,100*new_probs[jk]))

                (retval,baseLine) = cv2.getTextSize(textLabel,cv2.FONT_HERSHEY_COMPLEX,1,1)
                textOrg = (real_x1, real_y1+20)

                cv2.rectangle(img, (textOrg[0] - 5, textOrg[1]+baseLine - 5), (textOrg[0]+retval[0] + 5, textOrg[1]-retval[1] - 5), (0, 0, 0), 1)
                cv2.rectangle(img, (textOrg[0] - 5,textOrg[1]+baseLine - 5), (textOrg[0]+retval[0] + 5, textOrg[1]-retval[1] - 5), (255, 255, 255), -1)
                cv2.putText(img, textLabel, textOrg, cv2.FONT_HERSHEY_DUPLEX, 1, (0, 0, 0), 1)

    if save_output:
        print(idx)
        output_path = "../output/" + img_name[len(images_folder)+1:]
        cv2.imwrite(output_path,img)

        print('Elapsed time = {}'.format(time.time() - st))

    else:

        cv2.imshow('image',img)

        if show_rpn:
            cv2.imshow('rpn', img_rpn)

        while True:
            k = cv2.waitKey(0)
            if k == 27 or k == 113:         # wait for ESC or q key to exit
                cv2.destroyAllWindows()
                quit()
            if k == 110:
                break
            elif k == ord('s'): # wait for 's' key to save
                cv2.imwrite(img_name[len(images_folder)+1:],img)