gpu_optimizer.py

"""
Optimiseur batch GPU — Simulated Annealing avec évaluation vectorisée.

Stratégie
---------
À chaque itération, on génère N_batch voisins de l'allocation courante,
on les évalue tous en un seul appel GPU, puis on applique le critère SA
sur le meilleur du batch.

Cela remplace la boucle séquentielle "1 voisin → évaluation → acceptation"
par "N_batch voisins → N_batch évaluations en parallèle → acceptation du meilleur".

Gain : ×N_batch évaluations par unité de temps (pour N_batch ≤ ~2000 sur RTX 4000 Ada).
"""

from __future__ import annotations

import math
import random
import time
from dataclasses import dataclass, field
from typing import List

import numpy as np
import torch

from gpu_simulator import GPUBatchSimulator

# Types véhicules (doivent correspondre à data_loader)
POMPE_TYPE, VSAV_TYPE, PSE_TYPE = 0, 1, 2


@dataclass
class OptimResult:
    best_alloc:  torch.Tensor     # [V] int64 — CS index par véhicule
    best_score:  float
    history:     List[dict] = field(default_factory=list)
    elapsed_s:   float = 0.0


class GPUBatchSA:
    """
    Recuit simulé avec évaluation batch GPU.

    Parameters
    ----------
    simulator          : GPUBatchSimulator initialisé
    initial_alloc      : [V] int64 — allocation initiale (cs_idx par véhicule)
    N_cs               : nombre total de CS disponibles
    n_batch            : nombre de voisins évalués à chaque itération
    vehicle_type       : [V] int32 — type de chaque véhicule (pour contraintes)
    vsav_pse_pair_idx  : [V] int32 — pour chaque VSAV, index de son PSE pair (-1 si non couplé)
    pse_vsav_pair_idx  : [V] int32 — pour chaque PSE, index de son VSAV pair (-1 si non couplé)
    """

    def __init__(
        self,
        simulator:         GPUBatchSimulator,
        initial_alloc:     np.ndarray | torch.Tensor,
        N_cs:              int,
        n_batch:           int = 512,
        vehicle_type:      np.ndarray | None = None,
        vsav_pse_pair_idx: np.ndarray | None = None,
        pse_vsav_pair_idx: np.ndarray | None = None,
    ):
        self.sim     = simulator
        self.N_cs    = N_cs
        self.n_batch = n_batch
        self.device  = simulator.device
        self.V       = simulator.V

        if isinstance(initial_alloc, np.ndarray):
            initial_alloc = torch.from_numpy(initial_alloc.astype(np.int64))
        self.initial_alloc = initial_alloc.to(self.device)

        # --- Types véhicules (numpy CPU pour la génération de voisins) ---
        if vehicle_type is not None:
            self.vtype_np = vehicle_type.astype(np.int64)
        else:
            self.vtype_np = np.zeros(self.V, dtype=np.int64)

        # --- Indices des véhicules éligibles (non-POMPE) pour mutation ---
        self.movable_mask = self.vtype_np != POMPE_TYPE          # [V] bool
        self.movable_idx  = np.where(self.movable_mask)[0]       # indices des VSAV/PSE

        # --- Couplage PSE↔VSAV (numpy CPU) ---
        self.vsav_pse = vsav_pse_pair_idx if vsav_pse_pair_idx is not None \
            else np.full(self.V, -1, dtype=np.int32)
        self.pse_vsav = pse_vsav_pair_idx if pse_vsav_pair_idx is not None \
            else np.full(self.V, -1, dtype=np.int32)

        # pair_of[v] = index du pair de v (-1 si non couplé)
        self.pair_of = np.full(self.V, -1, dtype=np.int32)
        for v in range(self.V):
            if self.vtype_np[v] == VSAV_TYPE and self.vsav_pse[v] >= 0:
                self.pair_of[v] = self.vsav_pse[v]
            elif self.vtype_np[v] == PSE_TYPE and self.pse_vsav[v] >= 0:
                self.pair_of[v] = self.pse_vsav[v]

        # VSAV non couplés (indices)
        self.solo_vsav_idx = np.array([
            v for v in range(self.V)
            if self.vtype_np[v] == VSAV_TYPE and self.vsav_pse[v] < 0
        ], dtype=np.int64)

        # --- Poids de sélection : VSAV solo 3× plus probable que couplé ---
        solo_vsav_set = set(self.solo_vsav_idx)
        weights = np.ones(len(self.movable_idx), dtype=np.float64)
        for k, v in enumerate(self.movable_idx):
            if v in solo_vsav_set:
                weights[k] = 3.0
        self.movable_weights = weights / weights.sum()

        n_pompe   = int((self.vtype_np == POMPE_TYPE).sum())
        n_coupled = int((self.pair_of >= 0).sum()) // 2
        n_solo_v  = len(self.solo_vsav_idx)
        print(f"[SA] Contraintes: {n_pompe} POMPE gelés, "
              f"{n_coupled} couples PSE-VSAV, {n_solo_v} VSAV solo (poids ×3)")

    # ------------------------------------------------------------------
    def _generate_neighbors(
        self,
        current: torch.Tensor,   # [V] int64
        N:       int,
        n_moves: int,
    ) -> torch.Tensor:
        """
        Génère N voisins avec les contraintes :
          1. Les POMPE ne bougent jamais.
          2. Un couple PSE-VSAV se déplace toujours ensemble.
          3. Si le CS cible contient un VSAV solo, celui-ci fait le swap inverse.

        Returns
        -------
        [N, V] int64
        """
        cur_np = current.cpu().numpy()                           # [V] int64
        neighbors_np = np.tile(cur_np, (N, 1))                  # [N, V]

        movable = self.movable_idx
        if len(movable) == 0:
            return current.unsqueeze(0).expand(N, -1).clone()

        for _ in range(n_moves):
            # Choisir un véhicule éligible (VSAV solo en priorité, poids ×3)
            chosen = movable[np.random.choice(len(movable), size=N, p=self.movable_weights)]
            # CS cible aléatoire par voisin
            target_cs = np.random.randint(0, self.N_cs, size=N)

            for i in range(N):
                v = chosen[i]
                tgt = target_cs[i]
                src = neighbors_np[i, v]

                if tgt == src:
                    continue

                pair_v = self.pair_of[v]

                if pair_v >= 0:
                    # --- Véhicule couplé : déplacer le couple entier ---
                    # Chercher s'il y a un VSAV solo dans le CS cible
                    solo_in_tgt = self.solo_vsav_idx[
                        neighbors_np[i, self.solo_vsav_idx] == tgt
                    ] if len(self.solo_vsav_idx) > 0 else np.array([], dtype=np.int64)

                    if len(solo_in_tgt) > 0:
                        # Swap : le premier VSAV solo du CS cible revient au CS source
                        swap_v = solo_in_tgt[0]
                        neighbors_np[i, swap_v] = src

                    # Déplacer le couple (v + son pair) vers le CS cible
                    neighbors_np[i, v]      = tgt
                    neighbors_np[i, pair_v] = tgt

                else:
                    # --- VSAV solo ou PSE solo : déplacement simple ---
                    neighbors_np[i, v] = tgt

        return torch.from_numpy(neighbors_np).to(self.device)

    # ------------------------------------------------------------------
    def optimize(
        self,
        max_iter:      int   = 200,
        T_init:        float = 100.0,
        T_min:         float = 0.5,
        alpha:         float = 0.95,
        verbose:       bool  = True,
        random_window: bool  = False,
        window_size:   int | None = None,
        early_stop_window: int = 5,
        early_stop_eps:    float = 1.0,
        n_moves_init:      float = 3.0,
    ) -> OptimResult:
        """
        Lance l'optimisation SA batch.

        À chaque itération :
          1. Génère n_batch voisins
          2. Les évalue en parallèle sur GPU
          3. Prend le meilleur et applique le critère de Metropolis

        Parameters
        ----------
        max_iter  : nombre d'itérations
        T_init    : température initiale
        T_min     : température minimale (arrêt anticipé)
        alpha     : coefficient de refroidissement (T *= alpha)
        verbose   : afficher les logs par itération

        Returns
        -------
        OptimResult
        """
        t_start = time.time()
        history = []

        N_total = self.sim._N_total
        _wsize  = window_size or (self.sim.N_inter if not random_window else 10_000)

        if random_window and window_size is None:
            raise ValueError("random_window=True requiert window_size explicite")

        current = self.initial_alloc.clone()
        if random_window:
            start = random.randint(0, N_total - _wsize)
            self.sim.set_window(start, _wsize)
        current_score = self.sim.evaluate_single(current)

        best = current.clone()
        best_score = current_score

        T = T_init
        no_improve = 0

        if verbose:
            print(f"[SA] Départ  score={current_score:.1f}  device={self.device}")

        for it in range(max_iter):
            if T < T_min:
                if verbose:
                    print(f"[SA] Température minimale atteinte à l'itération {it}")
                break

            # Nombre de mutations proportionnel à la température
            n_moves = max(1, int(np.random.poisson(max(1.0, T / T_init * n_moves_init))))

            # 1) Changer de fenêtre si mode aléatoire
            if random_window:
                start = random.randint(0, N_total - _wsize)
                self.sim.set_window(start, _wsize)

            # 2) Générer et évaluer le batch
            neighbors = self._generate_neighbors(current, self.n_batch, n_moves)
            scores    = self.sim.simulate_batch(neighbors)   # [n_batch]

            # 2) Meilleur voisin du batch
            best_idx           = scores.argmin().item()
            best_neigh_score   = scores[best_idx].item()
            best_neigh         = neighbors[best_idx]

            # 3) Critère de Metropolis sur le meilleur du batch
            delta = best_neigh_score - current_score
            accept = delta < 0 or random.random() < math.exp(-delta / max(T, 1e-9))

            if accept:
                current       = best_neigh.clone()
                current_score = best_neigh_score
                no_improve    = 0

                if current_score < best_score:
                    best       = current.clone()
                    best_score = current_score
            else:
                no_improve += 1

            T *= alpha

            # Perturbation si bloqué (redémarrage partiel)
            if no_improve >= 50:
                if verbose:
                    print(f"[SA] it={it} — pas d'amélioration depuis 50 iter, perturbation")
                current = best.clone()
                current_score = best_score
                T = max(T_init * 0.3, T_min * 2)
                no_improve = 0

            elapsed = time.time() - t_start
            history.append(dict(
                iter=it, score=current_score, best=best_score, T=T, elapsed=elapsed
            ))

            # Arret anticipe : gain sur les N dernieres iter < eps
            if len(history) > early_stop_window:
                recent_best = [h["best"] for h in history[-early_stop_window-1:]]
                gain = recent_best[0] - recent_best[-1]
                if gain < early_stop_eps:
                    if verbose:
                        print(f"[SA] Arret anticipe a iter {it} (gain {early_stop_window} dernieres iter: {gain:.2f} < {early_stop_eps})")
                    break

            if verbose and (it % 10 == 0 or it < 5):
                print(
                    f"[SA] it={it:4d}  T={T:6.2f}  score={current_score:.1f}"
                    f"  best={best_score:.1f}  n_moves={n_moves}"
                    f"  accept={'✓' if accept else '✗'}"
                    f"  {elapsed:.0f}s"
                )

        elapsed_total = time.time() - t_start
        if verbose:
            print(f"\n[SA] Terminé — meilleur score: {best_score:.2f}  ({elapsed_total:.0f}s)")

        return OptimResult(
            best_alloc=best.cpu(),
            best_score=best_score,
            history=history,
            elapsed_s=elapsed_total,
        )


