"""
python test_RAP_hiarchical.py -m hiarBayesGoogLeNet -b 64 -g 1 -w ../models/imagenet_models/hiarBayesGoogLeNet_RAP
python test_RAP_hiarchical.py -m hiarGoogLeNet -b 64 -g 1 -w ../models/imagenet_models/GoogLeNet_RAP
python test_RAP_hiarchical.py -m hiarBayesGoogLeNet -b 64 -g 1 -w ../models/imagenet_models/hiarBayesGoogLeNet_RAP/binary3_epoch50_valloss0.40.hdf5
python test_RAP_hiarchical.py -m hiarBayesGoogLeNet -c 51 -g 1 -s 0 -wd 224 -hg 224 -w adam
"""
#from network.hiarGoogLenetSPP import hiarGoogLeNetSPP
#from network.hiarGoogLenetWAM import hiarGoogLeNetWAM
from network.hiarGoogLenet import hiarGoogLeNet
from network.hiarGoogLenetGAP import hiarGoogLeNetGAP
from network.hiarBayesGoogLenetGAP import hiarBayesGoogLeNetGAP
from network.hiarBayesGoogLenet_gap import hiarBayesGoogLeNet as hiarBayesGoogLeNet_gap
from network.hiarBayesGoogLenet_gap_v2 import hiarBayesGoogLeNet as hiarBayesGoogLeNet_gap_v2
from network.hiarBayesGoogLenet_gap_v3 import hiarBayesGoogLeNet as hiarBayesGoogLeNet_gap_v3
from network.hiarBayesGoogLenet_gap_v4 import hiarBayesGoogLeNet as hiarBayesGoogLeNet_gap_v4
from network.hiarBayesGoogLenet_gap_v5 import hiarBayesGoogLeNet as hiarBayesGoogLeNet_gap_v5
from network.hiarGoogLenet_high import hiarGoogLeNet_high
from network.hiarGoogLenet_mid import hiarGoogLeNet_mid
from network.hiarGoogLenet_low import hiarGoogLeNet_low
from network.hiarBayesGoogLenet import hiarBayesGoogLeNet
from network.hiarBayesResNet import hiarBayesResNet
from network.hiarBayesInception_v4 import hiarBayesInception_v4
from network.hiarBayesGoogLenet_inception import hiarBayesGoogLeNet as hiarBayesGoogLeNet_inception
from network.hiarGoogLenet_inception import hiarGoogLeNet as hiarGoogLeNet_inception
from network.hiarBayesGoogLenetv2 import hiarBayesGoogLeNet as hiarBayesGoogLeNetv2
from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing import image
from keras.utils import multi_gpu_model
from sklearn.model_selection import train_test_split
from keras.callbacks import ModelCheckpoint, LearningRateScheduler, ReduceLROnPlateau, EarlyStopping, TensorBoard, CSVLogger
import sys
import os
import argparse
import json
import numpy as np
import pandas as pd
import tqdm
from keras import backend as K
from angular_losses import bayes_binary_crossentropy
import re

def generate_imgdata_from_file(X_path, y, batch_size, image_height, image_width):
    while True:
        cnt = 0
        X = []
        Y = []
        for i in range(len(X_path)):
            img = image.load_img(X_path[i], target_size=(image_height, image_width, 3))
            img = image.img_to_array(img)
            X.append(img)
            Y.append(y[i])
            cnt += 1
            if cnt==batch_size:
                cnt = 0
                yield (np.array(X), np.array(Y))
                X = []
                Y = []

