# 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
    name = functools.partial('{}-{}'.format, name)

    self_pos = tf.placeholder(config.dtype, config.data_shape, name='self_pos')
    enemy_pos = tf.placeholder(config.dtype, config.data_shape, name='enemy_pos')
    self_ability = tf.placeholder(config.dtype, config.data_shape, name='self_ability')
    enemy_ability = tf.placeholder(config.dtype, config.data_shape, name='enemy_ability')
    self_protect = tf.placeholder(config.dtype, config.data_shape, name='self_protect')
    enemy_protect = tf.placeholder(config.dtype, config.data_shape, name='enemy_protect')

    input_label = tf.placeholder(config.dtype, config.label_shape, name='input_label')

    x = tf.concat(3, [self_pos, enemy_pos, self_ability, enemy_ability, self_protect, enemy_protect], name=name('input_concat'))
    y = input_label

    nl = tf.nn.tanh

    def conv_pip(name, x):
        name = functools.partial('{}_{}'.format, name)

        x = conv2d(name('subconv0'), x, config.data_shape[3]*2, kernel=3, stride=1, nl=nl)
        x = conv2d(name('subconv1'), x, config.data_shape[3], kernel=3, stride=1, nl=nl)
        x = fc_layer(name('subfc2'), x, config.unit_shape, config.unit_shape, nl)
        return x

    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'))

# 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-all/model.py
from config import config
from util.model import Model, conv2d, fc_layer
import tensorflow as tf
import functools
    import numpy as np

    name = functools.partial('{}-{}'.format, name)

    self_pos = tf.placeholder(config.dtype, config.data_shape, name='self_pos')
    enemy_pos = tf.placeholder(config.dtype, config.data_shape, name='enemy_pos')
    self_ability = tf.placeholder(config.dtype, config.data_shape, name='self_ability')
    enemy_ability = tf.placeholder(config.dtype, config.data_shape, name='enemy_ability')
    self_protect = tf.placeholder(config.dtype, config.data_shape, name='self_protect')
    enemy_protect = tf.placeholder(config.dtype, config.data_shape, name='enemy_protect')

    input_label = tf.placeholder(config.dtype, config.label_shape, name='input_label')

    x = tf.concat(3, [self_pos, enemy_pos, self_ability, enemy_ability, self_protect, enemy_protect], name=name('input_concat'))
    y = input_label

    nl = tf.nn.tanh

    def conv_pip(name, x):
        name = functools.partial('{}_{}'.format, name)

        x = conv2d(name('subconv0'), x, config.data_shape[3]*2, kernel=3, stride=1, nl=nl)
        x = conv2d(name('subconv1'), x, config.data_shape[3], kernel=3, stride=1, nl=nl)
        x = fc_layer(name('subfc2'), x, config.unit_shape, config.unit_shape, nl)
        return x

    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)