agent.py

# ──────────────────────────────────────────────────────────
#  Project  : Snake_AI
#  File     : agent.py
#  Author   : ProGen18
#  Created  : 22-02-2026
#  Modified : 27-03-2026
#  Purpose  : Main training entry point and AI agent logic
# ──────────────────────────────────────────────────────────


# ============================================================
#                          IMPORTS
# ============================================================

import random
import sys
import time
from collections import deque

import numpy as np
import pygame
import torch

from config import CONFIG
from dashboard import Dashboard
from game import JeuVectorise, Point, TAILLE_BLOC

# honesty_check: Ensure random seeds are set immediately after imports
from logger import JournalDeBord
from model import Entraineur, ReseauNeurones


# ============================================================
#                         SETUP & SEEDS
# ============================================================

random.seed(CONFIG.graine)
np.random.seed(CONFIG.graine)
torch.manual_seed(CONFIG.graine)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(CONFIG.graine)


# ============================================================
#                         UTILITIES
# ============================================================


def journal(message: str) -> None:
    """Print a timestamped log message to console."""
    heure = time.strftime("%H:%M:%S", time.localtime())
    print(f"[{heure}] {message}")


# ============================================================
#                        MEMORY SYSTEM
# ============================================================

# ╔══════════════════════════════════════════════════════════╗
# ║                  MemoireEfficace                         ║
# ╠══════════════════════════════════════════════════════════╣
# ║  Handles:                                                ║
# ║   • Pre-allocated NumPy buffers for transition storage   ║
# ║   • Circular buffer logic (FIFO)                         ║
# ║   • Fast batch sampling for training                     ║
# ╚══════════════════════════════════════════════════════════╝


class MemoireEfficace:
    """
    Highly optimized experience replay memory using fixed-size NumPy arrays.
    """

    def __init__(self, capacite: int, taille_etat: int = CONFIG.input_size):
        """
        Initialize the replay memory buffers.

        Args:
            capacite (int): Maximum number of transitions to store.
            taille_etat (int): Dimension of the state vector.
        """
        self.capacite = capacite
        self.etats = np.zeros((capacite, taille_etat), dtype=np.float32)
        self.actions = np.zeros(capacite, dtype=np.int32)
        self.recompenses = np.zeros(capacite, dtype=np.float32)
        self.etats_suivants = np.zeros((capacite, taille_etat), dtype=np.float32)
        self.finis = np.zeros(capacite, dtype=bool)
        self.position = 0
        self.taille = 0

    def stocker_batch(
        self,
        etats: np.ndarray,
        actions: np.ndarray,
        recompenses: np.ndarray,
        etats_suivants: np.ndarray,
        finis: np.ndarray,
    ) -> None:
        """
        Store a batch of transitions into the circular buffer.

        Args:
            etats (np.ndarray): Array of current states.
            actions (np.ndarray): Array of actions taken.
            recompenses (np.ndarray): Array of rewards received.
            etats_suivants (np.ndarray): Array of next states.
            finis (np.ndarray): Array of terminal flags.
        """
        n = len(etats)
        if self.position + n <= self.capacite:
            # Entire batch fits in the remaining space
            self.etats[self.position : self.position + n] = etats
            self.actions[self.position : self.position + n] = actions
            self.recompenses[self.position : self.position + n] = recompenses
            self.etats_suivants[self.position : self.position + n] = etats_suivants
            self.finis[self.position : self.position + n] = finis
        else:
            # Batch wraps around the buffer
            part1 = self.capacite - self.position
            part2 = n - part1

            # End of buffer
            self.etats[self.position : self.capacite] = etats[:part1]
            self.actions[self.position : self.capacite] = actions[:part1]
            self.recompenses[self.position : self.capacite] = recompenses[:part1]
            self.etats_suivants[self.position : self.capacite] = etats_suivants[:part1]
            self.finis[self.position : self.capacite] = finis[:part1]

            # Start of buffer
            self.etats[0:part2] = etats[part1:]
            self.actions[0:part2] = actions[part1:]
            self.recompenses[0:part2] = recompenses[part1:]
            self.etats_suivants[0:part2] = etats_suivants[part1:]
            self.finis[0:part2] = finis[part1:]

        self.position = (self.position + n) % self.capacite
        self.taille = min(self.taille + n, self.capacite)

    def echantillonner(self, batch_size: int) -> tuple:
        """Randomly sample a batch of transitions."""
        indices = np.random.randint(0, self.taille, size=batch_size)
        return (
            self.etats[indices],
            self.actions[indices],
            self.recompenses[indices],
            self.etats_suivants[indices],
            self.finis[indices],
        )

    def __len__(self) -> int:
        return self.taille


# ============================================================
#                        VISUALIZATION
# ============================================================


