import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Conv2D, Flatten, BatchNormalization, ReLU, InputLayer
from tensorflow.keras.optimizers import Adam
import gym
import numpy as np
from collections import deque
import random
import vizdoomgym
from datetime import datetime

env = gym.make('VizdoomDefendCenter-v0')
n_outputs = env.action_space.n
observation = env.reset()
plt.imshow(observation)
plt.show()

def preprocess_frame(frame):
  # Добавляем оттенок серого
  frame = np.mean(frame,-1)

  # Обрезать экран (удалить часть, не содержащую информации)
  # [вверх: вниз, влево: вправо]
  cropped_frame = frame[15:-5,20:-20]
    
  # Нормализация значений пикселей
  normalized_frame = cropped_frame/255.0
    
  # Изменение размера
  preprocessed_frame = transform.resize(cropped_frame, [240,320])
    
  return preprocessed_frame

stack_size = 4 # Складываем 4 кадра

stacked_frames = deque([np.zeros((240,320), dtype=np.int) 
                        for i in range(stack_size)], maxlen=4)

def stack_frames(stacked_frames, state, is_new_episode):
  frame = preprocess_frame(state)
  
  if is_new_episode:
    # Очистить наш stacked_frames
    stacked_frames = deque([np.zeros((240,320), dtype=np.int) 
                            for i in  range(stack_size)], maxlen=4)
    
    stacked_frames.append(frame)
    stacked_frames.append(frame)
    stacked_frames.append(frame)
    stacked_frames.append(frame)

  # Сложить кадры
  stacked_state = np.stack(stacked_frames, axis=2)
    
else:
  
  # Добавить кадр в двухстороннюю очередь, 
  # автоматически удалив самый старый фрейм
  stacked_frames.append(frame)
  
  # Создайте сложенное (stacked) состояние 
  stacked_state = np.stack(stacked_frames, axis=2) 
  
return stacked_state, stacked_frames

# Гиперпараметры
state_size = [240,320,4]
action_size = env.action_space.n
learning_rate =  0.00025 

total_episodes = 50  # Общее количество эпизодов
max_steps = 50       # Максимально возможное количество шагов в эпизоде
batch_size = 64 

explore_start = 1.0  # Вероятность исследования на старте
explore_stop = 0.01  # Минимальная вероятность исследования 
decay_rate = 0.0001  # Скорость затухания для исследований

# Q learning параметры
gamma = 0.95    # Коэффициент дисконтирования

# Количество опытов, сохраненных в памяти при первой инициализации
pretrain_length = batch_size

# Количество опытов, которые может сохранить память
memory_size = 1000000         

possible_actions = np.identity(action_size, dtype=int).tolist()
img_array = []
img_array_test= []
print(possible_actions)

[[1, 0, 0],
 [0, 1, 0],
 [0, 0, 1]]

#Веса модели будут сохраняться в папку Models.
Save_Path = 'Models'
if not os.path.exists(Save_Path): 
  os.makedirs(Save_Path)
path = '{}_PG'.format("Name")
Model_name = os.path.join(Save_Path, path)

class REINFORCEAgent:

  def __init__(self, env):
    self.env = gym.make(env)
    self.state_dim = [240,320,4]
    self.action_dim = self.env.action_space.n

    self.states, self.actions, self.rewards = [], [], []

    self.model = self.create_model(input_shape=self.state_dim, 
                                   action_space = self.action_dim)

def create_model(self, input_shape, action_space):
  X_input = Input(input_shape)
  X = Conv2D(filters = 32, kernel_size = [6,6], strides = [4,4], 
             padding = "VALID", activation="relu")(X_input)
  X = Conv2D(filters = 64, kernel_size = [3,3], strides = [2,2], 
             padding = "VALID", activation="relu")(X)
  X = Flatten(input_shape=input_shape)(X)
  X = Dense(512, activation="elu")(X)
  X = Dense(64, activation="elu")(X)
  action = Dense(action_space, activation="softmax")(X)
  model = tf.keras.Model(inputs = X_input, outputs = action)
  model.compile(loss='categorical_crossentropy', 
               optimizer=Adam(learning_rate=0.005))
  return model

def remember(self, state, action, reward):
  self.states.append(state)
  action_one_hot = np.zeros([self.action_dim])
  action_one_hot[action] = 1
  self.actions.append(action_one_hot)
  self.rewards.append(reward)
  
# Используем выходные данные policy network и 
# выбираем действие стохастически (Стохастическая стратегия)
def get_action(self, state):
  state = np.reshape(state, (1, *state.shape))
  policy = self.model.predict(state).flatten()
  return np.random.choice(self.action_dim, 1, p=policy)[0]

def discount_rewards(self, reward):
  running_add = 0
  discounted_r = np.zeros_like(reward)
    
  for i in reversed(range(0,len(reward))):
    if reward[i] != 0: 
      running_add = 0
    running_add = running_add * gamma + reward[i]
    discounted_r[i] = running_add

  # нормализация результата 
  # discounted_r =
  # делим его на среднеквадратичное отклонение
  # discounted_r =

  return discounted_r

def replay(self):
  states = np.array(self.states, ndmin=3)
  actions = np.array(self.actions)
    
  discounted_r = self.discount_rewards(self.rewards)
    
  self.model.fit(states, actions, sample_weight=discounted_r, 
                 epochs=1, verbose=0)
  self.states, self.actions, self.rewards = [], [], []
  
def load(self, Model_name):
  self.model = load_model(Model_name, compile=True)
    
def save(self, Model_name):
  self.model.save('ModelPG' + '.h5')

def train(self, max_episodes=10, stacked_frames=None):
  if os.path.exists("Models/Name_PG.h5"):
    self.model.load("Models/Name_PG")
    print("Модель загружена")
  
  for ep in range(max_episodes):
    done, total_reward = False, 0
    state = self.env.reset()
    state, stacked_frames = stack_frames(stacked_frames, state, True)
    while not done:
      action = self.get_action(state)
      next_state, reward, done, _ = self.env.step(action)
      img_array.append(next_state)
      self.remember(state, action, reward)
      next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)
      total_reward += reward
      state = next_state
      if done:
        self.replay()
    print('Эпизод {} Награда = {}'.format(ep, total_reward))
    
    if ep % 5 == 0:
      self.model.save("Models/Name_PG")
  self.env.close()

def test(self, Model_name, stacked_frames = stacked_frames):
  self.model.load(Model_name)
  for ep in range(2):
    done, total_reward = False, 0
    state = self.env.reset()
    state, stacked_frames = stack_frames(stacked_frames, state, True)
    while not done:
      action = self.get_action(state)
      next_state, reward, done, _ = self.env.step(action)
      img_array_test.append(next_state)
      self.remember(state, action, reward)
      next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)
      total_reward += reward
      state = next_state
      if done:
        print("Эпизод: {}/{}, общая награда: {}".format(e, 10, total_reward))
        break
  self.env.close()

env = gym.make('VizdoomDefendCenter-v0')
agent = REINFORCEAgent(env)

agent.train(max_episodes=10, stacked_frames = stacked_frames)

#Тест
agent.test("Models/Name_PG", stacked_frames = stacked_frames)

from random import choice
from google.colab.patches import cv2_imshow

import skvideo.io

out_video =  np.empty([len(img_array), 240, 320, 3], 
                      dtype = np.uint8)
out_video =  out_video.astype(np.uint8)

for i in range(len(img_array)):
  frame = img_array[i]
  out_video[i] = frame

skvideo.io.vwrite("doom_defend_center.avi", out_video)