import tensorflow as tf

from ucnn import input_image


def _conv2d(value, weight):
    """conv2d returns a 2d convolution layer with full stride."""
    return tf.nn.conv2d(
        value,
        weight,
        strides=[1, 1, 1, 1],
        padding='SAME'
    )


def _max_pool_2x2(value, name):
    """max_pool_2x2 downsamples a feature map by 2X."""
    return tf.nn.max_pool(
        value,
        ksize=[1, 2, 2, 1],
        strides=[1, 2, 2, 1],
        padding='SAME',
        name=name
    )


def _weight_variable(name, shape):
    """weight_variable generates a weight variable of a given shape."""
    with tf.device('/cpu:0'):
        initializer = tf.truncated_normal_initializer(stddev=0.1)
        var = tf.get_variable(
            name,
            shape,
            initializer=initializer,
            dtype=tf.float32
        )
    return var


def _bias_variable(name, shape):
    """bias_variable generates a bias variable of a given shape."""
    with tf.device('/cpu:0'):
        initializer = tf.constant_initializer(0.1)
        var = tf.get_variable(name, shape, initializer=initializer, dtype=tf.float32)
    return var


def loss(logits, labels):
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
        labels=labels, logits=logits, name='corss_entropy_per_example')
    cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')
    tf.add_to_collection('losses', cross_entropy_mean)
    return tf.add_n(tf.get_collection('losses'), name='total_loss')


def training(losses):
    optimizer = tf.train.AdamOptimizer(1e-4)
    train_op = optimizer.minimize(losses)
    return train_op


def evaluation(logits, labels):
    correct_prediction = tf.equal(tf.argmax(logits, 2), tf.argmax(labels, 2))
    correct_batch = tf.reduce_mean(tf.cast(correct_prediction, tf.int32), 1)
    return tf.reduce_sum(tf.cast(correct_batch, tf.float32))


def output(logits):
    return tf.argmax(logits, 2)


class CNNModel(object):
    def __init__(self, configuration):
        self.IMAGE_HEIGHT = configuration['image_height']
        self.IMAGE_WIDTH = configuration['image_width']
        self.CHARS_NUM = configuration['chars_length']
        self.CLASSES_NUM = configuration['charsets_length']

    def inputs(self, train, batch_size):
        return input_image.inputs(
            train,
            batch_size,
            self.IMAGE_HEIGHT,
            self.IMAGE_WIDTH,
            self.CHARS_NUM,
            self.CLASSES_NUM
        )

    def model(self, images, keep_prob):
        images = tf.reshape(images, [-1, self.IMAGE_HEIGHT, self.IMAGE_WIDTH, 1])

        with tf.variable_scope('conv1') as scope:
            kernel = _weight_variable('weights', shape=[3, 3, 1, 64])
            biases = _bias_variable('biases', [64])
            pre_activation = tf.nn.bias_add(_conv2d(images, kernel), biases)
            conv1 = tf.nn.relu(pre_activation, name=scope.name)

        pool1 = _max_pool_2x2(conv1, name='pool1')

        with tf.variable_scope('conv2') as scope:
            kernel = _weight_variable('weights', shape=[3, 3, 64, 64])
            biases = _bias_variable('biases', [64])
            pre_activation = tf.nn.bias_add(_conv2d(pool1, kernel), biases)
            conv2 = tf.nn.relu(pre_activation, name=scope.name)

        pool2 = _max_pool_2x2(conv2, name='pool2')

        with tf.variable_scope('conv3') as scope:
            kernel = _weight_variable('weights', shape=[3, 3, 64, 64])
            biases = _bias_variable('biases', [64])
            pre_activation = tf.nn.bias_add(_conv2d(pool2, kernel), biases)
            conv3 = tf.nn.relu(pre_activation, name=scope.name)

        pool3 = _max_pool_2x2(conv3, name='pool3')

        with tf.variable_scope('conv4') as scope:
            kernel = _weight_variable('weights', shape=[3, 3, 64, 64])
            biases = _bias_variable('biases', [64])
            pre_activation = tf.nn.bias_add(_conv2d(pool3, kernel), biases)
            conv4 = tf.nn.relu(pre_activation, name=scope.name)

        pool4 = _max_pool_2x2(conv4, name='pool4')

        with tf.variable_scope('local1') as scope:
            batch_size = images.get_shape()[0].value
            reshape = tf.reshape(pool4, [batch_size, -1])
            dim = reshape.get_shape()[1].value
            weights = _weight_variable('weights', shape=[dim, 1024])
            biases = _bias_variable('biases', [1024])
            local1 = tf.nn.relu(tf.matmul(reshape, weights) + biases, name=scope.name)

        local1_drop = tf.nn.dropout(local1, keep_prob)

        with tf.variable_scope('softmax_linear') as scope:
            weights = _weight_variable('weights', shape=[1024, self.CHARS_NUM * self.CLASSES_NUM])
            biases = _bias_variable('biases', [self.CHARS_NUM * self.CLASSES_NUM])
            softmax_linear = tf.add(tf.matmul(local1_drop, weights), biases, name=scope.name)

        return tf.reshape(softmax_linear, [-1, self.CHARS_NUM, self.CLASSES_NUM])
