# Path: rlvision/utils.py
# def load_grid_selection(path):
# def process_map_data(path, return_full=False):
# def mask_grid(pos, grid, radius, one_is_free=True):
# def accumulate_map(source_grid, new_grid, one_is_free=True):
# def plot_grid(data, imsize, start=None, pos=None, goal=None, title=None):
# def process_gridworld_data(data_in, imsize):
# def load_mat_data(file_path):
# def find_layout(data, verbose=0):
# def init_h5_db(db_name, save_dir):
# def add_h5_group(group_name, db, db_path=None):
# def add_h5_ds(data, ds_name, db, group_name=None, data_type=np.uint8,
#               renew=False):
# def create_grid_8_dataset(mat_file_path, db_name, save_dir):
# def create_ds_from_splits(mat_file_prefix, db, splits):
# def create_grid_16_dataset(mat_file_prefix, db_name, save_dir, splits=5):
# def create_grid_28_dataset(mat_file_prefix, db_name, save_dir, splits=20):
# def create_grid_40_dataset(mat_file_prefix, db_name, save_dir, splits=100):
# def load_grid8(return_imsize=True):
# def load_grid16(split=None, return_imsize=True):
# def load_grid28(split=None, return_imsize=True):
# def load_grid40(split=None, return_imsize=True):
# 
# Path: rlvision/vin.py
# def vin_model(l_s=16, k=10, l_h=150, l_q=10, l_a=8):
#     _handle_dim_ordering()
# 
#     def ext_start(inputs):
#         m = inputs[0]
#         s = inputs[1]
#         w = K.one_hot(s[:, 0] + l_s * s[:, 1], l_s * l_s)
#         return K.transpose(
#             K.sum(w * K.permute_dimensions(m, (1, 0, 2)), axis=2))
# 
#     map_in = Input(shape=(l_s, l_s, 2)
#                    if data_format == 'channels_last' else (2, l_s, l_s))
#     x = Conv2D(l_h, (3, 3), strides=(1, 1),
#                activation='relu',
#                padding='same')(map_in)
#     r = Conv2D(1, (1, 1), strides=(1, 1),
#                padding='valid',
#                use_bias=False, name='reward')(x)
#     conv3 = Conv2D(l_q, (3, 3), strides=(1, 1),
#                    padding='same',
#                    use_bias=False)
#     conv3b = Conv2D(l_q, (3, 3), strides=(1, 1),
#                     padding='same',
#                     use_bias=False)
#     q_ini = conv3(r)
#     q = q_ini
#     for idx in range(k):
#         v = Lambda(lambda x: K.max(x, axis=CHANNEL_AXIS, keepdims=True),
#                    output_shape=(l_s, l_s, 1)
#                    if data_format == 'channels_last' else (1, l_s, l_s),
#                    name='value{}'.format(idx + 1))(q)
#         q = add([q_ini, conv3b(v)])
# 
#     if data_format == "channels_last":
#         q = Lambda(lambda x: K.permute_dimensions(x, (0, 3, 1, 2)),
#                    output_shape=(l_q, l_s, l_s))(q)
# 
#     q = Reshape(target_shape=(l_q, l_s * l_s))(q)
#     s_in = Input(shape=(2,), dtype='int32')
#     q_out = merge([q, s_in], mode=ext_start, output_shape=(l_q,))
#     out = Dense(l_a, activation='softmax', use_bias=False)(q_out)
#     return Model(inputs=[map_in, s_in], outputs=out)


import os
import cPickle as pickle
import keras.backend as K
import rlvision
from keras.callbacks import ModelCheckpoint
from rlvision import utils
from rlvision.vin import vin_model
"""VIN training.

Grid 28x28

Author: Yuhuang Hu
Email : duguyue100@gmail.com
"""



# load data
file_name = os.path.join(rlvision.RLVISION_DATA,
                         "chain_data", "grid28_po.pkl")
model_path = os.path.join(rlvision.RLVISION_MODEL,
                          "grid28-po")
if not os.path.isdir(model_path):
    os.makedirs(model_path)

# parameters
batch_size = 256
nb_epochs = 80

print('# Minibatch-size: {}'.format(batch_size))
print('# epoch: {}'.format(nb_epochs))
print('')

train, test, _ = utils.process_map_data(file_name)

