from __future__ import print_function
import numpy as np
np.random.seed(42)                      

from keras.datasets import mnist
from keras.models import Sequential
from keras.layers import Dense, Activation, Flatten
from keras.layers import Convolution2D, MaxPooling2D
from keras.utils import np_utils
from keras import backend as K
import pickle
import sys

# Posto optimizacioni metod koristi mini-batch treniranje, definisemo velicinu svakog paketa.
batch_size = 10

# Definisemo broj klasa na koje vrsimo predikciju.
nb_classes = 2

# Broj epoha, tj. koliko puta ponavljamo proces optimizacije.
nb_epoch = 10

# Dimenzije fotografije
img_rows, img_cols = 112, 92

# Broj filtera svakog neurona iz konvolutivnog sloja
nb_filters = 8

# Velicina agregacije
pool_size = (2, 2)

# Velicina odsecka slike koji ce se posmatrati pri konvoluciji
kernel_size = (3, 3)

# Ucitavanje podataka
fr = open('images.data', 'rb')
(X, y) = pickle.load(fr)

# Podela na trening i test skup zarad procene kvaliteta modela
X_train = X[0:600,:,:]
y_train = y[0:600,]
X_test = X[600:,:,:]
y_test = y[600:,]

# Keras moze koristiti biblioteke Theano i TensorFlow za rad sa tenzorima
# API za treniranje je isti, medjutim format podataka na ulazu se razlikuje
# pa je potrebno formatirati ulaz u zavisnosti od toga koja biblioteka
# (backend) se koristi.
if K.image_dim_ordering() == 'th':
    X_train = X_train.reshape(X_train.shape[0], 1, img_rows, img_cols)
    X_test = X_test.reshape(X_test.shape[0], 1, img_rows, img_cols)
    input_shape = (1, img_rows, img_cols)
else:
    X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 1)
    X_test = X_test.reshape(X_test.shape[0], img_rows, img_cols, 1)
    input_shape = (img_rows, img_cols, 1)

# Ulaz moramo pretvoriti u numpy matrice i zarad brze konvergencije optimizatora
# skalirati ga na interval [0,1]
X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
X_train /= 255
X_test /= 255
print('X_train shape:', X_train.shape)
print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

# Ciljnu promenljivu definisemo kao kategoricku.
# Primer:
#   Imamo 5 klasa i nasa promenljiva uzima vrednosti
#   iz skupa {1, 2, 3, 4, 5}. Za potrebe racunanja
#   gradijenta je neophodno transformisati je u vektor
#   oblika:
#       [1, 0, 0, 0, 0] za vrednost 0,
#       [0, 1, 0, 0, 0] za vrednost 1,
#       ...
#       [0, 0, 0, 0, 1] za vrednost 5
Y_train = np_utils.to_categorical(y_train, nb_classes)
Y_test = np_utils.to_categorical(y_test, nb_classes)

# Pravimo sekvencijalni model (model sa propagacijom unapred)
model = Sequential()

# Dodajemo prvo konvolutivni sloj sa ReLU funkcijom aktivacije
model.add(Convolution2D(nb_filters, kernel_size[0], kernel_size[1], border_mode='valid', input_shape=input_shape, name='conv1'))
model.add(Activation('relu'))

# Nakon toga sledi sloj agregacije
model.add(MaxPooling2D(pool_size=pool_size))

# Jos jedan sloj konvolucije sa ReLU funkcijom aktivacije i nakon njega sloj agregacije
model.add(Convolution2D(nb_filters, kernel_size[0]/2, kernel_size[1]/2, name='conv2'))
model.add(Activation('relu'))
model.add(MaxPooling2D(pool_size=pool_size))

# Posle ovih slojeva dolaze "obicni" slojevi, tj. puno-povezana mreza. Ona na ulazu
# ocekuje vektor. Do sada smo imali matricne izlaze (filtere) koje treba rasporediti
# u vektor.
model.add(Flatten())
model.add(Dense(8))
model.add(Activation('relu'))
model.add(Dense(nb_classes))
# Koristimo funkciju aktivacije softmax kako bismo izlaz mogli da interpretiramo
# kao funkciju raspodele
model.add(Activation('softmax'))

# Prevodimo keras model u model koji odgovara backend biblioteci koja se koristi.
model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])

# Vrsimo treniranje modela
model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch, verbose=1)

# Cuvamo arhitekturu modela i dobijene (naucene) tezine.
f = open('model.arch', 'w')
model.save_weights('model.weights', overwrite=True)
f.write(model.to_yaml())
f.close()

# Vrsimo evaluaciju modela na test skupu.
score = model.evaluate(X_test, Y_test, verbose=0)
print('Test score:', score[0])
print('Test accuracy:', score[1])

