########################################################################################
# Davi Frossard, 2016                                                                  #
# VGG16 implementation in TensorFlow                                                   #
# Details:                                                                             #
# http://www.cs.toronto.edu/~frossard/post/vgg16/                                      #
#                                                                                      #
# Model from https://gist.github.com/ksimonyan/211839e770f7b538e2d8#file-readme-md     #
# Weights from Caffe converted using https://github.com/ethereon/caffe-tensorflow      #
########################################################################################

import tensorflow as tf
import numpy as np
from scipy.misc import imread, imresize
from imagenet_classes import class_names


class vgg16:
    def __init__(self, imgs, weights=None, sess=None):
        self.imgs = imgs
        self.convlayers()
        self.fc_layers()
        self.probs = self.fc4l# tf.nn.softmax(self.fc4l)
        #if weights is not None and sess is not None:
        #    self.load_weights(weights, sess)


    def convlayers(self):
        self.parameters = []

        # zero-mean input
        with tf.name_scope('preprocess') as scope:
            mean = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32, shape=[1, 1, 1, 3], name='img_mean')
            images = self.imgs-mean

        # conv1_1
        with tf.name_scope('conv1_1') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 3, 64], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(images, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[64], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv1_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv1_2
        with tf.name_scope('conv1_2') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 64, 64], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv1_1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[64], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv1_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool1
        self.pool1 = tf.nn.max_pool(self.conv1_2,
                               ksize=[1, 2, 2, 1],
                               strides=[1, 2, 2, 1],
                               padding='SAME',
                               name='pool1')

        # conv2_1
        with tf.name_scope('conv2_1') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 64, 128], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.pool1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[128], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv2_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv2_2
        with tf.name_scope('conv2_2') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 128, 128], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv2_1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[128], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv2_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool2
        self.pool2 = tf.nn.max_pool(self.conv2_2,
                               ksize=[1, 2, 2, 1],
                               strides=[1, 2, 2, 1],
                               padding='SAME',
                               name='pool2')

        # conv3_1
        with tf.name_scope('conv3_1') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 128, 256], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.pool2, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv3_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv3_2
        with tf.name_scope('conv3_2') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 256], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv3_1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv3_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv3_3
        with tf.name_scope('conv3_3') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 256], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv3_2, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv3_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool3
        self.pool3 = tf.nn.max_pool(self.conv3_3,
                               ksize=[1, 2, 2, 1],
                               strides=[1, 2, 2, 1],
                               padding='SAME',
                               name='pool3')

        # conv4_1
        with tf.name_scope('conv4_1') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.pool3, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv4_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv4_2
        with tf.name_scope('conv4_2') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv4_1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv4_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv4_3
        with tf.name_scope('conv4_3') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv4_2, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv4_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool4
        self.pool4 = tf.nn.max_pool(self.conv4_3,
                               ksize=[1, 2, 2, 1],
                               strides=[1, 2, 2, 1],
                               padding='SAME',
                               name='pool4')

        # conv5_1
        with tf.name_scope('conv5_1') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.pool4, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv5_1 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv5_2
        with tf.name_scope('conv5_2') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv5_1, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv5_2 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # conv5_3
        with tf.name_scope('conv5_3') as scope:
            kernel = tf.Variable(tf.truncated_normal([3, 3, 512, 512], dtype=tf.float32,
                                                     stddev=1e-1), name='weights')
            conv = tf.nn.conv2d(self.conv5_2, kernel, [1, 1, 1, 1], padding='SAME')
            biases = tf.Variable(tf.constant(0.0, shape=[512], dtype=tf.float32),
                                 trainable=True, name='biases')
            out = tf.nn.bias_add(conv, biases)
            self.conv5_3 = tf.nn.relu(out, name=scope)
            self.parameters += [kernel, biases]

        # pool5
        self.pool5 = tf.nn.max_pool(self.conv5_3,
                               ksize=[1, 2, 2, 1],
                               strides=[1, 2, 2, 1],
                               padding='SAME',
                               name='pool4')

    def fc_layers(self):
        # fc1
        with tf.name_scope('fc1') as scope:
            shape = int(np.prod(self.pool5.get_shape()[1:]))
            fc1w = tf.Variable(tf.truncated_normal([shape, 4096],
                                                         dtype=tf.float32,
                                                         stddev=1e-1), name='weights')
            fc1b = tf.Variable(tf.constant(1.0, shape=[4096], dtype=tf.float32),
                                 trainable=True, name='biases')
            pool5_flat = tf.reshape(self.pool5, [-1, shape])
            fc1l = tf.nn.bias_add(tf.matmul(pool5_flat, fc1w), fc1b)
            self.fc1 = tf.nn.relu(fc1l)
            self.parameters += [fc1w, fc1b]

        # fc2
        with tf.name_scope('fc2') as scope:
            fc2w = tf.Variable(tf.truncated_normal([4096, 4096],
                                                         dtype=tf.float32,
                                                         stddev=1e-1), name='weights')
            fc2b = tf.Variable(tf.constant(1.0, shape=[4096], dtype=tf.float32),
                                 trainable=True, name='biases')
            fc2l = tf.nn.bias_add(tf.matmul(self.fc1, fc2w), fc2b)
            self.fc2 = tf.nn.relu(fc2l)
            self.parameters += [fc2w, fc2b]

        # fc3
        with tf.name_scope('fc3') as scope:
            fc3w = tf.Variable(tf.truncated_normal([4096, 1000],
                                                         dtype=tf.float32,
                                                         stddev=1e-1), name='weights')
            fc3b = tf.Variable(tf.constant(1.0, shape=[1000], dtype=tf.float32),
                                 trainable=True, name='biases')
            self.fc3l = tf.nn.bias_add(tf.matmul(self.fc2, fc3w), fc3b)
            self.fc3 = tf.nn.relu(self.fc3l)
            self.parameters += [fc3w, fc3b]

        # fc4
        with tf.name_scope('fc4') as scope:
            fc4w = tf.Variable(tf.truncated_normal([1000, 12],
                                                         dtype=tf.float32,
                                                         stddev=1e-1), name='weights')
            fc4b = tf.Variable(tf.constant(1.0, shape=[12], dtype=tf.float32),
                                 trainable=True, name='biases')
            self.fc4l = tf.nn.bias_add(tf.matmul(self.fc3, fc4w), fc4b)
            self.fc4w = fc4w
            self.fc4b = fc4b

    def load_weights(self, weight_file, sess):
        weights = np.load(weight_file)
        keys = sorted(weights.keys())
        for i, k in enumerate(keys):
            print i, k, np.shape(weights[k])
            sess.run(self.parameters[i].assign(weights[k]))
        sess.run(self.fc4w.initializer)
        sess.run(self.fc4b.initializer)
        print "Pre-trained model loaded. "


