import keras
from keras.models import Sequential
from keras.layers import Dense, InputLayer, Flatten, Conv2D, MaxPooling2D, GlobalMaxPool2D, GlobalAveragePooling2D, BatchNormalization
from keras.preprocessing.image import ImageDataGenerator
from keras.layers.core import Activation

import os
from shutil import move , copy
from os.path import join
from os.path import split
path = join('data')
classes = os.listdir(join(path, 'train'))
classes

classes = [c for c in classes if c[0] != '.']
classes

import numpy as np
from numpy.random import choice

import glob

valid_ratio = 0.2
num_sample_images = 10
valid_path = join(path, 'valid')
if not os.path.exists(valid_path):
    os.makedirs(valid_path)

    for c in classes:
        cl_valid_path = join(valid_path, c)
        cl_train_path = join(path, 'train', c)
        cl_sample_path = join(path, 'sample', c)
        files = glob.glob(join(cl_train_path, '*.jpg'))
        os.makedirs(cl_valid_path)
        valid = choice(files, int(np.floor(0.2*len(files))), replace=False)
        [move(v, join(cl_valid_path, split(v)[1])) for v in valid]
        # We don't want the sample to contain the validation data, so redo the glob after moving. 
        files = glob.glob(join(cl_train_path, '*.jpg'))
        os.makedirs(cl_sample_path)
        sample = choice(files, num_sample_images, replace=False)
        [copy(s, join(cl_sample_path, split(s)[1])) for s in sample]

classes

os.listdir(join(path,'valid'))

bs = 1
gen = ImageDataGenerator()
train_gen = gen.flow_from_directory(join(path,'train'), batch_size=bs, shuffle=False)
valid_gen = gen.flow_from_directory(join(path,'valid'), batch_size=bs, shuffle=False)
sample_gen = gen.flow_from_directory(join(path,'sample'), batch_size=bs, shuffle=False)

import bcolz
from tqdm import tqdm
import os.path

def save_generator(gen, data_dir, labels_dir):
    """
    Save the output from a generator without loading all images into memory. 
    
    Does not return anything, instead writes data to disk.
    
    :gen: A Keras ImageDataGenerator object
    :data_dir: The folder name to store the bcolz array representing the features in. 
    :labels_dir: The folder name to store the bcolz array representing the labels in.
    :mode: the write mode. Set to 'a' for append, set to 'w' to overwrite existing data and 'r' to read only. 
    
    """
    for directory in [data_dir, labels_dir]:
        if not os.path.exists(directory):
            os.makedirs(directory)
    
    num_samples = gen.samples
    
    d,l = gen.__next__()
    
    data = bcolz.carray(d, rootdir=data_dir, mode='w')
    labels = bcolz.carray(l, rootdir=labels_dir, mode='w')

    for i in tqdm(range(num_samples-1)):
        d, l = gen.__next__()
        data.append(d)
        labels.append(l)
    data.flush()
    labels.flush()


trn_data = join('bdat','train','data')
trn_label = join('bdat','train','label')
val_data = join('bdat','valid','data')
val_label = join('bdat','valid','label')
samp_data = join('bdat','sample','data')
samp_label = join('bdat','sample','label')

#save_generator(train_gen, trn_data, trn_label)
#save_generator(valid_gen, val_data, val_label)
#save_generator(sample_gen, samp_data, samp_label)

data = bcolz.open(trn_data)
labels = bcolz.open(trn_label)
val_data = bcolz.open(val_data)
val_labels = bcolz.open(val_label)

print(data.shape)
print(labels.shape)

simple = Sequential([
    InputLayer((256,256,3)),
    Conv2D(16,(3,3), activation='relu'),
    MaxPooling2D(4),
    Conv2D(16, (3,3), activation='relu'),
    Conv2D(8, (3,3), activation='relu'),
    GlobalMaxPool2D()
  ])
simple.summary()

simple.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

simple.fit(data, labels, batch_size=16, validation_data=(val_data,val_labels))

print('blah')

from keras.applications.xception import Xception

xModel = Xception(include_top=False, input_shape=(256,256,3))

def save_predictions(model, data, rootdir, batch_size):
    """
    This function will use BColz to save the predictions from a model. This is useful when you want to get the features from a
    pretrained net and build something ontop of it without re-evaluating the network every time. 
    
    This function does not return anything and writes stuff to disk.
    
    :model: A keras model.
    :data: A Numpy dataframe, it is assumed that the first index is the batch index. 
    :roodir: The directory to store the bcolz data
    :batchsize: The number of samples to run. Will depend upon your hardware. 
    """
    output = bcolz.carray(model.predict(data[0:batch_size]), rootdir=rootdir, mode='w')
    
    for i in tqdm(range(batch_size, data.shape[0], batch_size)):
        end = i+batch_size if i+batch_size < data.shape[0] else data.shape[0]
        output.append(model.predict(data[i:end]))
    output.flush()
    

#os.makedirs(join('bdat','pretrained','data'))
save_predictions(xModel, data, join('bdat','pretrained','data'), 2)

save_predictions(xModel, val_data, join('bdat','pretrained','valdata'), 2)

xCdata = bcolz.open(join('bdat','pretrained','data'))
xValdata = bcolz.open(join('bdat','pretrained','valdata'))

xCdata = xModel.predict(data[:],verbose=1, batch_size=2)

xValdata = xModel.predict(val_data, verbose=1, batch_size=2)

# Comment out if you are reading off an SSD.
xCdata = xCdata[:]
xValdata = xValdata[:]

pretrained_top = Sequential([
    InputLayer((8,8,2048)),
    GlobalAveragePooling2D(),
    Dense(8, activation='softmax')
])
pretrained_top.compile(optimizer='adam', loss='categorical_crossentropy',metrics=['accuracy'])

pretrained_top.fit(xCdata[:], labels, validation_data=(xValdata, val_labels), batch_size=128)

conv_top = Sequential([InputLayer((8,8,2048)),
                      Conv2D(8, (3,3), padding='same', activation='relu'),
                       BatchNormalization(),
                      GlobalAveragePooling2D(),
                       Activation('softmax')
                      ])
conv_top.compile(optimizer='adam', loss='categorical_crossentropy',metrics=['accuracy'])

conv_top.summary()

conv_top.fit(xCdata, labels, validation_data=(xValdata, val_labels), batch_size=128)

full_layers = [xModel, conv_top.layers[1]]

full_model = Sequential(full_layers)
full_model.compile(optimizer='adam', loss='categorical_crossentropy',metrics=['accuracy'])
full_model.summary()

pred = full_model.predict(data[0][np.newaxis,:,:,:])[0]
pred.shape

import matplotlib.pyplot as pl
%matplotlib inline
import scipy.misc

i = np.argmax(conv_top.predict(xCdata[0][np.newaxis,:,:,:]))
pl.figure()
pl.imshow(data[0])
pl.imshow(scipy.misc.imresize(pred[i], (240,240), interp='nearest'), alpha=0.5)
pl.title(classes[i])