import random
import numpy as np
import matplotlib.pyplot as plt
import collections
import imageio
import IPython.display as display
import mindspore
import mindspore.nn as nn
import mindspore.ops as ops
import mindspore.numpy as mnp
from mindspore import Tensor, context
from mindspore.common.initializer import Normal
from mindspore.train import Model
from maze_generator import MazeGenerator
from maze_environment import MazeEnvironment
from dqn_model import DQN
from replay_buffer import ReplayBuffer
from agent import Agent

def train_dqn(env, agent, net, target_net, buffer, optimizer, batch_size, gamma, epsilon, device):
    state = env.reset(epsilon)
    total_reward = 0
    path = []

    for step in range(500):
        agent.make_a_move(net, epsilon, device)
        total_reward += agent.total_reward
        path.append(tuple(agent.env.current_position))

        if len(buffer) < batch_size:
            continue

        states, actions, next_states, rewards, dones = buffer.sample(batch_size)
        states_v = Tensor(states, mindspore.float32)
        next_states_v = Tensor(next_states, mindspore.float32)
        actions_v = Tensor(actions, mindspore.int32)
        rewards_v = Tensor(rewards, mindspore.float32)
        done_mask = Tensor(dones, mindspore.bool_)

        state_action_values = ops.GatherD()(net(states_v), actions_v.reshape((-1, 1))).squeeze(-1)
        next_state_values = ops.ReduceMax(keep_dims=True)(target_net(next_states_v), 1)
        next_state_values = next_state_values * (1.0 - done_mask.astype(mindspore.float32))
        expected_state_action_values = rewards_v + gamma * next_state_values

        loss = nn.loss.MSELoss()(state_action_values, expected_state_action_values)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if agent.isgameon is False:
            state = env.reset(epsilon)
            break
        else:
            state = next_states

    return total_reward, path

def create_animation(maze, path, init_position, goal, filename='maze_animation.gif'):
    frames = []
    for i, (x, y) in enumerate(path):
        fig, ax = plt.subplots()
        ax.imshow(maze, cmap='Paired', origin='upper')
        ax.text(init_position[1], init_position[0], 'Start', ha='center', va='center', color='black', fontsize=12, weight='bold', bbox=dict(facecolor='yellow', edgecolor='black', boxstyle='round,pad=0.5'))
        ax.text(goal[1], goal[0], 'End', ha='center', va='center', color='white', fontsize=12, weight='bold', bbox=dict(facecolor='blue', edgecolor='black', boxstyle='round,pad=0.5'))
        ax.plot([p[1] for p in path[:i+1]], [p[0] for p in path[:i+1]], marker='o', color='red', markersize=5)
        plt.axis('off')

        fig.canvas.draw()
        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        frames.append(image)
        plt.close(fig)

    imageio.mimsave(filename, frames, fps=2)

def main():
    context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
    mx, my = 20, 20
    maze_gen = MazeGenerator(mx, my)
    maze = maze_gen.get_maze()
    init_position = [0, 0]
    goal = [mx - 1, my - 1]
    env = MazeEnvironment(maze, init_position, goal)
    buffer = ReplayBuffer(10000)
    agent = Agent(env, buffer)

    device = 'cpu'
    net = DQN(env.state().size, 4)
    target_net = DQN(env.state().size, 4)
    target_net.set_train(False)
    optimizer = nn.Adam(net.trainable_params(), learning_rate=1e-4)

    num_episodes = 1000
    batch_size = 64
    gamma = 0.99
    epsilon = 1.0
    epsilon_decay = 0.995
    epsilon_min = 0.01

    best_path = []
    for episode in range(num_episodes):
        epsilon = max(epsilon * epsilon_decay, epsilon_min)
        total_reward, path = train_dqn(env, agent, net, target_net, buffer, optimizer, batch_size, gamma, epsilon, device)
        if episode % 10 == 0:
            target_net.set_train(False)
            target_net.set_train(net.train())
            print(f"Episode {episode}, Epsilon: {epsilon}")
        if total_reward > len(best_path):
            best_path = path

    create_animation(maze, best_path, init_position, goal)
    display.Image("/content/maze_animation.gif")

if __name__ == "__main__":
    main()