import random
from collections import deque

import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim


class DQN(nn.Module):
    def __init__(self, n_states=6, n_actions=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_states, 64), nn.ReLU(),
            nn.Linear(64, n_actions)
        )

    def forward(self, x):
        """
        Forward pass of the network.
        :param x:  torch.Tensor of shape [n_states]
        :return: torch.Tensor of shape [n_actions] with Q-Values for each action
        """
        return self.net(x)


def train_and_save(weights_path="acrobot_dqn.pth", episodes=2000, update_target_every=20):
    """
    Train a DQN agent on the Acrobot-v1 environnement.
    :param weights_path: file path to save learned network weights
    :param episodes: number of training episodes (complete games)
    :param update_target_every: how many episodes to wait before syncing the target network
    :return: trained Q-Network ready to be used for inference
    """

    # environnement setup
    env = gym.make("Acrobot-v1")
    n_states, n_actions = env.observation_space.shape[0], env.action_space.n

    # les DQN
    policy_net = DQN(n_states, n_actions) # Q Network
    target_net = DQN(n_states, n_actions) # Target network
    target_net.load_state_dict(policy_net.state_dict()) # same weights at start
    target_net.eval()

    # Optimizer et hyperparameters
    optimizer = optim.Adam(policy_net.parameters(), lr=1e-3)
    gamma = 0.99 # discount factor
    epsilon = 1.0 # Fréquence d'exploration initiale
    eps_min = 0.05 # Fréquence d'exploration minimale
    eps_decay = 0.995 # Facteur de réduction d'epsilon
    memory = deque(maxlen=5000)
    batch_size = 64


    # main training loop
    for ep in range(episodes):
        # env.reset() returns a tuple (initial_state, info_dict)
        s, _ = env.reset()
        s = torch.tensor(s, dtype=torch.float32)
        done, total_r = False, 0

        while not done:
            # epsilon-greedy
            if random.random() < epsilon:
                a = random.randrange(n_actions)
            else:
                a = torch.argmax(policy_net(s)).item()

            # env.step(s) returns (next_state, reward, terminated, truncated, info)
            ns, r, done, _, _ = env.step(a)
            ns = torch.tensor(ns, dtype=torch.float32)

            memory.append((s, a ,r ,ns, done))
            s, total_r = ns, total_r + r

            # learning phase
            if len(memory) >= batch_size:
                batch = random.sample(memory, batch_size)
                s_b, a_b, r_b, ns_b, d_b = zip(*batch)
                s_b = torch.stack(s_b)
                ns_b = torch.stack(ns_b)

                # Q values for chosen actions
                q_pred = policy_net(s_b).gather(1, torch.tensor(a_b).unsqueeze(1)).squeeze()

                # Target values using target network
                with torch.no_grad():
                    q_next = target_net(ns_b).max(1)[0]
                    q_target = torch.tensor(r_b, dtype=torch.float32) + \
                        gamma * q_next * (1 - torch.tensor(d_b, dtype=torch.float32))

                # MSE
                loss = ((q_pred - q_target)**2).mean()
                optimizer.zero_grad(); loss.backward(); optimizer.step()

        # decay epsilon to gradually reduce exploration
        epsilon = max(eps_min, epsilon * eps_decay)

        # Periodically synchronise target network with policy network
        if (ep + 1) % update_target_every == 0:
            target_net.load_state_dict(policy_net.state_dict())

        if (ep + 1) % 20 == 0:
            print(f'Episode {ep + 1}: total reward {total_r:.1f}, epsilon {epsilon:.2f}')
    env.close()

    # save trained policy network
    torch.save(policy_net.state_dict(), weights_path)
    print(f'Training finished. Weights saved to {weights_path}')
    return policy_net # <--- trained Q-network


def show(weights_path="acrobot_dqn.pth") -> None:
    """
    Load trained Q network and run a single episode to visually
    demonstrate the learned policy
    :param weights_path: path to the saved network weights
    :return:
    """
    env = gym.make("Acrobot-v1", render_mode="human")
    qnet = DQN()
    qnet.load_state_dict(torch.load(weights_path))
    qnet.eval()

    s, _ = env.reset()
    s = torch.tensor(s, dtype=torch.float32)
    done = False
    while not done:
        a = torch.argmax(qnet(s)).item()
        s_, r, done, _, _ = env.step(a)
        s = torch.tensor(s_, dtype=torch.float32)
    env.close()
    print('Demonstration finished.')


if __name__ == '__main__':
    trained_model = train_and_save()
    show()