def predict_example():
    sess = tf.Session()
    imgs = tf.placeholder(tf.float32, [None, 224, 224, 3])
    vgg = vgg16(imgs, 'vgg16_weights.npz', sess)

    img1 = imread('laska.png', mode='RGB')
    img1 = imresize(img1, (224, 224))

    prob = sess.run(vgg.probs, feed_dict={vgg.imgs: [img1]})[0]
    preds = (np.argsort(prob)[::-1])[0:5]
    for p in preds:
        print class_names[p], prob[p]


def load_train():
    train_label = open('ic-data/train/train.label', 'r')
    name_X = []
    label = []
    buff = train_label.readline()
    while buff:
        image_name = buff.split(' ')[0]
        label1 = buff.split(' ')[1]
        name_X.append(image_name)
        label.append(int(label1)-1)
        buff = train_label.readline()
    # print X
    # print y

    datalen = len(name_X)
    X = np.zeros((datalen, 224, 224, 3), dtype=np.float32)

    it = 0
    sl = 0
    temp_x = np.zeros((100, 224, 224, 3), dtype=np.float32)
    for name in name_X:
        img = imread('ic-data/train/'+name+'.jpg', mode='RGB')
        img = imresize(img, (224, 224))
        temp_x[it] = img
        it += 1
        if it == 100:
            sl += 100
            print sl
            it = 0
            X[sl-100:sl] = temp_x
    y = np.asarray(label)
    print X.shape
    print y.shape
    np.save('X_train.npy', X)
    np.save('y_train.npy', y)


