import numpy as np
import torch.nn as nn
import torch
from torch.optim import Adam
from collections import deque
import random
import matplotlib.pyplot as plt

class MazeEnv:
    def __init__(self):
        self.exit_coord = (3, 3)
        self.row_n, self.col_n = 4, 4
        self.walls = {(0, 3), (1, 0), (1, 2), (2, 2), (3, 0)}
        self.state = (0, 0)  # исходное состояние
        self.actions = [0, 1, 2, 3]  # вверх, вниз, влево, вправо
        self.reset()

    def reset(self):
        self.state = (0, 0)  # Сбросить положение агента до начальной точки
        return self.state

    def step(self, action):
        row, col = self.state
        if action == 0:  # начальство
            next_state = (max(row - 1, 0), col)
        elif action == 1:  # Вниз
            next_state = (min(row + 1, self.row_n - 1), col)
        elif action == 2:  # левый
            next_state = (row, max(col - 1, 0))
        elif action == 3:  # верно
            next_state = (row, min(col + 1, self.col_n - 1))
        else:
            raise ValueError("Invalid action")

        reward = -1  # Награда по умолчанию
        done = False
        if next_state in self.walls:
            reward = -10
        elif next_state == self.exit_coord:
            reward = 10
            done = True

        self.state = next_state if next_state not in self.walls else self.state
        return self.state, reward, done

    def get_state_size(self):
        return self.row_n * self.col_n

    def get_action_size(self):
        return len(self.actions)

# Модель DQN
class DQN(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(DQN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        output = self.fc3(x)
        return output

# Опыт повтора
class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        state, action, reward, next_state, done = zip(*random.sample(self.buffer, batch_size))
        return np.array(state), action, reward, np.array(next_state), done

    def __len__(self):
        return len(self.buffer)

def update_model(dqn_model, target_dqn_model, replay_buffer, optimizer, batch_size, gamma):
    # от Опыт Случайная выборка серии экспериментов из буфера повтора.
    state, action, reward, next_state, done = replay_buffer.sample(batch_size)

    # Преобразовать список в тензор
    state = torch.FloatTensor(state)
    action = torch.LongTensor(action)
    reward = torch.FloatTensor(reward)
    next_state = torch.FloatTensor(next_state)
    done = torch.FloatTensor(done)

    # вычислить当前состояние Вниз,принимать практические мерыдействиеизQценить
    current_q_values = dqn_model(state).gather(1, action.unsqueeze(1)).squeeze(1)

    # Используйте целевую модель DQNвычислить Вниз一состояниеиз最大Qценить
    next_q_values = target_dqn_model(next_state).max(1)[0]
    # Рассчитать ожидаемое значение Q
    expected_q_values = reward + gamma * next_q_values * (1 - done)

    # Рассчитать потери
    loss = torch.nn.functional.mse_loss(current_q_values, expected_q_values.detach())

    # Оптимизатор обнуления градиента
    optimizer.zero_grad()
    # Обратное распространение ошибки
    loss.backward()
    # Обновить вес модели
    optimizer.step()

# Обучение DQN
def train_dqn(epochs=300):
    global epsilon
    for epoch in range(epochs):
        state = env.reset()
        state = np.eye(env.row_n * env.col_n)[state[0] * env.col_n + state[1]]
        total_reward = 0
        done = False

        while not done:
            if random.random() < epsilon:
                action = random.choice(env.actions)
            else:
                q_values = dqn_model(torch.FloatTensor(state).unsqueeze(0))
                action = torch.argmax(q_values).item()

            next_state, reward, done = env.step(action)
            next_state = np.eye(env.row_n * env.col_n)[next_state[0] * env.col_n + next_state[1]]
            replay_buffer.push(state, action, reward, next_state, done)
            state = next_state
            total_reward += reward

            if len(replay_buffer) > batch_size:
                update_model(dqn_model, target_dqn_model, replay_buffer, optimizer, batch_size, gamma)

        epsilon = max(epsilon_min, epsilon_decay * epsilon)
        print(f'Epoch: {epoch}, Total Reward: {total_reward}')
        reward_list.append(total_reward)

        if epoch % 10 == 0:
            target_dqn_model.load_state_dict(dqn_model.state_dict())

if __name__ == "__main__":
    # Инициализировать среду
    env = MazeEnv()

    # Инициализировать DQN
    state_size = env.get_state_size()
    action_size = env.get_action_size()

    dqn_model = DQN(input_dim=state_size, output_dim=action_size)
    target_dqn_model = DQN(input_dim=state_size, output_dim=action_size)
    target_dqn_model.load_state_dict(dqn_model.state_dict())

    optimizer = Adam(dqn_model.parameters(), lr=1e-4)
    replay_buffer = ReplayBuffer(10000)
    batch_size = 64
    gamma = 0.99
    epsilon = 1.0
    epsilon_min = 0.01
    epsilon_decay = 0.995

    reward_list = []

    train_dqn()

    x = range(len(reward_list))
    plt.plot(x, reward_list, label='Line', color='blue')
    plt.xlabel("step")
    plt.ylabel("reward")
    plt.title("DQN to solve the 4*4 maze problem")
    plt.show()

def simulate_optimal_path(model, env):
    state = env.reset()
    optimal_path = [state]  # Запишите статус оптимального пути
    done = False
    while not done:
        state = np.eye(env.row_n * env.col_n)[state[0] * env.col_n + state[1]]
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            action = model(state_tensor).max(1)[1].view(1, 1)
        next_state, _, done = env.step(action.item())
        optimal_path.append(next_state)
        state = next_state
    return optimal_path

optimal_path = simulate_optimal_path(dqn_model, env)
print(optimal_path)

[(0, 0), (0, 1), (1, 1), (2, 1), (3, 1), (3, 2), (3, 3)]