# -*- coding: utf-8 -*-

import numpy as np
import warnings
import sys
import os
import argparse
import re

from keras.models import Model
from keras.models import Sequential
from keras.layers import Flatten
from keras.layers import Dense
from keras.layers import Input
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.preprocessing import image
from keras.utils import layer_utils
from keras import backend as K
from keras.applications.imagenet_utils import preprocess_input
from keras.applications.imagenet_utils import _obtain_input_shape
from keras.engine.topology import get_source_inputs

from sklearn.svm import SVC
from sklearn.decomposition import IncrementalPCA
from sklearn.decomposition import PCA
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score

def VGG16(include_top=True, weights='imagenet',
          input_tensor=None, input_shape=None,
          pooling=None,
          classes=1000):

    WEIGHTS_PATH = 'vgg16_weights_tf_dim_ordering_tf_kernels.h5'
    # Determine proper input shape
    input_shape = _obtain_input_shape(input_shape,
                                      default_size=224,
                                      min_size=48,
                                      data_format=K.image_data_format(),
                                      include_top=include_top)

    if input_tensor is None:
        img_input = Input(shape=input_shape)
    else:
        if not K.is_keras_tensor(input_tensor):
            img_input = Input(tensor=input_tensor, shape=input_shape)
        else:
            img_input = input_tensor
    # Block 1
    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1')(img_input)
    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv2')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x)

    # Block 2
    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv2')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x)

    # Block 3
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv1')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv2')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x)

    # Block 4
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x)

    # Block 5
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv1')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv2')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv3')(x)
    x = MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x)

    x = Flatten(name='flatten')(x)
    x = Dense(4096, activation='relu', name='fc1')(x)
    x = Dense(4096, activation='relu', name='fc2')(x)
    x = Dense(classes, activation='softmax', name='predictions')(x)

    # Ensure that the model takes into account
    # any potential predecessors of `input_tensor`.
    if input_tensor is not None:
        inputs = get_source_inputs(input_tensor)
    else:
        inputs = img_input
    # Create model.
    model = Model(inputs, x, name='vgg16')

    # load weights
    model.load_weights(WEIGHTS_PATH)

    return model

def natural_sort(l):
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ]
    return sorted(l, key=alphanum_key)

if __name__ == '__main__':
    #PARSE PROGRAM ARGUMENTS
    parser = argparse.ArgumentParser()
    parser.add_argument('imgBasePath', help="Insert the path to the Image Database you want to use\nI.E. \tCaltech101_600/\n\timages/Caltech101_600/")
    parser.add_argument('train', help='True for training, False for predicting', type=str)
    parser.add_argument('--n_components', help='Number of components to reduce in PCA', type=int, nargs='*', default=[16])
    parser.add_argument('-k', help='Number of neighbors in k-NN', nargs='*', type=int, default=[1])
    #parser.add_argument('--predict', help='Predict the class for the input Image Database, using the pre-trained VGG-16')
    args = parser.parse_args()

    if args.train.lower() == 'train':
        #LOAD THE PRE-TRAINED MODEL AND EXTRACT THE 15TH LAYER
        base_model = VGG16(weights='imagenet')
        model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc2').output)

        ipca_preds = np.array([[]])
        ipca_classes = np.array([])

        #GETS ALL IMAGES IN THE SUBDIRECTORIES
        for root, dirs, files in os.walk(args.imgBasePath):
            for f in natural_sort(files):
                if f.endswith(".jpg") or f.endswith('.jpeg'):
                    #print('File:\t' + os.path.join(root,f) + '\t|Class:\t' + root.split('/')[-1])

                    if ipca_classes.size > 0:
                        ipca_classes = np.append(ipca_classes, root.split('/')[-1])
                    else:
                        ipca_classes = np.array([root.split('/')[-1]])

                    #LOAD AND PREPROCESS THE IMAGE
                    img = image.load_img(os.path.join(root, f), target_size=(224, 224))
                    x = image.img_to_array(img)
                    x = np.expand_dims(x, axis=0)
                    x = preprocess_input(x)

                    prediction = model.predict(x)

                    #RUN THE IMAGE THROUGH THE MODEL
                    if ipca_preds.size > 0:
                        ipca_preds = np.append(ipca_preds, prediction, axis=0)
                    else:
                        ipca_preds = prediction

        print('Database: ', args.imgBasePath)

        f_name = args.imgBasePath.split('/')[-2];
        if not os.path.exists('build/'):
            os.makedirs('build/')

        with open('build/'+f_name+'_4096_weights.npy','wb+') as f_ipca :
            np.save(f_ipca, ipca_preds)

        with open('build/'+f_name+'_4096_classes.npy','wb+') as c_ipca:
            np.save(c_ipca, ipca_classes)

        f_ipca.close()
        c_ipca.close()

        print('N_Components: 4096')
        print('Predictions ', ipca_preds.shape)
        print('IPCA Output None')

        #RUN THE EXTRACTED VALUES THROUGH PCA TO REDUCE THE NUMBER OF FEATURES
        for n in args.n_components:
            #CREATE THE WEIGHTS AND CLASSES FILES

            ipca = IncrementalPCA(n_components=n, whiten=True, copy=True)
            ipca_output = ipca.fit_transform(ipca_preds)

            with open('build/'+f_name+'_'+str(n)+'_weights.npy','wb+') as f_ipca:
                np.save(f_ipca, ipca_output)

            with open('build/'+f_name+'_'+str(n)+'_classes.npy','wb+') as c_ipca:
                np.save(c_ipca, ipca_classes)

            f_ipca.close()
            c_ipca.close()

            print('N_Components: ', n)
            print('Predictions ', ipca_preds.shape)
            print('IPCA Output ', ipca_output.shape)

    if args.train.lower() == 'test':
        f_name = args.imgBasePath.split('/')[-2];
        print('Database: ', f_name)

        for n in args.n_components:
            print()
            for a in args.k:

                with open('build/'+f_name+'_'+str(n)+'_weights.npy','rb') as f_data :
                    data = np.load(f_data)
                    for d in f_data:
                        data = np.append(data, d)
                with open('build/'+f_name+'_'+str(n)+'_classes.npy','rb') as f_class:
                    classes = np.load(f_class)
                    for c in f_class:
                        classes = np.append(classes, c)
                f_data.close()
                f_class.close()

                knn = KNeighborsClassifier(n_neighbors=a)
                scores = cross_val_score(knn, data, classes, cv=5, n_jobs=-1)
                print('k-NN: %d Neighbors' % (a))
                print('\tN_Components: ', n)
                print('\tAccuracy: %0.4f +/- %0.4f' % (scores.mean(), scores.std()))

            svm = SVC(kernel='linear')
            scores = cross_val_score(svm, data, classes, cv=5, n_jobs=-1)
            print('SVC Linear')
            print('\tN_Components: ', n)
            print('\tAccuracy: %0.4f +/- %0.4f' % (scores.mean(), scores.std()))