import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.optimizers import Adam
import matplotlib.pyplot as plt
from collections import deque
import random
import os
import sys

# Ajouter le chemin du dossier parent pour importer environment_setup.py
sys.path.append('..')
from environment_setup import ScheduleEnv

# Vérifier la version de TensorFlow
print(f"TensorFlow version: {tf.__version__}")

TensorFlow version: 2.19.0

# Définir les chemins de fichiers
DATA_PATH = "../Data/processed/preprocessed_data.csv"
MODEL_PATH = "../models/dqn_model"
os.makedirs("../models", exist_ok=True)

# 1. Charger les données prétraitées
print("Chargement des données prétraitées...")
df = pd.read_csv(DATA_PATH)
print(f"Données chargées: {df.shape[0]} entrées")

Chargement des données prétraitées...
Données chargées: 37933 entrées

# Configurer l'environnement
env = ScheduleEnv(df)
print(f"Environnement créé avec {env.action_space.n} actions possibles")

Environnement créé avec 18 actions possibles

# 2. Implémentation de la mémoire d'expérience (Experience Replay)
class ReplayBuffer:
    def __init__(self, capacity=10000):
        self.buffer = deque(maxlen=capacity)
    
    def add(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
    
    def sample(self, batch_size):
        samples = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*samples)
        return np.array(states), np.array(actions), np.array(rewards), np.array(next_states), np.array(dones)
    
    def size(self):
        return len(self.buffer)

# 3. Implémentation de l'agent DQN
class DQNAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        
        # Paramètres de l'algorithme
        self.gamma = 0.95  # Facteur d'actualisation
        self.epsilon = 1.0  # Exploration initiale
        self.epsilon_min = 0.01  # Exploration minimale
        self.epsilon_decay = 0.995  # Taux de décroissance de l'exploration
        self.learning_rate = 0.001  # Taux d'apprentissage
        self.batch_size = 64  # Taille des batchs d'apprentissage
        
        # Mémoire d'expérience
        self.memory = ReplayBuffer(capacity=10000)
        
        # Réseaux de neurones (principal et cible)
        self.model = self._build_model()
        self.target_model = self._build_model()
        self.update_target_model()
    
    def _build_model(self):
        # Réseau de neurones pour approximer la fonction Q
        model = Sequential([
            Flatten(input_shape=(self.state_size,)),
            Dense(64, activation='relu'),
            Dense(64, activation='relu'),
            Dense(self.action_size, activation='linear')
        ])
        model.compile(loss='mse', optimizer=Adam(learning_rate=self.learning_rate))
        return model
    
    def update_target_model(self):
        # Copier les poids du modèle principal vers le modèle cible
        self.target_model.set_weights(self.model.get_weights())
    
    def remember(self, state, action, reward, next_state, done):
        # Stocker l'expérience dans la mémoire
        self.memory.add(state, action, reward, next_state, done)
    
    def act(self, state, training=True):
        # Stratégie epsilon-greedy pour l'exploration/exploitation
        if training and np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)
        
        # Prédire les valeurs Q pour toutes les actions
        q_values = self.model.predict(state.reshape(1, -1), verbose=0)
        return np.argmax(q_values[0])
    
    def replay(self):
        # Apprentissage par expérience replay
        if self.memory.size() < self.batch_size:
            return
        
        # Échantillonner un batch de la mémoire
        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
        
        # Calculer les valeurs Q cibles
        targets = self.model.predict(states, verbose=0)
        next_q_values = self.target_model.predict(next_states, verbose=0)
        
        for i in range(self.batch_size):
            if dones[i]:
                targets[i, actions[i]] = rewards[i]
            else:
                targets[i, actions[i]] = rewards[i] + self.gamma * np.max(next_q_values[i])
        
        # Entraîner le modèle
        self.model.fit(states, targets, epochs=1, verbose=0)
        
        # Mettre à jour epsilon pour réduire l'exploration
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay
    
    def load(self, path):
        self.model.load_weights(path)
        self.update_target_model()
    
    def save(self, path):
        self.model.save_weights(path)

# 4. Fonction d'entraînement
def train_dqn(env, agent, episodes=1000, update_target_every=10, max_steps=24):
    """
    Entraîne l'agent DQN sur l'environnement spécifié
    """
    rewards_history = []
    
    for episode in range(episodes):
        # Réinitialiser l'environnement
        state = env.reset()
        total_reward = 0
        
        for step in range(max_steps):
            # Choisir une action
            action = agent.act(state)
            
            # Exécuter l'action
            next_state, reward, done, _ = env.step(action)
            
            # Stocker l'expérience
            agent.remember(state, action, reward, next_state, done)
            
            # Mettre à jour l'état
            state = next_state
            total_reward += reward
            
            # Apprentissage
            agent.replay()
            
            if done:
                break
        
        # Mettre à jour le modèle cible périodiquement
        if episode % update_target_every == 0:
            agent.update_target_model()
        
        rewards_history.append(total_reward)
        
        # Afficher la progression
        if episode % 100 == 0:
            print(f"Episode: {episode}/{episodes}, Reward: {total_reward:.2f}, Epsilon: {agent.epsilon:.4f}")
    
    return rewards_history

