import gym
from gym import error, spaces, utils
from gym.utils import seeding
import itertools
import random
import time
class ShopsEnv(gym.Env):
  metadata = {'render.modes': ['human']}
  # конструктор класса, в котором происходит
  # инициализация среды
  def __init__(self):
    self.state = [0, 0, 0]  # текущее состояние
    self.next_state = [0, 0, 0]  # следующее состояние
    self.done = False  # флажок завершения эпизода
    self.actions = list(itertools.permutations([1, 2, 3]))  # массив возможных действий агента
    self.reward = 0  # текущая награда за действие
    self.time_tracker = 0  # трекер дня эпизода
    self.remembered_states = []  # очередь из трех последних состояний
    # для стохастичности среды
    t = int( time.time() * 1000.0 )
    random.seed( ((t & 0xff000000) >> 24) +
                 ((t & 0x00ff0000) >>  8) +
                 ((t & 0x0000ff00) <<  8) +
                 ((t & 0x000000ff) << 24)   )
  # метод позволяет агенту выполнить одно действие (шаг) в среде
  def step(self, action_num):
    # проверяем не завершен ли уже эпизод
    if self.done:
        return [self.state, self.reward, self.done, self.next_state]
    else:
        # выбираем следующее состояние текущим
        self.state = self.next_state
        # запоминаем состояние
        self.remembered_states.append(self.state)
        # инкрементируем трекер
        self.time_tracker += 1
        # выбираем действие в соответствии с полученным индексом
        action = self.actions[action_num]
        # обновляем состояние, используя выбранное действие (добавляем пироги)
        self.next_state = [x + y for x, y in zip(action, self.state)]
        # генерируем сколько будет куплено
        self.next_state[0] -= (3 + random.uniform(-0.1, 0.1))
        self.next_state[1] -= (1 + random.uniform(-0.1, 0.1))
        self.next_state[2] -= (2 + random.uniform(-0.1, 0.1))
        # вычисляем награду за действие
        if any([x < 0 for x in self.next_state]):
            self.reward = sum([x for x in self.next_state if x < 0])
        else:
            self.reward = 1
        # если накопилась очередь из минимум трех состояний
        # значит нужно убрать просроченные продукты
        # при этом если ушли в минус (не хватило пирогов для покупателей),
        # то также убираем данные отрицательные значения
        if self.time_tracker >= 3:
            remembered_state = self.remembered_states.pop(0)
            self.next_state = [max(x - y, 0) for x, y in zip(self.next_state, remembered_state)]
        else:
            self.next_state = [max(x, 0) for x in self.next_state]
        # проверяем прошло ли уже 30 дней
        self.done = self.time_tracker == 30
        # возвращаем результат шага агента в среде
        return [self.state, self.reward, self.done, self.next_state]
  # метод перезагрузки среды
  def reset(self):
    # устанавливаем все параметры в изначальное положение
    self.state = [0, 0, 0]
    self.next_state = [0, 0, 0]
    self.done = False
    self.reward = 0
    self.time_tracker = 0
    self.remembered_states = []
    t = int( time.time() * 1000.0 )
    random.seed( ((t & 0xff000000) >> 24) +
                 ((t & 0x00ff0000) >>  8) +
                 ((t & 0x0000ff00) <<  8) +
                 ((t & 0x000000ff) << 24)   )
    # возвращаем изначальное состояние
    return self.state
  # метод рендера текущего состояния среды:
  # сколько и в каком магазине пирогов
  def render(self, mode='human', close=False):
    print('-'*20)
    print('First shop')
    print('Pies:', self.state[0])
    print('Second shop')
    print('Pies:', self.state[1])
    print('Third shop')
    print('Pies:', self.state[2])
    print('-'*20)
    print('')

import numpy as np # линейная алгебра
import pandas as pd # препроцессинг данных
import gym # для сред
import gym_shops # для своей кастомной среды
from tqdm import tqdm # для прогресс бара
# для графиков
import matplotlib.pyplot as plt
import seaborn as sns
from IPython.display import clear_output
sns.set_color_codes()
# для моделирования
from collections import deque
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam
import random # для стохастичности среды