# Path: rlvision/utils.py
# def load_grid_selection(path):
# def process_map_data(path, return_full=False):
# def mask_grid(pos, grid, radius, one_is_free=True):
# def accumulate_map(source_grid, new_grid, one_is_free=True):
# def plot_grid(data, imsize, start=None, pos=None, goal=None, title=None):
# def process_gridworld_data(data_in, imsize):
# def load_mat_data(file_path):
# def find_layout(data, verbose=0):
# def init_h5_db(db_name, save_dir):
# def add_h5_group(group_name, db, db_path=None):
# def add_h5_ds(data, ds_name, db, group_name=None, data_type=np.uint8,
#               renew=False):
# def create_grid_8_dataset(mat_file_path, db_name, save_dir):
# def create_ds_from_splits(mat_file_prefix, db, splits):
# def create_grid_16_dataset(mat_file_prefix, db_name, save_dir, splits=5):
# def create_grid_28_dataset(mat_file_prefix, db_name, save_dir, splits=20):
# def create_grid_40_dataset(mat_file_prefix, db_name, save_dir, splits=100):
# def load_grid8(return_imsize=True):
# def load_grid16(split=None, return_imsize=True):
# def load_grid28(split=None, return_imsize=True):
# def load_grid40(split=None, return_imsize=True):
# 
# Path: rlvision/vin.py
# def vin_model(l_s=16, k=10, l_h=150, l_q=10, l_a=8):
#     _handle_dim_ordering()
# 
#     def ext_start(inputs):
#         m = inputs[0]
#         s = inputs[1]
#         w = K.one_hot(s[:, 0] + l_s * s[:, 1], l_s * l_s)
#         return K.transpose(
#             K.sum(w * K.permute_dimensions(m, (1, 0, 2)), axis=2))
# 
#     map_in = Input(shape=(l_s, l_s, 2)
#                    if data_format == 'channels_last' else (2, l_s, l_s))
#     x = Conv2D(l_h, (3, 3), strides=(1, 1),
#                activation='relu',
#                padding='same')(map_in)
#     r = Conv2D(1, (1, 1), strides=(1, 1),
#                padding='valid',
#                use_bias=False, name='reward')(x)
#     conv3 = Conv2D(l_q, (3, 3), strides=(1, 1),
#                    padding='same',
#                    use_bias=False)
#     conv3b = Conv2D(l_q, (3, 3), strides=(1, 1),
#                     padding='same',
#                     use_bias=False)
#     q_ini = conv3(r)
#     q = q_ini
#     for idx in range(k):
#         v = Lambda(lambda x: K.max(x, axis=CHANNEL_AXIS, keepdims=True),
#                    output_shape=(l_s, l_s, 1)
#                    if data_format == 'channels_last' else (1, l_s, l_s),
#                    name='value{}'.format(idx + 1))(q)
#         q = add([q_ini, conv3b(v)])
# 
#     if data_format == "channels_last":
#         q = Lambda(lambda x: K.permute_dimensions(x, (0, 3, 1, 2)),
#                    output_shape=(l_q, l_s, l_s))(q)
# 
#     q = Reshape(target_shape=(l_q, l_s * l_s))(q)
#     s_in = Input(shape=(2,), dtype='int32')
#     q_out = merge([q, s_in], mode=ext_start, output_shape=(l_q,))
#     out = Dense(l_a, activation='softmax', use_bias=False)(q_out)
#     return Model(inputs=[map_in, s_in], outputs=out)


# Path: rlvision/exps/vin_exp_po_28.py
import os
import cPickle as pickle
import keras.backend as K
import rlvision
from keras.callbacks import ModelCheckpoint
from rlvision import utils
from rlvision.vin import vin_model

"""VIN training.

Grid 28x28

Author: Yuhuang Hu
Email : duguyue100@gmail.com
"""



# load data
file_name = os.path.join(rlvision.RLVISION_DATA,
                         "chain_data", "grid28_po.pkl")
model_path = os.path.join(rlvision.RLVISION_MODEL,
                          "grid28-po")
if not os.path.isdir(model_path):
    os.makedirs(model_path)

# parameters
batch_size = 256
nb_epochs = 80

print('# Minibatch-size: {}'.format(batch_size))
print('# epoch: {}'.format(nb_epochs))
print('')

train, test, _ = utils.process_map_data(file_name)

model = vin_model(l_s=train[0].shape[2], k=20)