# Path: baselines/deepq/prediction/tfacvp/tf_ops.py
# def ReLu(x, name, reuse=False):
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         l = tf.nn.relu(x)
#     return l
# 
# def Conv2D(x, filter_shape, out_dim, strides, padding, name, reuse=False):
#     # x: input tensor (float32)[n, w, h, c]
#     # filter_shape: conv2d filter (int)[w, h]
#     # out_dim: output channels (int)
#     # strides: conv2d stride size (int)
#     # padding: padding type (str)
#     # name: variable scope (str)
#            
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=filter_shape + [in_dim, out_dim], initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=True))
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.nn.conv2d(x, w, strides=[1, strides, strides, 1], padding=padding, name='conv2d')
#         l = tf.nn.bias_add(l, b, name='bias_add')
#     return l
# 
# def FC(x, out_dim, name, initializer=tf.contrib.layers.xavier_initializer(uniform=True), reuse=False):
#     # x: input tensor (float32)[n, in_dim]
#     # out_dim: output channels (int)
#     # name: variable scope (str)
# 
#     x = tf.contrib.layers.flatten(x)
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=[in_dim, out_dim], initializer=initializer)
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.add(tf.matmul(x, w), b, name='add')
#     return l
# 
# def Deconv2D(x, filter_shape, output_shape, out_dim, strides, padding, name, reuse=False):
#     # x: input tensor (float32) [n, w, h, c]
#     # filter_shape: conv2d filter (int)[w, h]
#     # out_dim: output channels (int)
#     # strides: conv2d stride size (int)
#     # padding: padding type (str)
#     # name: variable scope (str)
# 
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=filter_shape + [out_dim, in_dim], initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=True))
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=[1, strides, strides, 1], padding=padding, name='deconv2d')
#         l = tf.nn.bias_add(l, b, name='bias_add')
#     return l


import tensorflow as tf
import numpy as np
import re
from .tf_ops import ReLu, Conv2D, FC, Deconv2D
        # x: input image (tensor([batch_size, 84, 84, 12]))
        l = Conv2D(x, [6, 6], 64, 2, 'VALID', 'conv1')
        l = ReLu(l, 'relu1')
        l = Conv2D(l, [6, 6], 64, 2, 'SAME', 'conv2')
        l = ReLu(l, 'relu2')
        l = Conv2D(l, [6, 6], 64, 2, 'SAME', 'conv3')
        l = ReLu(l, 'relu3')
        l = FC(l, 1024, 'ip1')
        l = ReLu(l, 'relu4')
        l = FC(l, 2048, 'enc-factor', initializer=tf.random_uniform_initializer(minval=-1.0, maxval=1.0))
        return l

    def _create_action_embedding(self, act):
        # act: action input (tensor([batch_size, num_act])) (one-hot vector)
        act = tf.cast(act, tf.float32)
        l = FC(act, 2048, 'act-embed', initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1))
        return l

    def _create_decoder(self, encode, act_embed):
        # encode: encode layer
        # act_embed: action embedding layer
        batch_size = tf.shape(encode)[0]
        l = tf.multiply(encode, act_embed, name='merge')
        l = FC(l, 1024, 'dec')
        l = FC(l, 64 * 10 * 10, 'ip4')
        l = ReLu(l, 'relu1')
        l = tf.reshape(l, [-1, 10, 10, 64], name='dec-reshape')
        l = Deconv2D(l, [6, 6], [batch_size, 20, 20, 64], 64, 2, 'SAME', 'deconv3')
        l = ReLu(l, 'relu2')
        l = Deconv2D(l, [6, 6], [batch_size, 40, 40, 64], 64, 2, 'SAME', 'deconv2')