def load_test():
    train_label = open('ic-data/extra/name.txt', 'r')
    name_X = []
    buff = train_label.readline()
    while buff:
        image_name = buff.strip('\n')
        name_X.append(image_name)
        buff = train_label.readline()
    # print X
    # print y

    datalen = len(name_X)
    X = np.zeros((datalen, 224, 224, 3), dtype=np.float32)

    it = 0
    sl = 0
    temp_x = np.zeros((100, 224, 224, 3), dtype=np.float32)
    for name in name_X:
        img = imread('ic-data/extra/'+name, mode='RGB')
        img = imresize(img, (224, 224))
        temp_x[it] = img
        it += 1
        if it == 100:
            sl += 100
            print sl
            it = 0
            X[sl-100:sl] = temp_x
    temp_x = temp_x.astype('float32')
    X[sl:sl+it] = temp_x[:it]
    print X.shape
    np.save('X_test.npy', X)

from sklearn.model_selection import train_test_split

def train():
    learning_rate = 0.001
    training_epoch = 10
    batch_size = 40
    display_step = 10
    n_classes = 12

    sess = tf.Session()
    X = np.load('X_train.npy')
    y = np.load('y_train.npy')
    onehot = tf.one_hot(y, n_classes)
    y = sess.run(onehot)
    print y.shape

    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size = 0.2, random_state = 47)
    print y_train.shape

    imgs = tf.placeholder(tf.float32, [None, 224, 224, 3])
    imgs_y = tf.placeholder(tf.float32, [None, n_classes])
    vgg = vgg16(imgs, 'vgg16_weights.npz', sess)

    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=vgg.probs, labels=imgs_y))
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)

    correct_pred = tf.equal(tf.argmax(vgg.probs, 1), tf.argmax(imgs_y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    #init = tf.global_variables_initializer()
    init = tf.initialize_all_variables()

    sess.run(init)
    if sess is not None:
        vgg.load_weights('vgg16_weights.npz', sess)

    train_size = X_train.shape[0]
    for epoch in range(training_epoch):
        step = 1
        print "Epoch: " + str(epoch)
        while step * batch_size < train_size:
            batch_x = X_train[(step-1)*batch_size:step*batch_size]
            batch_y = y_train[(step-1)*batch_size:step*batch_size]
            sess.run(optimizer, feed_dict={vgg.imgs: batch_x, imgs_y:batch_y})
            # print "Alive?"
            if step % display_step == 0:
                loss, acc = sess.run([cost, accuracy], feed_dict={vgg.imgs:batch_x, imgs_y:batch_y})
                print "Iter: " + str(step*batch_size) + ", Minibatch Loss= "+"{:.6f}".format(loss) + \
                    ", Training Accuracy= " + "{:.5f}".format(acc)
            step += 1
        batch_x = X_train[(step-1)*batch_size:]
        batch_y = y_train[(step-1)*batch_size:]
        sess.run(optimizer, feed_dict={vgg.imgs: batch_x, imgs_y:batch_y})


    print "Begin validation. "
    val_size = X_val.shape[0]
    step = 1
    res = 0
    while step * batch_size < val_size:
        batch_x = X_val[(step-1)*batch_size:step*batch_size]
        batch_y = y_val[(step-1)*batch_size:step*batch_size]
        acc = sess.run(accuracy, feed_dict={vgg.imgs:batch_x, imgs_y:batch_y})
        res += acc*batch_size
        step += 1

    batch_x = X_val[(step-1)*batch_size:]
    batch_y = y_val[(step-1)*batch_size:]
    acc = sess.run(accuracy, feed_dict={vgg.imgs:batch_x, imgs_y:batch_y})
    res += acc*(val_size-(step-1)*batch_size)
    res /= val_size
    print "Validation Accuracy: " + "{:.5f}".format(res)


if __name__ == '__main__':
    # load_train()
    # load_test()
    train()