class RenduPygame:
    """Handles rendering of a single game environment for the dashboard."""

    def __init__(self, env: JeuVectorise, index_env: int = 0):
        """
        Initialize the renderer.

        Args:
            env (JeuVectorise): The game engine.
            index_env (int): Which environment to visualize.
        """
        self.env = env
        self.idx = index_env
        self.largeur = env.l
        self.hauteur = env.h
        self.surface = pygame.Surface((self.largeur, self.hauteur))

    def dessiner(self) -> pygame.Surface:
        """
        Render the current frame to the internal surface.

        Returns:
            pygame.Surface: The rendered game frame.
        """
        self.surface.fill((0, 0, 0))

        # --- 1. Draw Snake ---
        points_serpent = self.serpent
        nb_points = len(points_serpent)
        for i, pt in enumerate(points_serpent):
            # Gradient effect: lighter towards the head
            ratio = 1 - (i / nb_points)
            luminosite = max(0.3, ratio)
            c = (int(50 * luminosite), int(200 * luminosite), int(50 * luminosite))

            pygame.draw.rect(self.surface, c, (pt.x, pt.y, TAILLE_BLOC, TAILLE_BLOC))
            pygame.draw.rect(
                self.surface, (0, 50, 0), (pt.x, pt.y, TAILLE_BLOC, TAILLE_BLOC), 1
            )

        # --- 2. Draw Food ---
        pomme = self.pomme
        pygame.draw.rect(
            self.surface, (255, 0, 0), (pomme.x, pomme.y, TAILLE_BLOC, TAILLE_BLOC)
        )

        return self.surface

    @property
    def serpent(self) -> list[Point]:
        """Convert vectorized body coordinates to list of Points."""
        longueur = self.env.longueurs[self.idx]
        corps = self.env.corps[self.idx, :longueur]
        return [Point(x * TAILLE_BLOC, y * TAILLE_BLOC) for x, y in corps]

    @property
    def tetes(self) -> Point:
        """Get head position in pixels."""
        hx, hy = self.env.tetes[self.idx]
        return Point(hx * TAILLE_BLOC, hy * TAILLE_BLOC)

    @property
    def pomme(self) -> Point:
        """Get food position in pixels."""
        fx, fy = self.env.pommes[self.idx]
        return Point(fx * TAILLE_BLOC, fy * TAILLE_BLOC)


# ============================================================
#                          AI AGENT
# ============================================================

# ╔══════════════════════════════════════════════════════════╗
# ║                      AgentIA                             ║
# ╠══════════════════════════════════════════════════════════╣
# ║  Handles:                                                ║
# ║   • Neural network management and optimization           ║
# ║   • Epsilon-greedy exploration policy                    ║
# ║   • Training loops (on-policy and experience replay)     ║
# ║   • Performance metrics and logging                      ║
# ╚══════════════════════════════════════════════════════════╝


