# Path: util/model.py
# class Model(object):
#     def __init__(self, inputs, label, loss, pred, debug=None):
#         self.inputs = inputs
#         self.label = label
#         self.loss = loss
#         self.pred = pred
#         self.debug = debug
# 
# def conv2d(name, x, out_channel, kernel, stride, nl=None):
#     name = functools.partial('{}-{}'.format, name)
#     # W
#     shape = [kernel, kernel, x.get_shape()[3].value, out_channel]
#     W = weight_variable(name('W'), shape)
#     # b
#     b = bias_variable(name('b'), [out_channel])
#     # conv
#     y = tf.nn.conv2d(x, W, strides=[1,stride,stride,1], padding='SAME', name=name('op'))
#     y = y + b
#     # nonlinearty
#     return nl(y) if nl else y
# 
# def fc_layer(name, x, in_shape, out_shape, nl=None):
#     name = functools.partial('{}-{}'.format, name)
# 
#     in_chanel = np.prod(in_shape)
#     out_chanel = np.prod(out_shape)
#     # W
#     W = weight_variable(name('W'), [in_chanel, out_chanel])
#     # b
#     b = bias_variable(name('b'), [out_chanel])
#     # fc
#     x = tf.reshape(x, [-1, in_chanel])
#     x = tf.matmul(x, W) + b
#     out_shape = [-1] + out_shape
#     x = tf.reshape(x, out_shape)
# 
#     return nl(x) if nl else x


from config import config
from util.model import Model, conv2d, fc_layer
import tensorflow as tf
import functools
    import numpy as np
    for layer in range(2):
        x_branch = conv_pip(name('conv%d'%layer), x)
        x = tf.concat(3, [x,x_branch], name=name('concate%d'%layer))

    x = conv_pip(name('conv5'), x)
    x = tf.tanh(x, name=name('control_tanh'))
    z = tf.mul(tf.exp(x), self_ability)
    z_sum = tf.reduce_sum(z, reduction_indices=[1,2,3], name=name('partition_function')) # partition function

    # another formula of y*logy
    loss = -tf.reduce_sum(tf.mul(x, y), reduction_indices=[1,2,3]) + tf.log(z_sum)
    z_sum = tf.reshape(z_sum, [-1, 1, 1, 1])
    pred = tf.div(z, z_sum, name=name('predict'))
    return Model([self_pos, enemy_pos, self_ability, enemy_protect], input_label, loss, pred, debug=z)

if __name__=='__main__':

    model = get_model('test')
    sess = tf.InteractiveSession()
    sess.run(tf.initialize_all_variables())

    x_data = np.random.randint(2, size=[4,100,9,10,16]).astype('float32')
    y_data = np.random.randint(2, size=[100,9,10,16]).astype('float32')

    input_dict = {}
    for var, data in zip(model.inputs, x_data):
        input_dict[var] = data
    input_dict[model.label] = y_data

    loss_val = model.loss.eval(feed_dict=input_dict)

# Path: util/model.py
# class Model(object):
#     def __init__(self, inputs, label, loss, pred, debug=None):
#         self.inputs = inputs
#         self.label = label
#         self.loss = loss
#         self.pred = pred
#         self.debug = debug
# 
# def conv2d(name, x, out_channel, kernel, stride, nl=None):
#     name = functools.partial('{}-{}'.format, name)
#     # W
#     shape = [kernel, kernel, x.get_shape()[3].value, out_channel]
#     W = weight_variable(name('W'), shape)
#     # b
#     b = bias_variable(name('b'), [out_channel])
#     # conv
#     y = tf.nn.conv2d(x, W, strides=[1,stride,stride,1], padding='SAME', name=name('op'))
#     y = y + b
#     # nonlinearty
#     return nl(y) if nl else y
# 
# def fc_layer(name, x, in_shape, out_shape, nl=None):
#     name = functools.partial('{}-{}'.format, name)
# 
#     in_chanel = np.prod(in_shape)
#     out_chanel = np.prod(out_shape)
#     # W
#     W = weight_variable(name('W'), [in_chanel, out_chanel])
#     # b
#     b = bias_variable(name('b'), [out_chanel])
#     # fc
#     x = tf.reshape(x, [-1, in_chanel])
#     x = tf.matmul(x, W) + b
#     out_shape = [-1] + out_shape
#     x = tf.reshape(x, out_shape)
# 
#     return nl(x) if nl else x


# Path: policy_experiment/policy.fc.add-enemyprot/model.py
from config import config
from util.model import Model, conv2d, fc_layer
import tensorflow as tf
import functools
    import numpy as np

    for layer in range(2):
        x_branch = conv_pip(name('conv%d'%layer), x)
        x = tf.concat(3, [x,x_branch], name=name('concate%d'%layer))

    x = conv_pip(name('conv5'), x)
    x = tf.tanh(x, name=name('control_tanh'))
    z = tf.mul(tf.exp(x), self_ability)
    z_sum = tf.reduce_sum(z, reduction_indices=[1,2,3], name=name('partition_function')) # partition function

    # another formula of y*logy
    loss = -tf.reduce_sum(tf.mul(x, y), reduction_indices=[1,2,3]) + tf.log(z_sum)
    z_sum = tf.reshape(z_sum, [-1, 1, 1, 1])
    pred = tf.div(z, z_sum, name=name('predict'))
    return Model([self_pos, enemy_pos, self_ability, enemy_protect], input_label, loss, pred, debug=z)

if __name__=='__main__':

    model = get_model('test')
    sess = tf.InteractiveSession()
    sess.run(tf.initialize_all_variables())

    x_data = np.random.randint(2, size=[4,100,9,10,16]).astype('float32')
    y_data = np.random.randint(2, size=[100,9,10,16]).astype('float32')

    input_dict = {}
    for var, data in zip(model.inputs, x_data):
        input_dict[var] = data
    input_dict[model.label] = y_data

    loss_val = model.loss.eval(feed_dict=input_dict)

    pred_val = model.pred.eval(feed_dict=input_dict)