# Path: baselines/deepq/prediction/tfacvp/tf_ops.py
# def ReLu(x, name, reuse=False):
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         l = tf.nn.relu(x)
#     return l
# 
# def Conv2D(x, filter_shape, out_dim, strides, padding, name, reuse=False):
#     # x: input tensor (float32)[n, w, h, c]
#     # filter_shape: conv2d filter (int)[w, h]
#     # out_dim: output channels (int)
#     # strides: conv2d stride size (int)
#     # padding: padding type (str)
#     # name: variable scope (str)
#            
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=filter_shape + [in_dim, out_dim], initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=True))
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.nn.conv2d(x, w, strides=[1, strides, strides, 1], padding=padding, name='conv2d')
#         l = tf.nn.bias_add(l, b, name='bias_add')
#     return l
# 
# def FC(x, out_dim, name, initializer=tf.contrib.layers.xavier_initializer(uniform=True), reuse=False):
#     # x: input tensor (float32)[n, in_dim]
#     # out_dim: output channels (int)
#     # name: variable scope (str)
# 
#     x = tf.contrib.layers.flatten(x)
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=[in_dim, out_dim], initializer=initializer)
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.add(tf.matmul(x, w), b, name='add')
#     return l
# 
# def Deconv2D(x, filter_shape, output_shape, out_dim, strides, padding, name, reuse=False):
#     # x: input tensor (float32) [n, w, h, c]
#     # filter_shape: conv2d filter (int)[w, h]
#     # out_dim: output channels (int)
#     # strides: conv2d stride size (int)
#     # padding: padding type (str)
#     # name: variable scope (str)
# 
#     with tf.variable_scope(name, reuse=reuse) as scope:
#         in_dim = x.get_shape()[-1]
#         w = tf.get_variable('w', shape=filter_shape + [out_dim, in_dim], initializer=tf.contrib.layers.xavier_initializer_conv2d(uniform=True))
#         b = tf.get_variable('b', shape=[out_dim], initializer=tf.constant_initializer(0.0))
#         l = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=[1, strides, strides, 1], padding=padding, name='deconv2d')
#         l = tf.nn.bias_add(l, b, name='bias_add')
#     return l


# Path: baselines/deepq/prediction/tfacvp/old_model.py
import tensorflow as tf
import numpy as np
import re
from .tf_ops import ReLu, Conv2D, FC, Deconv2D

        # x: input image (tensor([batch_size, 84, 84, 12]))
        l = Conv2D(x, [6, 6], 64, 2, 'VALID', 'conv1')
        l = ReLu(l, 'relu1')
        l = Conv2D(l, [6, 6], 64, 2, 'SAME', 'conv2')
        l = ReLu(l, 'relu2')
        l = Conv2D(l, [6, 6], 64, 2, 'SAME', 'conv3')
        l = ReLu(l, 'relu3')
        l = FC(l, 1024, 'ip1')
        l = ReLu(l, 'relu4')
        l = FC(l, 2048, 'enc-factor', initializer=tf.random_uniform_initializer(minval=-1.0, maxval=1.0))
        return l

    def _create_action_embedding(self, act):
        # act: action input (tensor([batch_size, num_act])) (one-hot vector)
        act = tf.cast(act, tf.float32)
        l = FC(act, 2048, 'act-embed', initializer=tf.random_uniform_initializer(minval=-0.1, maxval=0.1))
        return l

    def _create_decoder(self, encode, act_embed):
        # encode: encode layer
        # act_embed: action embedding layer
        batch_size = tf.shape(encode)[0]
        l = tf.multiply(encode, act_embed, name='merge')
        l = FC(l, 1024, 'dec')
        l = FC(l, 64 * 10 * 10, 'ip4')
        l = ReLu(l, 'relu1')
        l = tf.reshape(l, [-1, 10, 10, 64], name='dec-reshape')
        l = Deconv2D(l, [6, 6], [batch_size, 20, 20, 64], 64, 2, 'SAME', 'deconv3')
        l = ReLu(l, 'relu2')
        l = Deconv2D(l, [6, 6], [batch_size, 40, 40, 64], 64, 2, 'SAME', 'deconv2')

        l = ReLu(l, 'relu3')