class AgentIA:
    """
    High-level agent managing DQN training and inference.
    """

    def __init__(self):
        """Initialize models, optimizer, and memory buffers."""
        self.nb_parties = 0
        self.nb_frames = 0
        self.epsilon = CONFIG.epsilon_depart
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        journal(f"Device: {self.device}")

        # --- AI Core ---
        self.modele = ReseauNeurones(CONFIG.input_size, CONFIG.output_size).to(
            self.device
        )
        self.entraineur = Entraineur(
            self.modele,
            lr=CONFIG.taux_apprentissage,
            gamma=CONFIG.gamma,
            device=self.device,
            tau=CONFIG.tau,
        )

        # --- Metrics ---
        self.logger = JournalDeBord()
        self.record = 0
        self.scores_historique = deque(maxlen=CONFIG.scores_historique_maxlen)
        self.scores_test = deque(maxlen=CONFIG.scores_test_maxlen)
        self.debut_entrainement = time.time()

        # --- Memory ---
        self.memoire_pommes = MemoireEfficace(CONFIG.memoire_pommes_capacite)

        # --- Eval ---
        self.derniere_eval = 0
        self.eval_intervalle = CONFIG.eval_intervalle

    def evaluer_test_set(self, env_test: JeuVectorise) -> tuple[float, float]:
        """
        Evaluate agent performance on a clean test environment (greedy).

        Args:
            env_test (JeuVectorise): Isolated test environment.

        Returns:
            tuple: (Mean Score, Std Dev).
        """
        self.modele.eval()
        scores = []

        for _ in range(CONFIG.nb_episodes_test):
            env_test.reset()
            done = False
            while not done:
                etat = env_test.recuperer_etats()
                etat_tensor = torch.tensor(etat, dtype=torch.float).to(self.device)

                with torch.no_grad():
                    action = torch.argmax(self.modele(etat_tensor)).item()

                _, _, done_arr, score_arr = env_test.step(np.array([action]))
                done = done_arr[0]
                if done:
                    scores.append(score_arr[0])

        self.modele.train()
        return float(np.mean(scores)), float(np.std(scores))

    def epsilon_schedule(self) -> float:
        """Linear decay scheduler for exploration rate."""
        decay = self.nb_frames / CONFIG.epsilon_frames
        val = CONFIG.epsilon_depart - decay * (
            CONFIG.epsilon_depart - CONFIG.epsilon_fin
        )
        return max(CONFIG.epsilon_fin, val)

    def convertir_etat_tensor(self, etats_numpy: np.ndarray) -> torch.Tensor:
        """Convert NumPy state matrix to PyTorch tensor on device."""
        return torch.tensor(etats_numpy, dtype=torch.float).to(self.device)

    def entrainer_on_policy(
        self,
        etats: np.ndarray,
        actions: np.ndarray,
        recompenses: np.ndarray,
        etats_suivants: np.ndarray,
        finis: np.ndarray,
    ) -> None:
        """
        Perform a gradient step using fresh steps plus sampled apple-eating memories.

        Args:
            etats (np.ndarray): Fresh states.
            actions (np.ndarray): Fresh actions.
            recompenses (np.ndarray): Fresh rewards.
            etats_suivants (np.ndarray): Fresh next states.
            finis (np.ndarray): Fresh terminal flags.
        """
        n = len(etats)
        n_sample = min(n, CONFIG.taille_batch_pqn)
        indices = np.random.choice(n, n_sample, replace=False)

        # --- Hybrid Sampling ---
        if len(self.memoire_pommes) >= CONFIG.memoire_pommes_seuil:
            e2, a2, r2, n2, d2 = self.memoire_pommes.echantillonner(
                CONFIG.memoire_pommes_echantillons
            )
            E = np.concatenate([etats[indices], e2])
            A = np.concatenate([actions[indices], a2])
            R = np.concatenate([recompenses[indices], r2])
            N = np.concatenate([etats_suivants[indices], n2])
            D = np.concatenate([finis[indices], d2])
        else:
            E, A, R, N, D = (
                etats[indices],
                actions[indices],
                recompenses[indices],
                etats_suivants[indices],
                finis[indices],
            )

        self.entraineur.etape_d_apprentissage(E, A, R, N, D)

    def moyenne_mobile(self, n: int = 100) -> float:
        """Calculate moving average of last N scores."""
        if not self.scores_historique:
            return 0.0
        recent = list(self.scores_historique)[-n:]
        return sum(recent) / len(recent)


# ============================================================
#                       TRAINING LOOP
# ============================================================


