# 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


from config import Config
from util.model import Model, conv2d
import tensorflow as tf
import functools
    import numpy as np
        x = conv2d(name('0'), x, Config.data_shape[3]*2, kernel=3, stride=1, nl=nl)
        x = conv2d(name('1'), x, Config.data_shape[3], kernel=3, stride=1, nl=nl)
        return x

    for layer in range(5):
        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_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):

# 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


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

        x = conv2d(name('0'), x, Config.data_shape[3]*2, kernel=3, stride=1, nl=nl)
        x = conv2d(name('1'), x, Config.data_shape[3], kernel=3, stride=1, nl=nl)
        return x

    for layer in range(5):
        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_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