import numpy as np
import tensorflow as tf

class VGG16:
    def __init__(self, npy_file, npy_mean, trainable):
        self.npy_file = npy_file
        self.npy_mean = npy_mean
        self.var_dict = {}
        self.trainable = trainable

    def build(self, img_holder):
        b, g, r = tf.split(axis=3, num_or_size_splits=3, value=img_holder)
        bgr = tf.concat(axis=3, values=[b - self.npy_mean[0], g - self.npy_mean[1], r - self.npy_mean[2]])

        self.conv1_1 = self.conv_layer(bgr, 3, 64, 'conv1_1')
        self.conv1_2 = self.conv_layer(self.conv1_1, 64, 64, 'conv1_2')
        self.pool1 = self.max_pool(self.conv1_2, 'pool1')

        self.conv2_1 = self.conv_layer(self.pool1, 64, 128, 'conv2_1')
        self.conv2_2 = self.conv_layer(self.conv2_1, 128, 128, 'conv2_2')
        self.pool2 = self.max_pool(self.conv2_2, 'pool2')

        self.conv3_1 = self.conv_layer(self.pool2, 128, 256, 'conv3_1')
        self.conv3_2 = self.conv_layer(self.conv3_1, 256, 256, 'conv3_2')
        self.conv3_3 = self.conv_layer(self.conv3_2, 256, 256, 'conv3_3')
        self.pool3 = self.max_pool(self.conv3_3, 'pool3')

        self.conv4_1 = self.conv_layer(self.pool3, 256, 512, 'conv4_1')
        self.conv4_2 = self.conv_layer(self.conv4_1, 512, 512, 'conv4_2')
        self.conv4_3 = self.conv_layer(self.conv4_2, 512, 512, 'conv4_3')
        self.pool4 = self.max_pool(self.conv4_3, 'pool4')

        self.conv5_1 = self.conv_layer(self.pool4, 512, 512, 'conv5_1')
        self.conv5_2 = self.conv_layer(self.conv5_1, 512, 512, 'conv5_2')
        self.conv5_3 = self.conv_layer(self.conv5_2, 512, 512, 'conv5_3')
        self.pool5 = self.max_pool(self.conv5_3, 'pool5')

        self.fc6 = self.fc_layer(self.pool5, 25088, 4096, 'fc6')
        self.relu6 = tf.nn.relu(self.fc6)

        if self.trainable:
            self.relu6 = tf.nn.dropout(self.relu6, 0.5)

        self.fc7 = self.fc_layer(self.relu6, 4096, 4096, 'fc7')
        self.relu7 = tf.nn.relu(self.fc7)

        if self.trainable:
            self.relu7 = tf.nn.dropout(self.relu7, 0.5)

        self.fc8 = self.fc_layer(self.relu7, 4096, 1000, 'fc8')

        self.prob = tf.nn.softmax(self.fc8, name='prob')

        self.npy_file = None

    def get_var(self, initial_value, name, idx, var_name):
        if self.npy_file is not None and name in self.npy_file:
            value = self.npy_file[name][idx]
        else:
            value = initial_value

        var = tf.Variable(value, name=var_name)

        self.var_dict[(name, idx)] = var

        return var

    def get_conv_var(self, filter_size, in_channels, out_channels, name):
        initial_value = tf.truncated_normal([filter_size, filter_size, in_channels, out_channels], 0.0, 0.001)
        filters = self.get_var(initial_value, name, 0, name + '_filters')

        initial_value = tf.truncated_normal([out_channels], 0.0, 0.001)
        biases = self.get_var(initial_value, name, 1, name + '_biases')

        return filters, biases

    def get_fc_var(self, in_size, out_size, name):
        initial_value = tf.truncated_normal([in_size, out_size], 0.0, 0.001)
        weights = self.get_var(initial_value, name, 0, name + '_weights')

        initial_value = tf.truncated_normal([out_size], 0.0, 0.001)
        biases = self.get_var(initial_value, name, 1, name + '_biases')

        return weights, biases

    def conv_layer(self, bottom, in_channels, out_channels, name):
        with tf.variable_scope(name):
            filt, conv_biases = self.get_conv_var(3, in_channels, out_channels, name)

            conv = tf.nn.conv2d(bottom, filter=filt, strides=[1, 1, 1, 1], padding='SAME')
            bias = tf.nn.bias_add(conv, conv_biases)
            relu = tf.nn.relu(bias)

            return relu

    def fc_layer(self, bottom, in_size, out_size, name):
        with tf.variable_scope(name):
            weights, biases = self.get_fc_var(in_size, out_size, name)

            x = tf.reshape(bottom, [-1, in_size])
            fc = tf.nn.bias_add(tf.matmul(x, weights), biases)

            return fc

    def avg_pool(self, bottom, name):
        return tf.nn.avg_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)

    def max_pool(self, bottom, name):
        return tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)

    def get_var_count(self):
        count = 0
        for var in list(self.var_dict.values()):
            count += np.multiply(var.get_shape().as_list())
        return count
