# Location of digested data
input_directory = '/digested/'

# Desired location to save temporary PNG outputs:
png_directory = '/digested_png/'

# Desired location to save trained model
model_directory = '/model_directory/'

"""
Explanation: ------------- User's settings -------------
End of explanation
"""


import keras.applications

# Define native shape of the transferred model (refer to model documentation)
shape=(197,197,3)

loaded_model = keras.applications.ResNet50(include_top=False, 
                                           weights='imagenet',
                                           input_tensor = keras.layers.Input( shape ),
                                           classes = classes
                                          )

"""
Explanation: Download pre-trained model:
End of explanation
"""


import keras
import keras.preprocessing.image
import numpy
import pickle
import os.path
import tensorflow
import digest
import time

"""
Explanation: Pre-built Keras models:
Very Deep Convolutional Networks for Large-Scale Image Recognition https://arxiv.org/abs/1409.1556
- VGG16
- VGG19
Deep Residual Learning for Image Recognition https://arxiv.org/abs/1512.03385
- ResNet50
Inception Architecture for Computer Vision https://arxiv.org/abs/1512.00567
- Inception v3
Note:
Pre-trained networks are not well-suited for bioimaging analysis of multichannel fluorescent images, i.e. they were built for 3 channels of RGB images. In contrast to our built-in CNN, here one should selectively choose maximum 3 channels that provide sufficient information for making prediction, and should omit the channels that may introduce noise.
To make it plausible for the pre-trained network’s requirement of RGB and computationally efficient, single-channel images of an object will be stacked to form an RGB image and saved in PNG format. User should be informed that several PNG files will be stored in the hard disk as the location of png_directory.
------------- (semi)-Automatic -------------
End of explanation
"""


# -------- If using Tensorflow-GPU: -------- #

configuration = tensorflow.ConfigProto()

configuration.gpu_options.allow_growth = True

configuration.gpu_options.visible_device_list = "0"

session = tensorflow.Session(config=configuration)

keras.backend.set_session(session)


# -------- If using Tensorflow (CPU) : -------- #

# configuration = tensorflow.ConfigProto()

# session = tensorflow.Session(config=configuration)

# keras.backend.set_session(session)

if not os.path.exists(model_directory):
    os.makedirs(model_directory)

"""
Explanation: Configure CPU/GPU devices:
End of explanation
"""


class_weights = pickle.load(open(os.path.join(input_directory, "class_weights.sav"), 'rb'))

classes = len( pickle.load(open(os.path.join(input_directory, "class_names.sav"), 'rb')) )

training_x = numpy.load(os.path.join(input_directory, "training_x.npy"))

training_y = numpy.load(os.path.join(input_directory, "training_y.npy"))

validation_x = numpy.load(os.path.join(input_directory, "validation_x.npy"))

validation_y = numpy.load(os.path.join(input_directory, "validation_y.npy"))

"""
Explanation: Load data and labels:
End of explanation
"""


%%capture
digest.save_png(training_x, os.path.join(png_directory,"Training") )
digest.save_png(validation_x, os.path.join(png_directory,"Validation") )

"""
Explanation: Stack single-channel images (maximum 3 channels) and store to PNG files
End of explanation
"""


options = {
    "rotation_range" : 180
    "featurewise_center": True,
    "featurewise_std_normalization": True,
    "horizontal_flip": True,
    "vertical_flip": True
}

training_generator = keras.preprocessing.image.ImageDataGenerator( **options )

training_generator = training_generator.flow_from_directory(
    batch_size=32,
    color_mode="rgb",
    directory= os.path.join(png_directory,"Training"),
    target_size=(shape[0], shape[1])
)

validation_generator = keras.preprocessing.image.ImageDataGenerator()

validation_generator = validation_generator.flow_from_directory(
    batch_size=32,
    color_mode="rgb",
    directory= os.path.join(png_directory,"Validation"),
    target_size=(shape[0], shape[1])
)

loaded_model.summary()

# first: train only the top layers (which were randomly initialized)
# i.e. freeze all convolutional layers
for layer in loaded_model.layers:
    layer.trainable = False

x = loaded_model.output

x = keras.layers.GlobalAveragePooling2D()(x)

x = keras.layers.Dense(2048, activation='relu')(x)

# # A logistic prediction layer
predictions = keras.layers.Dense(classes, activation='softmax')(x)

# # this is the model we will train
model = keras.models.Model(input=loaded_model.input, output = predictions)

model.summary()

optimizer = keras.optimizers.Adam(0.00001)

loss = keras.losses.categorical_crossentropy

metrics = [
    "accuracy"
]

model.compile(optimizer, loss, metrics)

"""
Explanation: Note:
Since several PNG files will be stored during data digestion, if user desires to clean png_directory, for e.g. to start over, it is highly recommended to perform this clearance in command line.
In Linux/UNIX system:
mkdir empty_dir
rsync -a --delete empty_dir/ png_directory/
End of explanation
"""


start = time.time()

# -------- If using Tensorflow (CPU) : -------- #

# init = tensorflow.global_variables_initializer()

# session.run(init)

# -------------------------------------------- #

model.fit_generator(
    callbacks=[
        keras.callbacks.CSVLogger( os.path.join(model_directory, 'training.csv') ),
        keras.callbacks.ModelCheckpoint( os.path.join(model_directory, 'checkpoint.hdf5') )
    ],
    epochs=18,
    class_weight = class_weights,
    verbose = 0, # ON/OFF printing output
    generator = training_generator,
    max_q_size = 256,
    steps_per_epoch=2500,
    validation_data = validation_generator,
    validation_steps = 2500
)

stop = time.time()
print(stop - start, "s")

model.save( os.path.join(model_directory, 'model.h5') )

session.close()

"""
Explanation: Train the network
End of explanation
"""