def parse_arg():
    models = ['hiarGoogLeNet', 'hiarGoogLeNetGAP', 'hiarBayesGoogLeNet', 'hiarBayesGoogLeNet_inception', 'hiarGoogLeNet_inception']
    parser = argparse.ArgumentParser(description='training of the WPAL...')
    parser.add_argument('-g', '--gpus', type=str, default='',
                        help='The gpu device\'s ID need to be used')
    parser.add_argument('-c', '--classes', type=int, default=92,
                        help='The total number of classes to be predicted')
    parser.add_argument('-b', '--batch', type=int, default=1,
                        help='The batch size of the training process')
    parser.add_argument('-wd', '--width', type=int, default=224,
                        help='The width of thWPAL_PETAe picture')
    parser.add_argument('-hg', '--height', type=int, default=224,
                        help='The height of the picture')
    parser.add_argument('-w', '--weight', type=str, default='',
                        help='The weights file of the pre-training')
    parser.add_argument('-m', '--model', type=str, default='',
                        help='The model including: '+str(models))
    parser.add_argument('-s', '--split', type=int, default=0,
                        help='The split')
    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
    return args


    


if __name__ == "__main__":
    args = parse_arg()
    save_name = ""
    low_level = [11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level = [4,5,6,7,8,9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
    high_level = [0,1,2,3,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    low_level = [11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level = [9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
    high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    
    low_level = [11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level = [4,5,6,7,8,9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50]
    high_level = [0,1,2,3]#,51,52,53,54,55,56,57,58,59,60,61,62
    """
    save_name = "binary51_75v2_"
    low_level = [11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level = [9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
    high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    #"""
    """
    save_name = "binary92_oldhiar_newlossnoexp_split" + str(args.split)
    low_level = [11,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91]#
    mid_level = [9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
    high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62]#
    #"""
    if args.model[-1] == '3':
        #save_name = str(args.height) + "x" + str(args.width) + "binary51_newhier_newlossnoexp_split" + str(args.split)
        
        #low_level = [9,10,11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
        #mid_level = [4,5,6,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
        #high_level = [0,1,2,3,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
        mid_level_hs = [9,10,12,13,14]
        mid_level_ub = [15,16,17,18,19,20,21,22,23]
        mid_level_lb = [24,25,26,27,28,29]
        mid_level_sh = [30,31,32,33,34]
        mid_level_at = [35,36,37,38,39,40,41,42]
        mid_level_ot = [4,5,6,7,8,43,44,45,46,47,48,49,50]
        #"""
        ##################### newnewhier
        low_level = [9,10,11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
        mid_level_hs = [12,13,14]
        mid_level_ub = [15,16,17,18,19,20,21,22,23]
        mid_level_lb = [24,25,26,27,28,29]
        mid_level_sh = [30,31,32,33,34]
        mid_level_at = [35,36,37,38,39,40,41,42]
        mid_level_ot = [4,5,6]
        high_level = [0,1,2,3,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
        #"""
    """
    ##################### newnewnewhier
    low_level = [9,10,11,12,13,14,35,36,37,38,39,40,41,42]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level_ub = [15,16,17,18,19,20,21,22,23]
    mid_level_lb = [24,25,26,27,28,29]
    mid_level_sh = [30,31,32,33,34]
    high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    #"""
    if args.model[-1] == '5':
        #"""
        ##################### newnewnewnewhier
        low_level = [9,10,11,12,13,14]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
        mid_level_ub = [15,16,17,18,19,20,21,22,23]
        mid_level_lb = [24,25,26,27,28,29]
        mid_level_sh = [30,31,32,33,34]
        mid_level_at = [35,36,37,38,39,40,41,42]
        high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
        #"""
    
    """
    low_level = [9,10,11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level_hs = [12,13,14]
    mid_level_ub = [15,16,17,18,19,20,21,22,23]
    mid_level_lb = [24,25,26,27,28,29]
    mid_level_sh = [30,31,32,33,34]
    mid_level_at = [35,36,37,38,39,40,41,42]
    mid_level_ot = [4,5,6]
    high_level = [0,1,2,3,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    """
    """
    save_name = str(args.height) + "x" + str(args.width) + "binary51_beforreduction_oldhier_newlossnoexp_split" + str(args.split)
    low_level = [11]#,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91
    mid_level = [9,10,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42]
    high_level = [0,1,2,3,4,5,6,7,8,43,44,45,46,47,48,49,50]#,51,52,53,54,55,56,57,58,59,60,61,62
    """
    class_num = args.classes


    # Data augmentation to pre-processing
    heavy_augmentation = True
    if heavy_augmentation:
        datagen = ImageDataGenerator(
            featurewise_center=False,
            samplewise_center=False,
            featurewise_std_normalization=False,
            samplewise_std_normalization=False)
    else:
        datagen = ImageDataGenerator(
            featurewise_center=False,
            samplewise_center=False,
            featurewise_std_normalization=False,
            samplewise_std_normalization=False)
    image_width = args.width
    image_height = args.height
    #hiarBayesGoogLeNet
    filename = r"../results/RAP_labels_pd.csv"
    data = np.array(pd.read_csv(filename))[:, 1:]
    length = len(data)
    #global alpha
    load = True
    if load:
        data_x = np.zeros((length, image_height, image_width, 3))
    data_y = np.zeros((length, class_num))
    data_path = []
    for i in range(length):
        #img = image.load_img(path + m)
        data_path.append(data[i, 0])
        if load:
            img = image.load_img(data[i, 0], target_size=(image_height, image_width, 3))
            data_x[i] = image.img_to_array(img)
        data_y[i] = np.array(data[i, 1:1+class_num], dtype="float32")
    data_path = np.array(data_path)
    #X_train, X_test, y_train, y_test = train_test_split(data_x, data_y, test_size=0.3, random_state=0)
    split = np.load('../results/RAP_partion.npy', allow_pickle=True).item()
    if load:
        X_test = data_x[list(split['test'][args.split])]
    X_test_path = data_path[list(split['test'][args.split])]
    y_test = data_y[list(split['test'][args.split])]
    print("The shape of the X_test is: ", X_test_path.shape)
    print("The shape of the y_test is: ", y_test.shape)
    
    
    #googleNet默认输入32*32的图片
    if args.model == "hiarGoogLeNet":
        model_dir = "hiarGoogLeNet_RAP/"
        model = hiarGoogLeNet.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func = 'binary_crossentropy'#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet":
        model_dir = "hiarBayesGoogLeNet_RAP/"
        model = hiarBayesGoogLeNet.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)], test=True)
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesResNet":
        model_dir = "hiarBayesResNet_RAP/"
        model = hiarBayesResNet.build([len(low_level), len(mid_level), len(high_level)])#, test=True
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_inception":
        model_dir = "hiarBayesGoogLeNet_Inception_RAP/"
        model = hiarBayesGoogLeNet_inception.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarGoogLeNet_inception":
        model_dir = "hiarGoogLeNet_Inception_RAP/"
        model = hiarGoogLeNet_inception.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNetv2":
        model = hiarBayesGoogLeNetv2.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesInception_v4":
        model_dir = "hiarBayesInceptionV4_RAP/"
        model = hiarBayesInception_v4(image_width, image_height, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarGoogLeNetGAP":
        model_dir = 'hiarGoogLeNetGAP_RAP/'
        model = hiarGoogLeNetGAP.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func = 'binary_crossentropy'#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNetGAP":
        model_dir = 'hiarBayesGoogLeNetGAP_RAP/'
        model = hiarBayesGoogLeNetGAP.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func = 'binary_crossentropy'#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_gap":
        model_dir = 'hiarBayesGoogLeNetgap_RAP/'
        model = hiarBayesGoogLeNet_gap.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func = 'binary_crossentropy'#weighted_categorical_crossentropy(alpha)
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_gap_v2":
        model_dir = 'hiarBayesGoogLeNetgap_RAP/'
        model = hiarBayesGoogLeNet_gap_v2.build(image_height, image_width, 3, [len(low_level), len(mid_level), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_func = {"low":'binary_crossentropy', "middle":'binary_crossentropy', "high":'binary_crossentropy'}
        is_multi = [len(low_level), len(mid_level), len(high_level)]
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_gap_v3":
        model_dir = 'hiarBayesGoogLeNetgap_RAP/'
        model = hiarBayesGoogLeNet_gap_v3.build(image_height, image_width, 3, [len(low_level), len(mid_level_hs), len(mid_level_ub), len(mid_level_lb), len(mid_level_sh), len(mid_level_at), len(mid_level_ot), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_func = {"low":'binary_crossentropy', "middle_hs":'binary_crossentropy', "middle_ub":'binary_crossentropy', "middle_lb":'binary_crossentropy', "middle_sh":'binary_crossentropy', "middle_at":'binary_crossentropy', "middle_ot":'binary_crossentropy', "high":'binary_crossentropy'}
        is_multi = [len(low_level), len(mid_level_hs), len(mid_level_ub), len(mid_level_lb), len(mid_level_sh), len(mid_level_at), len(mid_level_ot), len(high_level)]
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_gap_v4":
        model_dir = 'hiarBayesGoogLeNetgap_RAP/'
        model = hiarBayesGoogLeNet_gap_v4.build(image_height, image_width, 3, [len(low_level), len(mid_level_ub), len(mid_level_lb), len(mid_level_sh), len(high_level)])
        loss_func ='binary_crossentropy'#bayes_binary_crossentropy(alpha, y_train)#weighted_categorical_crossentropy(alpha)
        loss_func = {"low":'binary_crossentropy', "middle_ub":'binary_crossentropy', "middle_lb":'binary_crossentropy', "middle_sh":'binary_crossentropy', "high":'binary_crossentropy'}
        is_multi = [len(low_level), len(mid_level_ub), len(mid_level_lb), len(mid_level_sh), len(high_level)]
        loss_weights = None
        metrics=['accuracy']
    elif args.model == "hiarBayesGoogLeNet_gap_v5":
        model_dir = 'hiarBayesGoogLeNetgap_RAP/'
        model = hiarBayesGoogLeNet_gap_v5.build(image_height, image_width, 3, [len(low_level), len(mid_level_ub), len(mid_level_lb), len(mid_level_sh), len(mid_level_at), len(high_level)])
    gpus_num = len(args.gpus.split(','))
    if gpus_num > 1:
        multi_gpu_model(model, gpus=gpus_num)
    #model.compile(loss="categorical_crossentropy", optimizer='adam', metrics=['accuracy'])
    #model.compile(loss=loss_func, optimizer='adam', loss_weights=loss_weights, metrics=metrics)
    model.summary()

    reg = args.weight + "_(e|f)1*"
    print(reg)
    print(model_dir)
    weights = [s for s in os.listdir("../models/imagenet_models/" + model_dir) 
          if re.match(reg, s)]
    print(weights)
    batch_size = args.batch
    if load:
        test_generator = datagen.flow(X_test, y_test, batch_size=batch_size)
    else:
        test_generator = generate_imgdata_from_file(X_test_path, y_test, batch_size, image_height, image_width)
    for w in tqdm.tqdm(weights):
        if os.path.exists("../results/predictions/" + args.model + "_RAP_" + w + ".npy"):
            continue
        model.load_weights("../models/imagenet_models/" + model_dir + w, by_name=True)
        #test_generator.reset()
        #predictions_list = model.predict_generator(test_generator, steps=y_test.shape[0]/batch_size)
        predictions_list = model.predict(X_test)
        #print(np.array(predictions_list).shape)
        for i in predictions_list:
            print(i.shape)
        if args.model[-1] == '2' or args.model[-1] == '3' or args.model[-1] == '4' or args.model[-1] == '5':
            predictions = np.hstack(predictions_list)
        else:
            predictions = np.array(predictions_list)
        print("The shape of the predictions_test is: ", predictions.shape)
        #w[w.rindex('/')+1:w.rindex('.')]
        np.save("../results/predictions/" + args.model + '_RAP_' + w + ".npy", predictions)
        print("../results/predictions/" + args.model + '_RAP_' + w + ".npy")