run_gpu_optim.py

"""
Point d'entrée pour l'optimisation GPU.

Usage
-----
/usr/local/Anaconda3-2024.10/bin/python3.12 run_gpu_optim.py \\
    --algo sa \\
    --n_batch 512 \\
    --max_iter 200 \\
    --metric temps_moyen \\
    --window_size 10000

Arguments
---------
--algo         sa | ga          algorithme (défaut: sa)
--n_batch      int              voisins évalués en parallèle par itération SA  (défaut: 512)
--max_iter     int              itérations SA ou générations GA (défaut: 200)
--metric       str              temps_moyen | taux_couverture | manque_local | mixte
--window_size  int              nombre d'interventions dans la fenêtre (défaut: 10000)
--window_start int              première intervention de la fenêtre (défaut: 0)
--T_init       float            température initiale SA (défaut: 100)
--alpha        float            coefficient de refroidissement SA (défaut: 0.95)
--pop_size     int              taille population GA (défaut: 200)
--rebuild_cache                 force le recalcul des matrices de trajet
--cpu                           forcer l'exécution sur CPU (debug)
--output       str              chemin du CSV résultat (défaut: results/best_allocation.csv)
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch  # AVANT numpy (sinon fbgemm.dll crash sur Windows)
import numpy as np
import pandas as pd

# Ajouter le dossier courant et la racine au path
sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent.parent))

from data_loader import load_data
from precompute  import build_travel_matrices
from gpu_simulator import GPUBatchSimulator
from gpu_optimizer import GPUBatchSA, GPUBatchGA


def parse_args():
    p = argparse.ArgumentParser(description="Optimisation allocation engins BSPP (GPU)")
    p.add_argument("--algo",          default="sa",          choices=["sa", "ga"])
    p.add_argument("--n_batch",       type=int,   default=512)
    p.add_argument("--max_iter",      type=int,   default=200)
    p.add_argument("--metric",        default="temps_moyen",
                   choices=["temps_moyen", "taux_couverture", "manque_local", "mixte", "couverture_0", "p95_vsav", "p99_vsav", "std_q95_pression", "composite"])
    p.add_argument("--window_size",    type=int,   default=10_000)
    p.add_argument("--window_start",  type=int,   default=0)
    p.add_argument("--random_window", action="store_true",
                   help="Fenêtre glissante aléatoire à chaque itération SA")
    p.add_argument("--T_init",        type=float, default=100.0)
    p.add_argument("--T_min",         type=float, default=0.5)
    p.add_argument("--alpha",         type=float, default=0.95)
    p.add_argument("--pop_size",      type=int,   default=200)
    p.add_argument("--no_cstc_priority", action="store_true",
                   help="Désactiver la priorité CSTC (toujours envoyer le véhicule le plus proche)")
    p.add_argument("--no_early_stop", action="store_true",
                   help="Désactiver l'arrêt anticipé SA")
    p.add_argument("--n_moves_init",   type=float, default=3.0,
                   help="Moyenne du nb de mutations à T=T_init (SA, défaut 3)")
    p.add_argument("--rebuild_cache",  action="store_true")
    p.add_argument("--cpu",            action="store_true")
    p.add_argument("--output",         default=None,
                   help="Chemin CSV résultat (défaut: results/allocation_<metric>_optimale.csv)")
    p.add_argument("--valhalla_url",   default="http://localhost:8002",
                   help="URL du serveur Valhalla (défaut: http://localhost:8002)")
    return p.parse_args()


def main():
    args = parse_args()
    device = "cpu" if args.cpu else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[run] Device: {device}")
    if device == "cuda":
        print(f"[run] GPU: {torch.cuda.get_device_name(0)}")
        print(f"[run] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

    # ── 1. Chargement des données ────────────────────────────────────────────
    print("\n[run] Chargement des données...")
    interventions_np, vehicles_np, cs_list, cs_coords, vehicle_unavail = load_data()
    N_inter = len(interventions_np["timestamps"])
    V       = len(vehicles_np["ids"])
    N_cs    = len(cs_list)
    print(f"[run]   {N_inter} interventions | {V} véhicules | {N_cs} CS")

    # ── 2. Matrices de trajet ────────────────────────────────────────────────
    print("\n[run] Matrices de trajet...")
    travel_mat, return_mat = build_travel_matrices(
        cs_coords,
        interventions_np["x"],
        interventions_np["y"],
        interventions_np["x_final"],
        interventions_np["y_final"],
        timestamps=interventions_np["timestamps"],
        force_rebuild=args.rebuild_cache,
        valhalla_url=args.valhalla_url,
    )
    print(f"[run]   travel_matrix: {travel_mat.shape}  {travel_mat.nbytes/1e6:.1f} MB")

    # ── 3. Simulateur GPU ────────────────────────────────────────────────────
    print(f"\n[run] Initialisation simulateur (fenêtre [{args.window_start}:"
          f"{args.window_start + args.window_size}])...")
    sim = GPUBatchSimulator(
        interventions_np     = interventions_np,
        vehicles_np          = vehicles_np,
        travel_matrix        = travel_mat,
        return_travel_matrix = return_mat,
        vehicle_unavail      = vehicle_unavail,
        window_start         = args.window_start,
        window_size          = args.window_size,
        metric               = args.metric,
        device               = device,
        cstc_priority        = not args.no_cstc_priority,
    )
    print(f"[run]   {sim.N_inter} interventions valides dans la fenêtre")

    # Afficher la mémoire GPU utilisée
    if device == "cuda":
        used_mb = torch.cuda.memory_allocated() / 1e6
        print(f"[run]   VRAM utilisée: {used_mb:.0f} MB")

    # ── 4. Score de référence (allocation initiale) ──────────────────────────
    initial_alloc = torch.from_numpy(vehicles_np["initial_cs_idx"].astype(np.int64))
    ref_score = sim.evaluate_single(initial_alloc)
    print(f"\n[run] Score allocation initiale: {ref_score:.2f}")

    # ── 5. Estimation de la capacité batch ───────────────────────────────────
    if device == "cuda":
        free_vram = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()
        # Estimation : ~N_inter × V × 4 bytes par batch de 1 alloc
        mem_per_alloc = sim.N_inter * V * 4 * 4   # 4 tenseurs principaux × float32
        max_batch_est = int(free_vram / max(mem_per_alloc, 1))
        print(f"[run] VRAM libre: {free_vram/1e9:.1f} GB — batch max estimé: ~{max_batch_est}")
        if args.n_batch > max_batch_est:
            print(f"[run] Avertissement: --n_batch {args.n_batch} > estimé {max_batch_est},"
                  " risque d'OOM. Réduire si nécessaire.")

    # ── 6. Optimisation ─────────────────────────────────────────────────────
    cstc_label = " (sans priorité CSTC)" if args.no_cstc_priority else ""
    print(f"\n[run] Lancement {args.algo.upper()}  n_batch={args.n_batch}"
          f"  max_iter={args.max_iter}  metric={args.metric}{cstc_label}")

    if args.algo == "sa":
        optimizer = GPUBatchSA(
            simulator         = sim,
            initial_alloc     = initial_alloc,
            N_cs              = N_cs,
            n_batch           = args.n_batch,
            vehicle_type      = vehicles_np["type"],
            vsav_pse_pair_idx = vehicles_np["vsav_pse_pair_idx"],
            pse_vsav_pair_idx = vehicles_np["pse_vsav_pair_idx"],
        )
        result = optimizer.optimize(
            max_iter        = args.max_iter,
            T_init          = args.T_init,
            T_min           = args.T_min,
            alpha           = args.alpha,
            random_window   = args.random_window,
            window_size     = args.window_size if args.random_window else None,
            early_stop_eps  = -1.0 if args.no_early_stop else 1.0,
            n_moves_init    = args.n_moves_init,
        )
    else:
        optimizer = GPUBatchGA(
            simulator         = sim,
            initial_alloc     = initial_alloc,
            N_cs              = N_cs,
            pop_size          = args.n_batch,
            vehicle_type      = vehicles_np["type"],
            vsav_pse_pair_idx = vehicles_np["vsav_pse_pair_idx"],
            pse_vsav_pair_idx = vehicles_np["pse_vsav_pair_idx"],
        )
        result = optimizer.optimize(n_generations=args.max_iter)

    # ── 7. Export résultat ───────────────────────────────────────────────────
    import json

    improvement = (ref_score - result.best_score) / ref_score * 100
    print(f"\n[run] Score initial:   {ref_score:.2f}")
    print(f"[run] Meilleur score:  {result.best_score:.2f}")
    print(f"[run] Amélioration:    {improvement:+.1f}%")
    print(f"[run] Temps total:     {result.elapsed_s:.0f}s")

    best_alloc_np = result.best_alloc.numpy()
    output = args.output or f"results/allocation_{args.metric}_optimale.csv"
    out_path = Path(__file__).parent / output
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # ── CSV : allocation des véhicules ─────────────────────────────────────
    df_result = pd.DataFrame({
        "vehicle_id": vehicles_np["ids"],
        "cs":         [cs_list[idx] for idx in best_alloc_np],
        "type":       vehicles_np["type"],
        "cs_initial": [cs_list[idx] for idx in vehicles_np["initial_cs_idx"]],
        "changed":    best_alloc_np != vehicles_np["initial_cs_idx"],
    })
    df_result.to_csv(str(out_path), index=False)
    print(f"[run] Allocation sauvegardée: {out_path}")

    from collections import Counter
    init_c, final_c = Counter(), Counter()
    for t, ci, cf in zip(vehicles_np["type"], vehicles_np["initial_cs_idx"], best_alloc_np):
        init_c[(int(ci), int(t))]  += 1
        final_c[(int(cf), int(t))] += 1
    n_changed = sum(abs(final_c[k] - init_c[k]) for k in set(init_c) | set(final_c)) // 2
    print(f"[run] Véhicules déplacés (net par CS×type): {n_changed} / {V}")
    if df_result["changed"].any():
        print(df_result[df_result["changed"]][["vehicle_id", "cs_initial", "cs"]].to_string())

    # ── JSON : métriques + hyperparamètres ─────────────────────────────────
    # Métriques annualisées (sur toutes les interventions, pas juste la fenêtre)
    print("[run] Calcul des métriques annualisées (année complète)...")
    sim.set_window(0, None)  # fenêtre = toutes les interventions
    print(f"[run]   Évaluation sur {sim.N_inter} interventions")
    df_init  = sim.simulate_single_detailed(initial_alloc.to(device), vehicles_np, cs_list)
    df_final = sim.simulate_single_detailed(result.best_alloc.to(device), vehicles_np, cs_list)
    metrics_init  = GPUBatchSimulator.compute_all_metrics(df_init)["global"]
    metrics_final = GPUBatchSimulator.compute_all_metrics(df_final)["global"]

    # Composite + std_q95_pression sur l'année entière via simulate_batch
    print("[run] Scores composite/std_q95_pression annualisés...")
    alloc_batch = torch.stack([initial_alloc, result.best_alloc]).to(device)
    for extra_metric in ("composite", "std_q95_pression"):
        sim.metric = extra_metric
        sc = sim.simulate_batch(alloc_batch).cpu().tolist()
        metrics_init[extra_metric]  = round(float(sc[0]), 4)
        metrics_final[extra_metric] = round(float(sc[1]), 4)

    # Convergence
    bests = [h["best"] for h in result.history]
    last_improve_it = max((i for i, b in enumerate(bests) if b == bests[-1]), default=-1) if bests else -1

    report = {
        "algo": args.algo,
        "metric_optimisee": args.metric,
        "cstc_priority": not args.no_cstc_priority,
        "hyperparametres": {
            "max_iter": args.max_iter,
            "n_batch": args.n_batch,
            "window_size": args.window_size,
            "window_start": args.window_start,
            "random_window": args.random_window,
            **({"T_init": args.T_init, "T_min": args.T_min, "alpha": args.alpha}
               if args.algo == "sa" else {"pop_size": args.pop_size}),
        },
        "donnees": {
            "n_interventions_total": N_inter,
            "n_interventions_fenetre": sim.N_inter,
            "n_vehicules": V,
            "n_cs": N_cs,
        },
        "metriques_avant": {k: (round(v, 4) if isinstance(v, float) else v)
                           for k, v in metrics_init.items()},
        "metriques_apres": {k: (round(v, 4) if isinstance(v, float) else v)
                           for k, v in metrics_final.items()},
        "score_optimise_avant": round(ref_score, 4),
        "score_optimise_apres": round(result.best_score, 4),
        "amelioration_pct": round(improvement, 2),
        "convergence": {
            "iterations_effectuees": len(result.history),
            "derniere_amelioration_iter": last_improve_it,
            "T_finale": round(result.history[-1]["T"], 6) if result.history else None,
            "delta_20_dernieres_iter": round(bests[-21] - bests[-1], 2)
                                      if len(bests) > 20 else None,
        },
        "vehicules_deplaces": n_changed,
        "temps_total_s": round(result.elapsed_s, 1),
        "historique": result.history,
    }

    json_path = out_path.with_suffix(".json")
    with open(str(json_path), "w") as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    print(f"[run] Rapport sauvegardé: {json_path}")


if __name__ == "__main__":
    main()