"""
Implementing a Q-learning agent from scratch
"""
import logging

import gym
import numpy as np

#######################################################################################################################
# Create Log
#######################################################################################################################

# Config log file
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Create Handlers
c_handler = logging.StreamHandler()
f_handler = logging.FileHandler(filename='./Logs/MountainCar.log')

# Create Log Format(s)
f_format = logging.Formatter('%(asctime)s:%(processName)s:%(name)s:%(levelname)s:%(message)s')

# Set handler levels
c_handler.setLevel(logging.INFO)
f_handler.setLevel(logging.INFO)

# Set handler formats
c_handler.setFormatter(f_format)
f_handler.setFormatter(f_format)

# Add handlers to the logger
logger.addHandler(c_handler)
logger.addHandler(f_handler)


#######################################################################################################################
# Create Environment and set global constants
#######################################################################################################################

MAX_NUM_EPISODES = 50000
STEPS_PER_EPISODE = 200
EPSILON_MIN = 0.005
max_num_steps = MAX_NUM_EPISODES * STEPS_PER_EPISODE
EPSILON_DECAY = 500 * EPSILON_MIN / max_num_steps
ALPHA = 0.05  # Learning rate
GAMMA = 0.98  # Discount Factor
NUM_DISCRETE_BINS = 30


#######################################################################################################################
# Classes
#######################################################################################################################

class QLearner(object):
    def __init__(self, env):
        self.obs_shape = env.observation_space.shape
        self.obs_high = env.observation_space.high
        self.obs_low = env.observation_space.low
        self.obs_bins = NUM_DISCRETE_BINS
        self.bin_width = (self.obs_high - self.obs_low) / self.obs_bins
        self.action_shape = env.action_space.n
        self.Q = np.zeros((self.obs_bins + 1, self.obs_bins + 1, self.action_shape))
        self.alpha = ALPHA
        self.gamma = GAMMA
        self.epsilon = 1.0

    def discretize(self, obs):
        return tuple(((obs - self.obs_low) / self.bin_width).astype(int))

    def get_action(self, obs):
        discretized_obs = self.discretize(obs)
        if self.epsilon > EPSILON_MIN:
            self.epsilon -= EPSILON_DECAY
        if np.random.random() > self.epsilon:
            return np.argmax(self.Q[discretized_obs])
        else:  # Choose a random action
            return np.random.choice([a for a in range(self.action_shape)])

    def learn(self, obs, action, reward, next_obs):
        discretized_obs = self.discretize(obs)
        discretized_next_obs = self.discretize(next_obs)
        td_target = reward + self.gamma * np.max(self.Q[discretized_next_obs])
        td_error = td_target - self.Q[discretized_obs][action]
        self.Q[discretized_obs][action] += self.alpha * td_error


def train(agent, env):
    """
    Trains an agent given an environment
    :param agent: the agent to train
    :param env: the environment to train in
    :return: Trained agent policy
    """
    best_reward = -float('inf')
    for episode in range(MAX_NUM_EPISODES):
        done = False
        obs = env.reset()
        total_reward = 0.0
        while not done:
            action = agent.get_action(obs)
            next_obs, reward, done, info = env.step(action)  # Send the action to the environment
            agent.learn(obs, action, reward, next_obs)
            obs = next_obs
            total_reward += reward
        if total_reward > best_reward:
            best_reward = total_reward
        logger.info(f'Episode #: {episode} reward: {total_reward} best_reward: {best_reward} eps: {agent.epsilon}')
    return np.argmax(agent.Q, axis=2)


def test(agent, env, policy):
    """
    Tests an agent given an environment and the agent's policy
    :param agent: the agent to test
    :param env: the environment to test in
    :param policy: the policy to follow
    :return: agent's total reward as float
    """
    done = False
    obs = env.reset()
    total_reward = 0.0
    while not done:
        action = policy[agent.discretize(obs)]
        next_obs, reward, done, info = env.step(action)
        obs = next_obs
        total_reward += reward
    return total_reward


#######################################################################################################################
# Simple environment
#######################################################################################################################

if __name__ == '__main__':
    env = gym.make('MountainCar-v0')

    agent = QLearner(env)
    learned_policy = train(agent, env)

    # Monitor agent in evironment
    gym_monitor_path = './gym_monitor_output'
    # env = gym.wrappers.Monitor(env, gym_monitor_path, force=True)
    for _ in range(1000):
        test(agent, env, learned_policy)
    env.close()
