from __future__ import division

from lib import env_config
from lib.senior_env import BetterEnvironment
from keras.optimizers import Adam
from rl.agents.dqn import DQNAgent
from rl.policy import LinearAnnealedPolicy, BoltzmannQPolicy, EpsGreedyQPolicy
from rl.memory import SequentialMemory
from lib import models
import random

choices = [0,1,2]

def gen_action():
    return random.choice(choices)

if __name__ == '__main__':
    config_ = env_config.EnvConfig('config/debug.json')
    env = BetterEnvironment(config_)

    INPUT_SHAPE = (30, 180)
    WINDOW_LENGTH = 4

    model = models.build_paper_model()

    # Get the environment and extract the number of actions.
    nb_actions = 3

    # Next, we build our model. We use the same model that was described by Mnih et al. (2015).
    input_shape = (WINDOW_LENGTH,) + INPUT_SHAPE


    # Finally, we configure and compile our agent. You can use every built-in Keras optimizer and
    # even the metrics!
    memory = SequentialMemory(limit=10000000, window_length=WINDOW_LENGTH)

    # Select a policy. We use eps-greedy action selection, which means that a random action is selected
    # with probability eps. We anneal eps from 1.0 to 0.1 over the course of 1M steps. This is done so that
    # the agent initially explores the environment (high eps) and then gradually sticks to what it knows
    # (low eps). We also set a dedicated eps value that is used during testing. Note that we set it to 0.05
    # so that the agent still performs some random actions. This ensures that the agent cannot get stuck.
    policy = LinearAnnealedPolicy(EpsGreedyQPolicy(), attr='eps', value_max=1., value_min=.1, value_test=.05,
                                  nb_steps=1000000)

    # The trade-off between exploration and exploitation is difficult and an on-going research topic.
    # If you want, you can experiment with the parameters or use a different policy. Another popular one
    # is Boltzmann-style exploration:
    # policy = BoltzmannQPolicy(tau=1.)
    # Feel free to give it a try!

    dqn = DQNAgent(model=model, nb_actions=nb_actions, policy=policy, memory=memory,
                nb_steps_warmup=50000, gamma=.99, target_model_update=10000,
                train_interval=4, delta_clip=1.)
    dqn.compile(Adam(lr=.00025), metrics=['mae'])


        # Okay, now it's time to learn something! We capture the interrupt exception so that training
    # can be prematurely aborted. Notice that now you can use the built-in Keras callbacks!
    weights_filename = 'dqn_{}_weights.h5f'.format('god_help_me.weights')
    dqn.fit(env, nb_steps=100000, log_interval=10000)

    print(env.portfolio.print_portfolio_results())

[((15, 11, 15, 33), 'random.choice', 'random.choice', ({(15, 25, 15, 32): 'choices'}, {}), '(choices)', False, 'import random\n'), ((18, 14, 18, 55), 'lib.env_config.EnvConfig', 'env_config.EnvConfig', ({(18, 35, 18, 54): '"""config/debug.json"""'}, {}), "('config/debug.json')", False, 'from lib import env_config\n'), ((19, 10, 19, 36), 'lib.senior_env.BetterEnvironment', 'BetterEnvironment', ({(19, 28, 19, 35): 'config_'}, {}), '(config_)', False, 'from lib.senior_env import BetterEnvironment\n'), ((24, 12, 24, 38), 'lib.models.build_paper_model', 'models.build_paper_model', ({}, {}), '()', False, 'from lib import models\n'), ((35, 13, 35, 74), 'rl.memory.SequentialMemory', 'SequentialMemory', (), '', False, 'from rl.memory import SequentialMemory\n'), ((51, 10, 53, 48), 'rl.agents.dqn.DQNAgent', 'DQNAgent', (), '', False, 'from rl.agents.dqn import DQNAgent\n'), ((42, 34, 42, 52), 'rl.policy.EpsGreedyQPolicy', 'EpsGreedyQPolicy', ({}, {}), '()', False, 'from rl.policy import LinearAnnealedPolicy, BoltzmannQPolicy, EpsGreedyQPolicy\n'), ((54, 16, 54, 31), 'keras.optimizers.Adam', 'Adam', (), '', False, 'from keras.optimizers import Adam\n')]