def lancer_entrainement() -> None:
    """Main execution thread for training the Snake AI."""

    # --- Initialization ---
    env = JeuVectorise(n_envs=CONFIG.nb_environnements)
    env_test = JeuVectorise(n_envs=1)
    agent = AgentIA()
    dashboard = Dashboard()
    visu = RenduPygame(env, index_env=0)

    t0 = time.time()
    frames = 0
    last_screen_time = time.time()
    derniere_eval_test = 0

    etats = env.recuperer_etats()

    journal(f"Démarrage avec {CONFIG.nb_environnements} environnements parallèles")
    journal(
        f"Epsilon schedule: {CONFIG.epsilon_depart} → {CONFIG.epsilon_fin} sur {CONFIG.epsilon_frames} frames"
    )
    journal(f"Test set: évaluation toutes les {CONFIG.eval_intervalle} parties")

    # --- Evolution Loop ---
    while True:
        # --- 1. Event Handling ---
        evenements = pygame.event.get()
        action_user = None
        for event in evenements:
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            act = dashboard.handle_input(event)
            if act:
                action_user = act

        # Dashboard states (PAUSED, etc.)
        if dashboard.state != "RUNNING":
            dashboard.update()
            continue

        # Process user commands
        if action_user:
            if action_user == "QUIT":
                pygame.quit()
                sys.exit()
            elif action_user == "EXPORT":
                agent.logger.exporter_excel()
            elif isinstance(action_user, tuple):
                cmd, fichier = action_user
                if cmd == "SAVE":
                    temps_jeu = time.time() - agent.debut_entrainement
                    agent.modele.sauvegarder(
                        nom_fichier=fichier,
                        nb_parties=agent.nb_parties,
                        temps_total=temps_jeu,
                        etat_optimiseur=agent.entraineur.optimiseur.state_dict(),
                        epsilon=agent.epsilon,
                        record=agent.record,
                    )
                    journal(f"Sauvegardé: {fichier}")
                elif cmd == "LOAD":
                    res = agent.modele.charger(fichier, agent.device)
                    if res is not None:
                        nb, t, opt, eps, rec = res
                        agent.nb_parties = nb
                        agent.debut_entrainement = time.time() - t
                        agent.record = rec
                        if eps is not None:
                            agent.epsilon = eps
                        if opt is not None:
                            try:
                                agent.entraineur.optimiseur.load_state_dict(opt)
                            except Exception as e:
                                journal(f"Optimiseur non chargé: {e}")
                        agent.entraineur.target_model.load_state_dict(
                            agent.modele.state_dict()
                        )
                        journal(f"Chargé: {fichier}")

        # --- 2. Action Selection ---
        agent.epsilon = agent.epsilon_schedule()
        etat_tensor = agent.convertir_etat_tensor(etats)

        agent.modele.eval()
        with torch.no_grad():
            prediction = agent.modele(etat_tensor)
        agent.modele.train()

        # Policy: Mixture of Model, Epsilon-Random, and Heuristic-Random
        masque_exploration = np.random.random(CONFIG.nb_environnements) < agent.epsilon
        actions_modele = torch.argmax(prediction, dim=1).cpu().numpy()

        p_heuristic = max(0.0, 1.0 - agent.nb_frames / CONFIG.blend_frames)
        if p_heuristic > 0:
            masque_h = np.random.random(CONFIG.nb_environnements) < p_heuristic
            actions_h = env.actions_gloutonnes()
            actions_s = env.actions_aleatoires_sures()
            actions_aleatoires = np.where(masque_h, actions_h, actions_s)
        else:
            actions_aleatoires = env.actions_aleatoires_sures()

        coups_finaux = np.where(masque_exploration, actions_aleatoires, actions_modele)

        # --- 3. Step Environment ---
        etats_suivants, recompenses, finis, scores = env.step(coups_finaux)

        # Sparse memory for important events (eating)
        masque_positif = recompenses > 0.5
        if np.any(masque_positif):
            agent.memoire_pommes.stocker_batch(
                etats[masque_positif],
                coups_finaux[masque_positif],
                recompenses[masque_positif],
                etats_suivants[masque_positif],
                finis[masque_positif],
            )

        # Optimization step
        agent.entrainer_on_policy(
            etats, coups_finaux, recompenses, etats_suivants, finis
        )
        etats = etats_suivants
        agent.nb_frames += 1

        # --- 4. Stat Tracking ---
        nb_morts = np.sum(finis)
        if nb_morts > 0:
            scores_morts = scores[finis]
            for i_m, s in enumerate(scores_morts):
                agent.scores_historique.append(s)
                num = agent.nb_parties - nb_morts + i_m + 1
                dashboard.update_courbe(
                    float(s), agent.moyenne_mobile(100), num, float(agent.record)
                )
            agent.nb_parties += nb_morts

        max_actuel = np.max(scores)
        if max_actuel > agent.record:
            agent.record = max_actuel
            journal(f"🏆 Nouveau Record: {agent.record}")
            agent.modele.sauvegarder(
                nb_parties=agent.nb_parties,
                temps_total=time.time() - agent.debut_entrainement,
                etat_optimiseur=agent.entraineur.optimiseur.state_dict(),
                epsilon=agent.epsilon,
                record=agent.record,
            )

        frames += 1

        # --- 5. Logging & Dashboard ---
        if time.time() - t0 > CONFIG.log_intervalle_sec:
            tps = frames * CONFIG.nb_environnements
            moyenne = agent.moyenne_mobile(100)

            journal(
                f"{tps} TPS | Parties: {agent.nb_parties} | Eps: {agent.epsilon:.3f} | Moy100: {moyenne:.1f} | Record: {agent.record}"
            )
            agent.logger.noter_stats(
                agent.nb_parties, agent.epsilon, agent.record, moyenne, tps
            )
            dashboard.update_stats(
                agent.nb_parties,
                time.time() - agent.debut_entrainement,
                agent.epsilon,
                agent.record,
                tps,
            )

            if agent.epsilon < CONFIG.lr_scheduler_epsilon_seuil:
                agent.entraineur.scheduler.step(moyenne)

            if agent.nb_parties - derniere_eval_test >= agent.eval_intervalle:
                derniere_eval_test = agent.nb_parties
                m_t, s_t = agent.evaluer_test_set(env_test)
                journal(f"📊 TEST SET: {m_t:.1f} ± {s_t:.1f} (100 parties greedy)")

            frames = 0
            t0 = time.time()

        # Screenshot logic
        if dashboard.auto_screen_active:
            if time.time() - last_screen_time >= dashboard.screen_interval:
                dashboard._take_screenshot()
                last_screen_time = time.time()

        # Render cycle
        if dashboard.state == "RUNNING":
            surface_jeu = visu.dessiner()
            dashboard.update_game(surface_jeu)
            dashboard.update_nn(etats[0])

        dashboard.update()


# ============================================================
#                            MAIN
# ============================================================

if __name__ == "__main__":
    lancer_entrainement()