#!/usr/bin/env python3

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
import tensorflow as tf
import tensorflow.keras as keras
tf.get_logger().setLevel('ERROR')

from rl.games import tromis
from rl.agents import dqn
from rl.memory import uniqmemory
from rl.callbacks import history

width, height = 7, 10
nb_frames = 1

game = tromis.Tromis(width, height, max_turn=512)

inp = keras.layers.Input(shape=(nb_frames, height, width, 3))
x = keras.layers.Conv3D(32,5,padding='same',strides=1,activation='relu')(inp)
x = keras.layers.AveragePooling3D(padding='same')(x)
x = keras.layers.Conv3D(64,3,padding='same',strides=1,activation='relu')(x)
x = keras.layers.GlobalAveragePooling3D()(x)
x = keras.layers.Dense(128, activation='relu')(x)
act = keras.layers.Dense(game.nb_actions, activation='linear')(x)

model = keras.models.Model(inputs=inp, outputs=act)
model.compile(keras.optimizers.RMSprop(), keras.losses.LogCosh())
model.summary()

params = {
    'batch_size': 256,
    'epochs': 200,
    'episodes': 100,
    'train_freq': 32,
    'target_sync': 512,
    'epsilon_start': 0.5,
    'epsilon_decay': 0.75,
    'epsilon_final': 0.0,
    'gamma': 0.92,
    'reset_memory': False,
    'observe': 100
}

rlparams = {
    'rl.memory': 'UniqMemory',
    'rl.memory_size': 65536,
    'rl.optimizer': 'RMSprop',
    'rl.with_target': True,
    'rl.nb_frames': nb_frames
}

gameparams = {
    'game.width': game.width,
    'game.height': game.height,
    'game.max_turn': game.max_turn
}

memory = uniqmemory.UniqMemory(memory_size=rlparams['rl.memory_size'])
agent = dqn.Agent(model, memory, with_target=rlparams['rl.with_target'])
#history = history.HistoryLog("tromis", {**params, **rlparams, **gameparams})

agent.train(game, verbose=1, callbacks=[], **params)

tensorflow.keras.models.Model tensorflow.keras.layers.GlobalAveragePooling3D tensorflow.keras.layers.Dense tensorflow.keras.optimizers.RMSprop tensorflow.keras.layers.Conv3D tensorflow.keras.losses.LogCosh tensorflow.get_logger tensorflow.keras.layers.AveragePooling3D tensorflow.keras.layers.Input
[(18, 'rl.games.tromis.Tromis', 'tromis.Tromis', (['width', 'height'], {'max_turn': '(512)'}), False, 'from rl.games import tromis\n'), (20, 'tensorflow.keras.layers.Input', 'keras.layers.Input', ([], {'shape': '(nb_frames, height, width, 3)'}), True, 'import tensorflow.keras as keras\n'), (28, 'tensorflow.keras.models.Model', 'keras.models.Model', ([], {'inputs': 'inp', 'outputs': 'act'}), True, 'import tensorflow.keras as keras\n'), (60, 'rl.memory.uniqmemory.UniqMemory', 'uniqmemory.UniqMemory', ([], {'memory_size': "rlparams['rl.memory_size']"}), False, 'from rl.memory import uniqmemory\n'), (61, 'rl.agents.dqn.Agent', 'dqn.Agent', (['model', 'memory'], {'with_target': "rlparams['rl.with_target']"}), False, 'from rl.agents import dqn\n'), (21, 'tensorflow.keras.layers.Conv3D', 'keras.layers.Conv3D', (['(32)', '(5)'], {'padding': '"""same"""', 'strides': '(1)', 'activation': '"""relu"""'}), True, 'import tensorflow.keras as keras\n'), (22, 'tensorflow.keras.layers.AveragePooling3D', 'keras.layers.AveragePooling3D', ([], {'padding': '"""same"""'}), True, 'import tensorflow.keras as keras\n'), (23, 'tensorflow.keras.layers.Conv3D', 'keras.layers.Conv3D', (['(64)', '(3)'], {'padding': '"""same"""', 'strides': '(1)', 'activation': '"""relu"""'}), True, 'import tensorflow.keras as keras\n'), (24, 'tensorflow.keras.layers.GlobalAveragePooling3D', 'keras.layers.GlobalAveragePooling3D', ([], {}), True, 'import tensorflow.keras as keras\n'), (25, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['(128)'], {'activation': '"""relu"""'}), True, 'import tensorflow.keras as keras\n'), (26, 'tensorflow.keras.layers.Dense', 'keras.layers.Dense', (['game.nb_actions'], {'activation': '"""linear"""'}), True, 'import tensorflow.keras as keras\n'), (29, 'tensorflow.keras.optimizers.RMSprop', 'keras.optimizers.RMSprop', ([], {}), True, 'import tensorflow.keras as keras\n'), (29, 'tensorflow.keras.losses.LogCosh', 'keras.losses.LogCosh', ([], {}), True, 'import tensorflow.keras as keras\n'), (8, 'tensorflow.get_logger', 'tf.get_logger', ([], {}), True, 'import tensorflow as tf\n')]