import shap
##import statements
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from tensorflow.python.keras.utils.data_utils import Sequence
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import cv2 as cv
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.optimizers import RMSprop as rmsprop
from tensorflow.keras import backend as k
# from tensorflow.keras import backend as K
import CustomDataGenerator
from tensorflow.keras.layers import Input
import os,time,math,sys
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.applications.xception import Xception
from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2
import AlexNet
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.densenet import DenseNet121
from collections import OrderedDict
from sklearn import metrics
from tensorflow.keras.preprocessing import image as Image
import cv2
import time

BATCH_SIZE = 64
GPU_ID = "2"
GPU_FRACTION = 0.5


##set GPU options
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=GPU_ID
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = GPU_FRACTION
k.set_session(tf.Session(config=config))

def get_filename(dir):
    (_, _, filenames) = next(os.walk(dir))
    for i in filenames:
        if(i.find('data')!=-1):
            data = i
            break
    return(data.split('.')[0])

def load_image(path, preprocess=True):
    """Load and preprocess image."""
    x = Image.load_img(path, target_size=(240, 320))
    if preprocess:
        x = Image.img_to_array(x)
        x = x/255.0
        x = np.expand_dims(x, axis=0)
    return x[:,:,:,0][...,None]


def normalize(x):
    """Utility function to normalize a tensor by its L2 norm"""
    return (x + 1e-10) / (k.sqrt(k.mean(k.square(x))) + 1e-10)


##create data generator for test
test_path = "/home/vedantb/datav5/Val_v53_3c.txt"
test_df = pd.read_csv(test_path,sep=".png ",header=None)
test_df[0] = test_df[0].astype(str)+'.png'
test_df[0]=test_df[0].replace("/media/Newhdd2tb","/mnt/2tbhdd",regex=True)
test_list_IDs = test_df[0].values
test_labels = test_df[1].values
test_gen = CustomDataGenerator.DataGenerator(test_list_IDs,test_labels,batch_size=BATCH_SIZE,n_channels=1,shuffle=False)

val_path = "/home/vedantb/datav5/Val_v53_3c_by5.txt"
val_df = pd.read_csv(val_path,sep=".png ",header=None)
val_df[0] = val_df[0].astype(str)+'.png'
val_df[0]=val_df[0].replace("/media/Newhdd2tb","/mnt/2tbhdd",regex=True)
val_list_IDs = val_df[0].values
val_labels = val_df[1].values

##define models dictionary
models = {}
#models['resnet50_2']=ResNet50
#models['mobilenet'] = MobileNetV2
#models['alexnet']=AlexNet
#models['vgg16']=VGG16
models['densenet121'] = DenseNet121
#models['xception']=Xception


##import model
def get_model(name):
    model=models[name]
    if(name!="alexnet"):
        model = model(weights=None,input_shape=(240,320,1),include_top=True,classes=3)
    else:
        model=model.alexnet_model()
    checkpoint_path = get_filename('./best_models/'+name)
    model.load_weights('./best_models/{0}/{1}'.format(name,checkpoint_path))
    model.compile(loss='categorical_crossentropy',
        optimizer=keras.optimizers.Adam(lr=0.00001),
        metrics=['accuracy'])
    for i in reversed(model.layers):
        if(len(i.output_shape)==4):
            layer_to_analyze=i.name
            break
    return(model,layer_to_analyze)


def computeInferenceTime(model,num_images):
    total_time=0
    img = load_image(test_list_IDs[0])
    trial_run = model.predict(img)
    t0 = time.time()
    for i in range(1,num_images+1):
        img = load_image(test_list_IDs[i])
        t1 = time.time()
        model.predict(img)
        t2= time.time()
#         print(" ",t2-t1)
        total_time+= t2-t1
    t3 = time.time()
#     print("total",t3-t0)
    return(total_time/num_images)
dict_time = {}
for name,m in models.items():
    print("name",name)
    model,layer_name = get_model(name)
    a = computeInferenceTime(model,10)
    dict_time[name]=a
    print(a)
#     break
model_sess= None
export_dir = "../XAI/utils/portal/tf_googlenet.model"
# tf.device('/device:GPU:2')
global sess
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.4)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options,allow_soft_placement=True,device_count = {'GPU': 1}))
# sess = tf.Session()
print("googlenet")
tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], export_dir)
x = sess.graph.get_tensor_by_name('data:0')
y = sess.graph.get_tensor_by_name('softmax:0')
img = load_image(test_list_IDs[0])
trial_run = sess.run(y,feed_dict={x:img})
t0 = time.time()
total_time=0
for i in range(1,10+1):
    img = load_image(test_list_IDs[i])
    t1 = time.time()
    sess.run(y,feed_dict={x:img})
    t2= time.time()
#     print(" ",t2-t1)
    total_time+= t2-t1
t3 = time.time()
# print("total",t3-t0)
dict_time['googlenet']=total_time/10
np.save('./best_models/timing_stats-densenet121-10.npy',dict_time)
    