class DQLAgent():
    def __init__(self, env):
        # определяем параметры и гиперпараметры
        self.state_size = 3 # размер входа нейронной сети
        self.action_size = 6 # размер выхода нейронной сети
        # эта часть для replay()
        self.gamma = 0.99
        self.learning_rate = 0.01
        # эта часть для adaptiveEGreedy()
        self.epsilon = 0.99
        self.epsilon_decay = 0.99
        self.epsilon_min = 0.0001
        self.memory = deque(maxlen = 5000) # дек с 5000 ячейками памяти, если он переполнится - удалятся первые ячейки
        # собираем модель (NN)
        self.model = self.build_model()
    # метод сборки нейронной сети для Deep Q Learning
    def build_model(self):
        model = Sequential()
        model.add(Dense(10, input_dim = self.state_size, activation = 'sigmoid')) # первый скрытый слой
        model.add(Dense(50, activation = 'sigmoid')) # второй слой
        model.add(Dense(10, activation = 'sigmoid')) # третий слой
        model.add(Dense(self.action_size, activation = 'sigmoid')) # выходной слой
        model.compile(loss = 'mse', optimizer = Adam(lr = self.learning_rate))
        return model
    # метод для запоминания состояния
    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
    # метод выбора действия
    def act(self, state):
        # если случайное число от 0 до 1 меньше epsilon
        # то выбираем действие случайно (exploration)
        if random.uniform(0,1) <= self.epsilon:
            return random.choice(range(6))
        else:
            # иначе нейронная сеть предсказывает следующее действие на основе текущего состояния
            act_values = self.model.predict(state)
            return np.argmax(act_values[0])
    # метод для тренировки нейронной сети
    def replay(self, batch_size):
        # выходим из метода, если еще на накопили достаточно опыта в памяти
        if len(self.memory) < batch_size:
            return
        minibatch = random.sample(self.memory, batch_size) # берем batch_size примеров рандомно из памяти
        # обучаемся на каждой записи батча
        for state, action, reward, next_state, done in minibatch:
            if done: # если эпизод закончен - тогда у нас есть только награда
                target = reward
            else:
                # иначе таргет формируем с помощью следующего состояния
                target = reward + self.gamma * np.amax(self.model.predict(next_state)[0])
                # target = R(s,a) + gamma * max Q`(s`,a`)
                # target (max Q` value) это выход из нейронной сети, которая принимает s` на вход
            train_target = self.model.predict(state) # s --> NN --> Q(s,a) = train_target
            train_target[0][action] = target
            self.model.fit(state, train_target, verbose = 0)
    # метод для уменьшения exploration rate,
    # то есть epsilon
    def adaptiveEGreedy(self):
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

# инициализация gym среды и агента
env = gym.make('shops-v0')
agent = DQLAgent(env)
# устанавливаем параметры тренировки
batch_size = 100
episodes = 1000
# начинаем тренировку
progress_bar = tqdm(range(episodes), position=0, leave=True)
for e in progress_bar:
    # инициализируем среду
    state = env.reset()
    state = np.reshape(state, [1, 3])
    # запоминаем текущий день симуляции, id выбранных действий и сумму наград за эпизод
    time = 0
    taken_actions = []
    sum_rewards = 0
    # симулируем эпизод среды
    while True:
        # выбираем действие
        action = agent.act(state)
        # запоминаем действие
        taken_actions.append(action)
        # выполняем шаг агентом в среде
        next_state, reward, done, _ = env.step(action)
        next_state = np.reshape(next_state, [1, 3])
        # добавляем полученную награду к остальным
        sum_rewards += reward
        # запоминаем результат шага
        agent.remember(state, action, reward, next_state, done)
        # переходим к следующему состоянию
        state = next_state
        # выполняем replay
        agent.replay(batch_size)
        # обновляем epsilon
        agent.adaptiveEGreedy()
        # инкрементируем счетчик времени
        time += 1
        # выводим прогресс тренировки
        progress_bar.set_postfix_str(s='mean reward: {}, time: {}, epsilon: {}'.format(round(sum_rewards/time, 3), time, round(agent.epsilon, 3)), refresh=True)
        # проверяем не завершился ли эпизод
        if done:
            # выводим распределение выбранных действий в течении эпизода
            clear_output(wait=True)
            sns.distplot(taken_actions, color="y")
            plt.title('Episode: ' + str(e))
            plt.xlabel('Action number')
            plt.ylabel('Occurrence in %')
            plt.show()
            break

import time
trained_model = agent  # теперь мы имеем натренированного агента
state = env.reset()  # перезапускаем среду
state = np.reshape(state, [1,3])
# следим за основными параметрами в течении тестового эпизода
time_t = 0
MAX_EPISOD_LENGTH = 1000  # для прогресс бара
taken_actions = []
mean_reward = 0
# симулируем тестовый эпизод
progress_bar = tqdm(range(MAX_EPISOD_LENGTH), position=0, leave=True)
for time_t in progress_bar:
    # выполняем шаг агентом в среде
    action = trained_model.act(state)
    next_state, reward, done, _ = env.step(action)
    next_state = np.reshape(next_state, [1,3])
    state = next_state
    taken_actions.append(action)
    # выводим результат шага
    clear_output(wait=True)
    env.render()
    progress_bar.set_postfix_str(s='time: {}'.format(time_t), refresh=True)
    print('Reward:', round(env.reward, 3))
    time.sleep(0.5)
    mean_reward += env.reward
    if done:
        break
# выводим распределение выбранных действий
sns.distplot(taken_actions, color='y')
plt.title('Test episode - mean reward: ' + str(round(mean_reward/(time_t+1), 3)))
plt.xlabel('Action number')
plt.ylabel('Occurrence in %')
plt.show()