from ddqn import DDQN
from keras.models import Sequential
from keras.layers import Dense, Convolution2D, Flatten, ZeroPadding2D
from keras.optimizers import RMSprop
from keras import backend as K
K.set_image_dim_ordering('tf')

# Use any environment here. OpenAI's Gym can be used
# Must ensure that model fits the dimensions of data returned by environment
env = some_arbitrary_environment.load()

cols, rows = env.shape()
frames = 4

shape = (cols, rows, frames) # specific to arbitrary environment
num_actions = 7              # specific to arbitrary environment

def create_model():
    model = Sequential()

    # Layer 1
    model.add(ZeroPadding2D(input_shape=shape))
    model.add(Convolution2D(16, 8, strides=(4,4), activation='relu'))

    # Layer 2
    model.add(ZeroPadding2D())
    model.add(Convolution2D(32, 4, strides=(2,2), activation='relu'))

    # Layer 3
    model.add(Flatten())
    model.add(Dense(units=256, activation='relu'))

    # Output Layer
    model.add(Dense(units=num_actions, activation='linear'))

    # optimizer and loss
    learning_rate = 0.00025
    loss = "mse"
    rmsprop = RMSprop(lr=learning_rate, clipvalue=1)

    model.compile(loss=loss, optimizer=rmsprop)

    return rmsprop, loss, learning_rate, model

optimizer, loss, learning_rate, model = create_model()

ddqn = DDQN(env, model, loss, optimizer, learning_rate, num_actions, 500, shape)

while True:
    show = env.display()
    ddqn.act_and_learn()
    # Environment automatically resets
    # Check ddqn.py for functions which can print statistics here