run_simulation.py

"""Entry point for running NIAC prototype simulations."""
from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Tuple

from niac_prototype.environment.number_game import NumberGameModel
from niac_prototype.core import USING_MESA_STUB
from niac_prototype.utils.metrics import EpisodeMetrics, MetricsSink

DEFAULT_CONFIG: Dict[str, Any] = {
    "num_agents": 3,
    "episodes": 5,
    "communication_rounds": 2,
    "latent_dim": 128,
    "seed": 42,
}

logger = logging.getLogger("niac.simulation")


def load_config(path: Path | None) -> Dict[str, Any]:
    if path is None or not path.exists():
        return dict(DEFAULT_CONFIG)
    try:
        import yaml  # type: ignore
    except ImportError:
        print("PyYAML indisponible, configuration par défaut utilisée.", file=sys.stderr)
        return dict(DEFAULT_CONFIG)
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle) or {}
    config = dict(DEFAULT_CONFIG)
    config.update(data.get("simulation", {}))
    return config


def simulate_variant(mode: str, config: Dict[str, Any]) -> MetricsSink:
    metrics = MetricsSink()
    seed = config.get("seed")
    for episode_idx in range(config["episodes"]):
        episode_seed = None if seed is None else seed + episode_idx
        logger.info(
            "Début épisode %d/%d mode=%s seed=%s",
            episode_idx + 1,
            config["episodes"],
            mode,
            episode_seed,
        )
        model = NumberGameModel(
            num_agents=config["num_agents"],
            mode=mode,
            latent_dim=config["latent_dim"],
            seed=episode_seed,
        )
        start = time.perf_counter()
        for _ in range(config["communication_rounds"]):
            model.step()
            if model.episode_over:
                break
        duration = time.perf_counter() - start
        guesses = model.current_guesses
        success = any(value == model.true_sum() for value in guesses.values())
        if not guesses:
            # Force one guess based on agent zero observation to keep stats comparable
            guesses[0] = model.current_numbers[0]
        text_tokens = 0
        latent_messages = 0
        if mode == "text":
            for message in model.channel.text_history():
                text_tokens += len(message.content.split())
        else:
            latent_messages = len(model.channel.neural_history())
        metrics.log_episode(
            EpisodeMetrics(
                episode_id=episode_idx,
                success=success,
                text_tokens=text_tokens,
                latent_messages=latent_messages,
                latency_seconds=duration,
                notes={"true_sum": model.true_sum(), "guesses": len(guesses)},
            )
        )
        logger.info(
            "Fin épisode %d mode=%s success=%s guesses=%d true_sum=%d",
            episode_idx + 1,
            mode,
            success,
            len(guesses),
            model.true_sum(),
        )
    return metrics


def ensure_output_dir(path: Path | None) -> Path:
    target = path or Path("reports")
    target.mkdir(parents=True, exist_ok=True)
    return target


def write_json_report(
    summaries: Dict[str, Dict[str, float]],
    sinks: Dict[str, MetricsSink],
    output_dir: Path,
) -> Path:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    data = {
        "generated_at": timestamp,
        "summaries": summaries,
        "episodes": {
            mode: [
                {
                    "episode_id": ep.episode_id,
                    "success": ep.success,
                    "text_tokens": ep.text_tokens,
                    "latent_messages": ep.latent_messages,
                    "latency_seconds": ep.latency_seconds,
                    "notes": ep.notes,
                }
                for ep in sink.episodes()
            ]
            for mode, sink in sinks.items()
        },
    }
    json_path = output_dir / f"niac_metrics_{timestamp}.json"
    json_path.write_text(json.dumps(data, indent=2), encoding="utf-8")
    logger.info("Rapport JSON sauvegardé: %s", json_path)
    return json_path


def create_plot(
    summaries: Dict[str, Dict[str, float]],
    output_dir: Path,
) -> Tuple[Path | None, bool]:
    try:
        import matplotlib.pyplot as plt  # type: ignore
    except ImportError:  # pragma: no cover - optional dependency
        logger.warning("Matplotlib n'est pas installé; graphique non généré.")
        return None, False

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    variants = list(summaries.keys())
    success = [summaries[v]["success_rate"] for v in variants]
    latent = [summaries[v]["avg_latent_messages"] for v in variants]
    tokens = [summaries[v]["avg_text_tokens"] for v in variants]
    latency = [summaries[v]["avg_latency_seconds"] for v in variants]

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    axes[0].bar(variants, success, color=["#4caf50", "#2196f3"])
    axes[0].set_title("Taux de succès")
    axes[0].set_ylim(0, 1)

    axes[1].bar(variants, latent, color="#9c27b0", label="Latent")
    axes[1].bar(variants, tokens, color="#ff9800", bottom=latent, label="Texte")
    axes[1].set_title("Messages moyens")
    axes[1].legend()

    axes[2].bar(variants, latency, color="#607d8b")
    axes[2].set_title("Latence moyenne (s)")

    for ax in axes:
        ax.set_xlabel("Mode")
        ax.grid(axis="y", linestyle="--", alpha=0.3)

    fig.suptitle("Comparaison NIAC : texte vs neurones")
    fig.tight_layout()
    output_path = output_dir / f"niac_metrics_{timestamp}.png"
    fig.savefig(output_path, dpi=200)
    plt.close(fig)
    logger.info("Graphique sauvegardé: %s", output_path)
    return output_path, True


def main() -> None:
    parser = argparse.ArgumentParser(description="NIAC number-game simulation runner")
    parser.add_argument("--config", type=Path, default=None, help="Chemin vers config.yaml")
    parser.add_argument(
        "--mode",
        choices=["neural", "text", "both"],
        default="both",
        help="Type de simulation à exécuter",
    )
    parser.add_argument(
        "--plot",
        action="store_true",
        help="Générer un graphique comparatif (nécessite matplotlib).",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Répertoire de sortie pour les rapports/graphes (défaut: reports/).",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Sauvegarder les métriques détaillées au format JSON.",
    )
    args = parser.parse_args()
    config = load_config(args.config)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )
    logger.info(
        "Configuration chargée: %s (Mesa stub=%s)",
        {**config, "mode": args.mode},
        USING_MESA_STUB,
    )

    variants = [args.mode] if args.mode != "both" else ["text", "neural"]
    summaries: Dict[str, Dict[str, float]] = {}
    sinks: Dict[str, MetricsSink] = {}
    for variant in variants:
        logger.info("---- Simulation mode %s ----", variant)
        metrics = simulate_variant(variant, config)
        summary = metrics.summary()
        print(f"=== Résultats mode {variant} ===")
        if summary:
            for key, value in summary.items():
                print(f"{key}: {value}")
            logger.info("Synthèse mode %s: %s", variant, summary)
            summaries[variant] = summary
            sinks[variant] = metrics
        else:
            print("Aucune donnée collectée.")
            logger.warning("Aucune donnée collectée pour le mode %s", variant)

    if summaries:
        output_dir = ensure_output_dir(args.output_dir) if (args.plot or args.json) else None
        if args.json and output_dir:
            write_json_report(summaries, sinks, output_dir)
        if args.plot and output_dir:
            create_plot(summaries, output_dir)


if __name__ == "__main__":
    main()