import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import random
from collections import deque
import matplotlib.pyplot as plt
import pandas as pd
import os
from tqdm.notebook import tqdm
import pickle

from environment_setup import ScheduleEnv, df, NUM_ACTIVITIES, NUM_TIME_SLOTS

# 2. Charger ou redéfinir le modèle DQN
class DQNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(DQNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, output_dim)

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

# 3. Configurer les paramètres d'entraînement
state_size = NUM_TIME_SLOTS  # 24
action_size = NUM_ACTIVITIES
gamma = 0.95            # Facteur de décompte
epsilon = 1.0           # Exploration vs exploitation
epsilon_decay = 0.995   # Décroissance d'epsilon
epsilon_min = 0.1       # Epsilon minimum
learning_rate = 0.001   # Taux d'apprentissage
batch_size = 64         # Taille du batch pour l'apprentissage
memory_size = 10000     # Taille de la mémoire de replay
num_episodes = 1000     # Nombre d'épisodes d'entraînement

# 4. Initialisation du modèle et de l'environnement
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Utilisation de l'appareil: {device}")

model = DQNetwork(state_size, action_size).to(device)
target_model = DQNetwork(state_size, action_size).to(device)
target_model.load_state_dict(model.state_dict())
target_model.eval()

optimizer = optim.Adam(model.parameters(), lr=learning_rate)
loss_fn = nn.MSELoss()
memory = deque(maxlen=memory_size)

Utilisation de l'appareil: cpu

env = ScheduleEnv(df)

# 5. Fonction pour choisir une action (epsilon-greedy)
def choose_action(state, epsilon):
    if np.random.rand() <= epsilon:
        return random.randrange(action_size)
    state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
    with torch.no_grad():
        q_values = model(state_tensor)
    return torch.argmax(q_values).item()

# 5. Fonction pour choisir une action (epsilon-greedy)
def choose_action(state, epsilon):
    if np.random.rand() <= epsilon:
        return random.randrange(action_size)
    state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
    with torch.no_grad():
        q_values = model(state_tensor)
    return torch.argmax(q_values).item()

# 7. Fonction pour mettre à jour le modèle cible
def update_target_model():
    target_model.load_state_dict(model.state_dict())

# 8. Entraînement du modèle
episode_rewards = []
losses = []
epsilons = []

update_target_every = 10