#!/usr/bin/python

data_augmentation = True

import numpy
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Convolution2D, MaxPooling2D
from keras.optimizers import SGD
from keras.utils import np_utils

batch_size = 32
nb_classes = 10
nb_epoch = 200

img_channels, img_rows, img_cols = 3, 32, 32

orig_trn_X = numpy.loadtxt("trn_X.csv", delimiter = ",", dtype = "uint8")
trn_X = orig_trn_X.reshape(50000, 3, 32, 32)
orig_trn_y = numpy.loadtxt("trn_y.txt", delimiter = ",", dtype = "uint8")
trn_y = orig_trn_y.reshape(50000, 1)
orig_tst_X = numpy.loadtxt("tst_X.csv", delimiter = ",", dtype = "uint8")
tst_X = orig_tst_X.reshape(10000, 3, 32, 32)

print("trn_X shape:", trn_X.shape)
print(trn_X.shape[0], " train samples")
print(tst_X.shape[0], " test samples")

trn_Y = np_utils.to_categorical(trn_y, nb_classes)

model = Sequential()
model.add(Convolution2D(32, 3, 3, border_mode="same", input_shape = (img_channels, img_rows, img_cols)))
model.add(Activation("relu"))
model.add(Convolution2D(32, 3, 3))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size = (2, 2)))
model.add(Dropout(0.25))
model.add(Convolution2D(64, 3, 3, border_mode = "same"))
model.add(Activation("relu"))
model.add(Convolution2D(64, 3, 3))
model.add(Activation("relu"))
model.add(MaxPooling2D(pool_size = (2, 2)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(512))
model.add(Activation("relu"))
model.add(Dropout(0.5))
model.add(Dense(nb_classes))
model.add(Activation("softmax"))

sgd = SGD(lr = 0.01, decay = 1e-6, momentum = 0.9, nesterov = True)
model.compile(loss = "categorical_crossentropy", optimizer = sgd, metrics = ["accuracy"])

trn_X = trn_X.astype("float32")
tst_X = tst_X.astype("float32")
trn_X /= 255
tst_X /= 255

if not data_augmentation:
    print("Training without data augmentation.")
    model.fit(trn_X, trn_Y,
              batch_size = batch_size,
              nb_epoch = nb_epoch,
              shuffle = True)
else:
    print("Training with real-time data augmentation.")
    datagen = ImageDataGenerator(
                  featurewise_center = False,             # set input mean to 0 over the dataset
                  samplewise_center = False,              # set each sample mean to 0
                  featurewise_std_normalization = False,  # divide inputs by std of the dataset
                  samplewise_std_normalization = False,   # divide each input by its std
                  zca_whitening = False,                  # apply ZCA whitening
                  rotation_range = 0,                     # randomly rotate images in the range (degrees, 0 to 180)
                  width_shift_range = 0.1,                # randomly shift images horizontally (fraction of total width)
                  height_shift_range = 0.1,               # randomly shift images vertically (fraction of total height)
                  horizontal_flip = True,                 # randomly flip images
                  vertical_flip = False)                  # randomly flip images

    # compute quantities required for featurewise normalization
    datagen.fit(trn_X)

    # fit the model on the batches generated by datagen.flow()
    model.fit_generator(datagen.flow(trn_X, trn_Y,
                                     batch_size = batch_size),
                                     samples_per_epoch = trn_X.shape[0],
                                     nb_epoch = nb_epoch)

predictions = model.predict(tst_X)
output = open("predictions.csv", "w")
output.write("index," + ",".join([ "probability" + str(i) for i in range(predictions.shape[1]) ]) + "\n")
for i in range(predictions.shape[0]):
    output.write(str(i) + "," + ",".join([ str(prediction) for prediction in predictions[i,:] ]) + "\n")
output.close()