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('--test_label', type=str,
			    help='CSV file with annotation')
parser.add_argument('--config', type=str,
                help='configuration file')
parser.add_argument('--images_path', type=str,
			    help='images folder')
parser.add_argument('--mode', type=str,
			help='training_mode')    

parser.add_argument('-IIOU','--IIOU', action='store_true')

args = parser.parse_args()

test_path = args.test_label
base_path = os.getcwd()
images_path = args.images_path
mode = args.mode
use_iiou = args.IIOU


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

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

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)

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
test_data, label_nb_total_test_images, label_nb_total_test_pigs, label_nb_total_test_others = get_label(test_path)
tmp = 0

rpn_nms_range = [0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95]
max_boxes_range = [50,100,150,200,250,300,350,400,450,500,550,600]

rpn_confusion_matrix = np.zeros((2,3,len(rpn_nms_range)*len(max_boxes_range)))

detect_nms_range = [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]
detect_score_range = [0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95]
detect_nms_iiou_range = [-0.9,-0.8,-0.7,-0.6,-0.5,-0.4,-0.3,-0.2,-0.1,0,0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]

detect_confusion_matrix = np.zeros((3,3,len(detect_nms_range)*len(detect_score_range)))
detect_confusion_matrix_iiou = np.zeros((3,3,len(detect_nms_iiou_range)*len(detect_score_range)))

# Get data
test_imgs, classes_count_test, _ = get_data(test_path,images_path)
epoch_length = len(test_imgs)
print("")
data = get_anchor_gt(test_imgs, C, get_img_output_length, mode="PaO")
iter_num = 0
s = time.time()
while True:
    X, Y, img_data, debug_img, debug_num_pos, _ = next(data)
    tmp = tmp + 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=600, overlap_thresh=0.99)
    s = time.time()
    for i in range (len(rpn_nms_range)):
        for j in range (len(max_boxes_range)):
            R_rpn_nms, probs_nms, valid = non_max_suppression_fast( R,
                                                                    probs,
                                                                    overlap_thresh=rpn_nms_range[i],
                                                                    max_boxes=max_boxes_range[j])
            
            tmp_confusion = compare_rpn_to_groundtruth(img_data['bboxes'],R_rpn_nms*C.rpn_stride,num_boxes=max_boxes_range[j])
            rpn_confusion_matrix[:,:,i*len(max_boxes_range)+j] = rpn_confusion_matrix[:,:,i*len(max_boxes_range)+j] + tmp_confusion

    print("time processed rpn /image {} : {}".format(tmp,time.time()-s))
    iter_num = iter_num + 1

    # convert from (x1,y1,x2,y2) to (x,y,w,h)
    R[:, 2] -= R[:, 0]
    R[:, 3] -= R[:, 1]
    ROIs = np.expand_dims(R, axis=0)
    [P_cls, P_regr] = model_classifier.predict([F, ROIs])
    
    Y_detection_train = []
    for i in range (P_cls.shape[1]):     
        Y_detection_train.append( (ROIs[0][i], P_cls[0][i], P_regr[0][i]) )
    
    bboxes, probs = get_detections_boxes(Y_detection_train, C, class_mapping)
    
    for i in range (len(detect_nms_range)):
        for j in range (len(detect_score_range)):
            bboxes_nms, probs_nms, valid = non_max_suppression_fast(    bboxes,
                                                                        probs,
                                                                        overlap_thresh=detect_nms_range[i],
                                                                        max_boxes=600)
            
            tmp_confusion = compare_detection_to_groundtruth(   img_data['bboxes'],
                                                                bboxes_nms,
                                                                probs_nms,
                                                                iou_min=0.5, 
                                                                score_min=detect_score_range[j])
            detect_confusion_matrix[:,:,i*len(detect_score_range)+j] = detect_confusion_matrix[:,:,i*len(detect_score_range)+j] + tmp_confusion
            if i == 3 and j == 0:
                if tmp_confusion[0,1] > 0 or tmp_confusion[0,2] > 0:
                    print("False positive on image : ",img_data["filepath"])
                if tmp_confusion[1,0] > 0 or tmp_confusion[2,0] > 0:
                    print("False negative on image : ",img_data["filepath"])

    for i in range (len(detect_nms_iiou_range)):
        for j in range (len(detect_score_range)):
            bboxes_nms, probs_nms, valid = non_max_suppression_fast_iiou(       bboxes,
                                                                                probs,
                                                                                overlap_thresh=detect_nms_iiou_range[i],
                                                                                max_boxes=600)
                    
            tmp_confusion = compare_detection_to_groundtruth(   img_data['bboxes'],
                                                                bboxes_nms,
                                                                probs_nms,
                                                                iou_min=0.5, 
                                                                score_min=detect_score_range[j])
            detect_confusion_matrix_iiou[:,:,i*len(detect_score_range)+j] = detect_confusion_matrix_iiou[:,:,i*len(detect_score_range)+j] + tmp_confusion



    print("time processed class /image : {}".format(time.time()-s))

    if iter_num == epoch_length:
        np.save('../rpn.npy', rpn_confusion_matrix)
        np.save('../detect.npy', detect_confusion_matrix)
        np.save('../detect_iiou.npy', detect_confusion_matrix_iiou)
        quit()