# ---------------------------------------------------------------------------
# Algorithme génétique batch (bonus)
# ---------------------------------------------------------------------------

class GPUBatchGA:
    """
    Algorithme génétique avec évaluation batch GPU.

    Population entière évaluée en un seul appel GPU par génération.
    Respecte les mêmes contraintes que GPUBatchSA (freeze POMPE, couples, swap).
    """

    def __init__(
        self,
        simulator:         GPUBatchSimulator,
        initial_alloc:     np.ndarray | torch.Tensor,
        N_cs:              int,
        pop_size:          int = 200,
        elite_size:        int = 10,
        mutation_rate:     float = 0.05,
        vehicle_type:      np.ndarray | None = None,
        vsav_pse_pair_idx: np.ndarray | None = None,
        pse_vsav_pair_idx: np.ndarray | None = None,
    ):
        self.sim           = simulator
        self.N_cs          = N_cs
        self.pop_size      = pop_size
        self.elite_size    = elite_size
        self.mutation_rate = mutation_rate
        self.device        = simulator.device
        self.V             = simulator.V

        if isinstance(initial_alloc, np.ndarray):
            initial_alloc = torch.from_numpy(initial_alloc.astype(np.int64))
        self.initial_alloc = initial_alloc.to(self.device)

        # Réutiliser la même logique de contraintes que SA
        if vehicle_type is not None:
            self.vtype_np = vehicle_type.astype(np.int64)
        else:
            self.vtype_np = np.zeros(self.V, dtype=np.int64)

        self.movable_mask = self.vtype_np != POMPE_TYPE
        self.movable_idx  = np.where(self.movable_mask)[0]

        vsav_pse = vsav_pse_pair_idx if vsav_pse_pair_idx is not None \
            else np.full(self.V, -1, dtype=np.int32)
        pse_vsav = pse_vsav_pair_idx if pse_vsav_pair_idx is not None \
            else np.full(self.V, -1, dtype=np.int32)

        self.pair_of = np.full(self.V, -1, dtype=np.int32)
        for v in range(self.V):
            if self.vtype_np[v] == VSAV_TYPE and vsav_pse[v] >= 0:
                self.pair_of[v] = vsav_pse[v]
            elif self.vtype_np[v] == PSE_TYPE and pse_vsav[v] >= 0:
                self.pair_of[v] = pse_vsav[v]

        self.solo_vsav_idx = np.array([
            v for v in range(self.V)
            if self.vtype_np[v] == VSAV_TYPE and vsav_pse[v] < 0
        ], dtype=np.int64)

        # Poids de sélection : VSAV solo ×3
        solo_vsav_set = set(self.solo_vsav_idx)
        weights = np.ones(len(self.movable_idx), dtype=np.float64)
        for k, v in enumerate(self.movable_idx):
            if v in solo_vsav_set:
                weights[k] = 3.0
        self.movable_weights = weights / weights.sum()

    def _apply_constrained_mutation(self, alloc_np: np.ndarray, n_mut: int) -> np.ndarray:
        """Applique n_mut mutations avec contraintes (freeze, couples, swap)."""
        movable = self.movable_idx
        if len(movable) == 0:
            return alloc_np

        chosen    = movable[np.random.choice(len(movable), size=n_mut, p=self.movable_weights)]
        target_cs = np.random.randint(0, self.N_cs, size=n_mut)

        for j in range(n_mut):
            v   = chosen[j]
            tgt = target_cs[j]
            src = alloc_np[v]
            if tgt == src:
                continue

            pair_v = self.pair_of[v]
            if pair_v >= 0:
                solo_in_tgt = self.solo_vsav_idx[
                    alloc_np[self.solo_vsav_idx] == tgt
                ] if len(self.solo_vsav_idx) > 0 else np.array([], dtype=np.int64)
                if len(solo_in_tgt) > 0:
                    alloc_np[solo_in_tgt[0]] = src
                alloc_np[v]      = tgt
                alloc_np[pair_v] = tgt
            else:
                alloc_np[v] = tgt

        return alloc_np

    def _init_population(self) -> torch.Tensor:
        """[pop_size, V] int64 — initial alloc + pop_size-1 mutations contraintes."""
        init_np = self.initial_alloc.cpu().numpy()
        pop_list = [init_np.copy()]
        for _ in range(1, self.pop_size):
            ind = init_np.copy()
            n_mut = max(1, int(np.random.poisson(len(self.movable_idx) * 0.05)))
            self._apply_constrained_mutation(ind, n_mut)
            pop_list.append(ind)
        return torch.from_numpy(np.stack(pop_list)).to(self.device)

    def _crossover(self, p1_np: np.ndarray, p2_np: np.ndarray) -> np.ndarray:
        """Crossover respectant les contraintes : on prend p1 comme base,
        puis pour chaque couple/véhicule de p2 après le point de coupure,
        on applique le déplacement comme une mutation."""
        child = p1_np.copy()
        cut = random.randint(1, self.V - 1)
        # Appliquer les différences de p2 après le point de coupure
        for v in range(cut, self.V):
            if not self.movable_mask[v]:
                continue  # POMPE : on ne touche pas
            if p2_np[v] != child[v]:
                tgt = int(p2_np[v])
                src = int(child[v])
                pair_v = self.pair_of[v]
                if pair_v >= 0:
                    solo_in_tgt = self.solo_vsav_idx[
                        child[self.solo_vsav_idx] == tgt
                    ] if len(self.solo_vsav_idx) > 0 else np.array([], dtype=np.int64)
                    if len(solo_in_tgt) > 0:
                        child[solo_in_tgt[0]] = src
                    child[v]      = tgt
                    child[pair_v] = tgt
                else:
                    child[v] = tgt
        return child

    def _mutate(self, ind_np: np.ndarray) -> np.ndarray:
        """Mutations aléatoires avec contraintes."""
        n_mut = int((np.random.rand(self.V) < self.mutation_rate).sum())
        if n_mut > 0:
            ind_np = ind_np.copy()
            self._apply_constrained_mutation(ind_np, n_mut)
        return ind_np

    def optimize(
        self,
        n_generations: int  = 100,
        verbose:       bool = True,
    ) -> OptimResult:
        t_start = time.time()
        history = []

        pop    = self._init_population()
        scores = self.sim.simulate_batch(pop)
        best_idx   = scores.argmin().item()
        best_score = scores[best_idx].item()
        best       = pop[best_idx].clone()

        if verbose:
            print(f"[GA] Génération 0 — best={best_score:.2f}  device={self.device}")

        for gen in range(n_generations):
            # Sélection par tournoi
            sorted_idx = scores.argsort().cpu().numpy()
            elite_np = [pop[i].cpu().numpy() for i in sorted_idx[:self.elite_size]]
            pop_np   = pop.cpu().numpy()
            scores_np = scores.cpu().numpy()

            new_pop_np = list(elite_np)

            while len(new_pop_np) < self.pop_size:
                t1 = random.randint(0, self.pop_size - 1)
                t2 = random.randint(0, self.pop_size - 1)
                p1 = pop_np[t1 if scores_np[t1] < scores_np[t2] else t2]

                t3 = random.randint(0, self.pop_size - 1)
                t4 = random.randint(0, self.pop_size - 1)
                p2 = pop_np[t3 if scores_np[t3] < scores_np[t4] else t4]

                child = self._mutate(self._crossover(p1, p2))
                new_pop_np.append(child)

            pop    = torch.from_numpy(np.stack(new_pop_np[:self.pop_size])).to(self.device)
            scores = self.sim.simulate_batch(pop)
            bi     = scores.argmin().item()
            gs     = scores[bi].item()

            if gs < best_score:
                best_score = gs
                best       = pop[bi].clone()

            elapsed = time.time() - t_start
            history.append(dict(gen=gen, best=best_score, elapsed=elapsed))

            if verbose and gen % 10 == 0:
                print(f"[GA] gen={gen:4d}  best={best_score:.2f}  {elapsed:.0f}s")

        elapsed_total = time.time() - t_start
        if verbose:
            print(f"\n[GA] Terminé — best={best_score:.2f}  ({elapsed_total:.0f}s)")

        return OptimResult(
            best_alloc=best.cpu(),
            best_score=best_score,
            history=history,
            elapsed_s=elapsed_total,
        )