import os
import sys
import glob

import numpy as np
import tensorflow as tf
import scipy
import scipy.io
import keras
from keras.models import Model, Sequential
from keras.layers import *
from keras.optimizers import Adam
from keras import regularizers
from keras import backend as K
from keras.utils import to_categorical
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import LearningRateScheduler

from Benchmark import Benchmark
from Config import MODEL_PARAMS_DIR



IMAGENET_DIR = '/home/nz11/ILSVRC2012/'
NUM_TUNE_CLASSES = 200
IMAGES_PER_CLASS = 50



class AlexNet(Benchmark):

    def data_preprocess(self):

        X_val = np.fromfile(MODEL_PARAMS_DIR + '/alexnet_imagenet/test_input.bin', dtype=np.float32)
        y_val = np.fromfile(MODEL_PARAMS_DIR + '/alexnet_imagenet/test_labels.bin', dtype=np.uint32)
        
        X_val = X_val.reshape((-1, 3, 224, 224)) 
        X_train, y_train = None, None
        
            
        X_test = X_val[0:5000]
        y_test = y_val[0:5000]
        X_tuner = X_val[5000:]
        y_tuner = y_val[5000:]
        
        return X_train, y_train, X_test, y_test, X_tuner, y_tuner
    
    
    def buildModel(self):

        input_layer = Input((3, 224, 224))

        x = ZeroPadding2D((2, 2))(input_layer)
        x = Conv2D(64, (11, 11), strides=4, padding='valid')(x)
        x = Activation('relu')(x)
        x = MaxPooling2D(3, 2)(x)

        x = ZeroPadding2D((2, 2))(x)
        x = Conv2D(192, (5, 5), padding='valid')(x)
        x = Activation('relu')(x)
        x = MaxPooling2D(3, 2)(x)

        x = Conv2D(384, (3, 3), padding='same')(x)
        x = Activation('relu')(x)

        x = Conv2D(256, (3, 3), padding='same')(x)
        x = Activation('relu')(x)

        x = Conv2D(256, (3, 3), padding='same')(x)
        x = Activation('relu')(x)

        x = MaxPooling2D(3, 2)(x)

        x = Flatten()(x)
        x = Dropout(0.5)(x)
        x = Dense(4096)(x)
        x = Activation('relu')(x)
        x = Dropout(0.5)(x)
        x = Dense(4096)(x) 
        x = Activation('relu')(x)
        x = Dense(self.num_classes)(x)
        x = Activation('softmax')(x)
        
        model = Model(input_layer, x)

        return model


    def trainModel(self, model, X_train, y_train, X_test, y_test):

        assert False, "ImageNet training not supported - use Pretrained weights"


    
if __name__ == '__main__':

      
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    # Changing to NCHW format
    K.set_image_data_format('channels_first')


    ### Parameters specific to each benchmark
    reload_dir = MODEL_PARAMS_DIR + '/alexnet_imagenet/'
    keras_model_file = MODEL_PARAMS_DIR + '/alexnet_imagenet/weights.h5'
    data_dir = '/alexnet_imagenet/' 
    src_dir = 'data/alexnet_imagenet_src/'
    num_classes = 1000
    batch_size = 50

    model = AlexNet('AlexNet_Imagenet', reload_dir, keras_model_file, data_dir, src_dir, num_classes, batch_size)
    
    model.run(sys.argv)