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, 'random.choice', 'random.choice', (['choices'], {}), False, 'import random\n'), (18, 'lib.env_config.EnvConfig', 'env_config.EnvConfig', (['"""config/debug.json"""'], {}), False, 'from lib import env_config\n'), (19, 'lib.senior_env.BetterEnvironment', 'BetterEnvironment', (['config_'], {}), False, 'from lib.senior_env import BetterEnvironment\n'), (24, 'lib.models.build_paper_model', 'models.build_paper_model', ([], {}), False, 'from lib import models\n'), (35, 'rl.memory.SequentialMemory', 'SequentialMemory', ([], {'limit': '(10000000)', 'window_length': 'WINDOW_LENGTH'}), False, 'from rl.memory import SequentialMemory\n'), (51, 'rl.agents.dqn.DQNAgent', 'DQNAgent', ([], {'model': 'model', 'nb_actions': 'nb_actions', 'policy': 'policy', 'memory': 'memory', 'nb_steps_warmup': '(50000)', 'gamma': '(0.99)', 'target_model_update': '(10000)', 'train_interval': '(4)', 'delta_clip': '(1.0)'}), False, 'from rl.agents.dqn import DQNAgent\n'), (42, 'rl.policy.EpsGreedyQPolicy', 'EpsGreedyQPolicy', ([], {}), False, 'from rl.policy import LinearAnnealedPolicy, BoltzmannQPolicy, EpsGreedyQPolicy\n'), (54, 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.00025)'}), False, 'from keras.optimizers import Adam\n')]