# 5. Visualisation des résultats
def plot_rewards(rewards):
    plt.figure(figsize=(10, 6))
    plt.plot(rewards)
    plt.title('Récompenses par épisode')
    plt.xlabel('Épisode')
    plt.ylabel('Récompense totale')
    plt.grid(True)
    plt.savefig('../outputs/dqn_rewards.png')
    plt.show()

# 6. Évaluation du modèle
def evaluate_agent(env, agent, episodes=10):
    """
    Évalue les performances de l'agent entraîné
    """
    print("\n--- Évaluation du modèle ---")
    total_rewards = []
    
    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False
        
        while not done:
            action = agent.act(state, training=False)
            next_state, reward, done, _ = env.step(action)
            state = next_state
            total_reward += reward
        
        total_rewards.append(total_reward)
        print(f"Épisode {episode+1}: Récompense = {total_reward:.2f}")
    
    print(f"\nRécompense moyenne: {np.mean(total_rewards):.2f}")
    return total_rewards

# 7. Exécution de l'entraînement et de l'évaluation
if __name__ == "__main__":
    # Paramètres
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n
    print(f"État: {state_size} dimensions, Actions: {action_size} possibilités")
    
    # Créer l'agent
    agent = DQNAgent(state_size, action_size)
    
    # Entraîner l'agent
    print("\n--- Début de l'entraînement ---")
    rewards = train_dqn(env, agent, episodes=500, update_target_every=5)
    
    # Sauvegarder le modèle
    agent.save(MODEL_PATH)
    print(f"\nModèle sauvegardé dans {MODEL_PATH}")
    
    # Visualiser les résultats
    plot_rewards(rewards)
    
    # Évaluer l'agent
    evaluate_agent(env, agent)

État: 27 dimensions, Actions: 18 possibilités

c:\Users\Abdel\anaconda3\Lib\site-packages\keras\src\layers\reshaping\flatten.py:37: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
  super().__init__(**kwargs)

--- Début de l'entraînement ---
Episode: 0/500, Reward: 13.50, Epsilon: 1.0000

# 8. Visualisation des plannings générés
def visualize_schedule(env, agent):
    """
    Génère et visualise un planning optimisé par l'agent
    """
    state = env.reset()
    done = False
    
    activities = []
    slots = []
    
    for slot in range(24):
        action = agent.act(state, training=False)
        next_state, reward, done, _ = env.step(action)
        
        # Récupérer le nom de l'activité
        activity_name = df[df['ACTIVITY_NAME_ENC'] == action]['ACTIVITY_NAME'].mode()
        if len(activity_name) > 0:
            activity = activity_name.iloc[0]
        else:
            activity = f"Activity {action}"
        
        activities.append(activity)
        slots.append(slot)
        
        state = next_state
        
        if done:
            break
    
    # Créer un DataFrame pour la visualisation
    schedule_df = pd.DataFrame({
        'Heure': slots,
        'Activité': activities
    })
    
    # Visualisation
    plt.figure(figsize=(12, 8))
    plt.barh(schedule_df['Heure'], [1] * len(schedule_df), color='skyblue')
    
    # Ajouter les noms des activités
    for i, (hour, activity) in enumerate(zip(schedule_df['Heure'], schedule_df['Activité'])):
        plt.text(0.5, hour, activity, ha='center', va='center')
    
    plt.yticks(slots, [f"{h}:00" for h in slots])
    plt.xlabel('Activité')
    plt.ylabel('Heure')
    plt.title('Planning journalier optimisé')
    plt.grid(True, axis='y', alpha=0.3)
    plt.savefig('../outputs/optimized_schedule.png')
    plt.show()

# Créer le dossier de sortie si nécessaire
os.makedirs("../outputs", exist_ok=True)

# Générer un exemple de planning
print("\n--- Génération d'un planning optimisé ---")
visualize_schedule(env, agent)

# Afficher les informations sur la taille de la mémoire d'expérience
print(f"\nTaille de la mémoire d'expérience: {agent.memory.size()} échantillons")
print(f"Taux d'exploration final: {agent.epsilon:.4f}")