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

env = gym.make('VizdoomCorridor-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

print(stacked_frames)

array([[0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]])

# Гиперпараметры
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()
print(possible_actions)

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

img_array= []
img_array_test= []

class ReplayBuffer:
  def __init__(self, capacity=10000):
    self.buffer = []
    self.buffer_size = capacity
      
  def add(self, state, action, reward, next_state, done):
    if len(self.buffer) + 1 >= self.buffer_size:
      self.buffer[0:(1+len(self.buffer))-self.buffer_size] = []
        
    self.buffer.append((state, action, reward, next_state, done))

def sample(self):
 buffer_size = len(self.buffer)
 index = np.random.choice(np.arange(buffer_size),
                          size = 32,
                          replace = False)
 return [self.buffer[i] for i in index]
    
def size(self):
 return len(self.buffer)

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

class DQN:
  def __init__(self, state_dim, aciton_dim):
    self.state_dim  = state_dim
    self.action_dim = aciton_dim
    self.epsilon = 1.0

    self.model = self.create_model()

def create_model(self):
  model = tf.keras.Sequential([
    Input(shape=(*self.state_dim,), name="inputs"),
    Conv2D(filters = 32, kernel_size = [8,8], strides = [4,4], padding = "VALID"),
    BatchNormalization(epsilon = 1e-5),
    ReLU(),
    Flatten(),
    Dense(units = 320, activation="relu"),
    Dense(units = self.action_dim, activation = None)
  ])
  model.compile(loss='mse', optimizer=RMSprop(learning_rate=0.03))
  return model

def predict(self, state):
  return self.model.predict(state)

def get_action(self, explore_start, explore_stop, decay_rate, decay_step, state, possible_actions):
  state = np.reshape(state, (1, *state.shape))
  self.epsilon = explore_stop + (explore_start - explore_stop) * np.exp(-decay_rate * decay_step)
  
  q_value = self.predict(state)
  
  if np.random.random() < self.epsilon:
    action = np.random.choice(range(self.action_dim), 1)[0]
    return action

  action = np.max(q_value)
  index = np.where(q_value == action)
  
  return int(index[3])

def train(self, states, targets):
  self.model.fit(states, targets, epochs=10, verbose=0)

def load(self, Model_name):
  self.model = load_model(Model_name + ".h5", compile=True)

def save(self, Model_name):
  self.model.save(Model_name + ".h5")
  
# Не входят в класс DQN  
print(env.action_space.n)

print(env.observation_space.shape)

3
(240, 320, 3)

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

    self.model = DQN(self.state_dim, self.action_dim)

    self.buffer = ReplayBuffer()

def replay(self):
  for _ in range(10):
    # Выбор значений по индексам
    batch = self.buffer.sample()
    states = np.array([each[0] for each in batch], ndmin=3)
    actions = np.array([each[1] for each in batch])
    rewards = np.array([each[2] for each in batch]) 
    next_states = np.array([each[3] for each in batch], ndmin=3)
    done = np.array([each[4] for each in batch])
    
    targets = self.model.predict(states)
    target_next = self.model.predict(next_states)
   

for i in range(0, 32):
  terminal = done[i]
  if terminal:
    targets[i] = rewards[i]
  else:
    targets[i] = rewards[i] + 0.95 * np.max(target_next[i])

self.model.train(states, targets)

def train(self, max_episodes=10, stacked_frames=None):
  if os.path.exists("Models/Name_DQN.h5"):
    self.model.load("Models/Name_DQN")
    print("Модель загружена")
  
  decay_step = 0
  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:
      decay_step += 1
      action = self.model.get_action(explore_start, explore_stop, decay_rate, decay_step, state, possible_actions)
      next_state, reward, done, _ = self.env.step(action)
      img_array.append(next_state)
      next_state, stacked_frames = stack_frames(stacked_frames, next_state, False)
      self.buffer.add(state, action, reward*0.01, next_state, done)
      total_reward += reward
      state = next_state
    if self.buffer.size() >= 32:
      self.replay()
    self.target_update()
    print('Эпизод {}; Награда за эпизод = {}; Эпсилион={} '.format(ep, total_reward, self.model.epsilon))
    
    if ep % 5 == 0:
      self.model.save("Models/Name_DQN")

def test(self, Model_name, stacked_frames = stacked_frames):      
  self.model.load(Model_name)
  decay_step = 0
  for e in range(2):
    done, total_reward = False, 0
    state = self.env.reset()
    state, stacked_frames = stack_frames(stacked_frames, state, True)
    while not done:
      decay_step += 1
      action = self.model.get_action(explore_start, explore_stop, decay_rate, decay_step, state, possible_actions)
      next_state, reward, done, _ = self.env.step(action)
      img_array_test.append(next_state)
      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('VizdoomCorridor-v0')
agent = Agent(env)

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

#Тест
agent.test("Models/Name_DQN", 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_corridor.avi", out_video)