#!/usr/bin/env python3
"""IteraGPU Lab v1 — arithmétique CPU et mesure d'un MLP synthétique sur UN GPU.

Aucun téléchargement ni installation. Python >= 3.10. Licence MIT.
Les fonctions précédant main() sont aussi intégrées au notebook autonome.
"""
import argparse
import gc
import json
import platform
import re
import sys
import time
from datetime import datetime, timezone
from decimal import Decimal, InvalidOperation, localcontext
from pathlib import Path


class MeasurementUnavailable(RuntimeError):
    """Prérequis manquant ou mesure impossible, jamais remplacé par un résultat CPU."""


def positive_integer(value, name):
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} doit être un entier strictement positif.")
    return value


def estimate_weights(parameters, bits, reserve_gib="0"):
    """Poids théoriques seulement ; la réserve est une hypothèse fournie par l'utilisateur."""
    positive_integer(parameters, "parameters")
    if isinstance(bits, bool) or bits not in (4, 8, 16, 32):
        raise ValueError("bits doit valoir 4, 8, 16 ou 32.")
    try:
        reserve = Decimal(str(reserve_gib))
    except InvalidOperation as error:
        raise ValueError("reserve_gib doit être un nombre décimal fini positif ou nul.") from error
    if not reserve.is_finite() or reserve < 0:
        raise ValueError("reserve_gib doit être un nombre décimal fini positif ou nul.")
    weight_bytes = (parameters * bits + 7) // 8
    with localcontext() as ctx:
        ctx.prec = max(40, len(str(parameters)) + 20)
        weight_gib = Decimal(weight_bytes) / Decimal(2**30)
        total = weight_gib + reserve
        return {
            "kind": "arithmetic_only",
            "parameters": parameters,
            "bits_per_parameter": bits,
            "theoretical_weight_bytes": weight_bytes,
            "theoretical_weight_gib": str(weight_gib),
            "assumed_reserve_gib": str(reserve),
            "weights_plus_assumed_reserve_gib": str(total),
            "scope": "Poids stockés, arrondis à l'octet supérieur ; pas de pic mesuré. "
                     "Métadonnées de quantification, activations, cache, gradients et optimiseur exclus. "
                     "La réserve ajoutée est une hypothèse, pas une garantie de capacité.",
        }


def parse_device(device):
    if not isinstance(device, str) or not re.fullmatch(r"cuda:[0-9]+", device):
        raise ValueError("Choisir un seul GPU explicitement : cuda:0, cuda:1, etc. CPU non mesuré.")
    return int(device.split(":")[1])


def load_torch():
    try:
        import torch
    except (ImportError, OSError) as error:
        raise MeasurementUnavailable(
            "PyTorch est absent ou ne peut pas être chargé. Aucune mesure GPU effectuée. "
            "Utilisez estimate sans PyTorch ; préparez séparément un environnement compatible "
            "avec votre GPU et votre pilote pour measure. Ce script n'installe rien."
        ) from error
    return torch


def environment_report(device="cuda:0"):
    index = parse_device(device)
    report = {
        "python_version": platform.python_version(),
        "platform": platform.system(),
        "requested_device": device,
        "torch_importable": False,
        "torch_version": None,
        "cuda_build_version": None,
        "hip_build_version": None,
        "gpu_available": False,
        "visible_device_count": 0,
        "selected_device_available": False,
        "selected_device_name": None,
        "selected_device_total_bytes": None,
        "driver_version": None,
        "driver_note": "Relever séparément la version du pilote avec l'outil du fournisseur.",
    }
    try:
        torch = load_torch()
        report.update(torch_importable=True, torch_version=str(torch.__version__),
                      cuda_build_version=torch.version.cuda, hip_build_version=torch.version.hip)
        report["gpu_available"] = bool(torch.cuda.is_available())
        report["visible_device_count"] = torch.cuda.device_count() if report["gpu_available"] else 0
        if report["gpu_available"] and index < report["visible_device_count"]:
            properties = torch.cuda.get_device_properties(index)
            report.update(selected_device_available=True, selected_device_name=properties.name,
                          selected_device_total_bytes=properties.total_memory)
    except MeasurementUnavailable as error:
        report["reason"] = str(error)
    except RuntimeError:
        report["reason"] = "Le runtime GPU ne répond pas correctement. Aucune mesure effectuée."
    return report


def record_phase(torch, device, phase, run_index, operation):
    """Pics absolus depuis reset, avec baseline ; le résultat reste vivant à la lecture."""
    torch.cuda.synchronize(device)
    baseline_allocated = torch.cuda.memory_allocated(device)
    baseline_reserved = torch.cuda.memory_reserved(device)
    torch.cuda.reset_peak_memory_stats(device)
    start = time.perf_counter()
    result = operation()
    torch.cuda.synchronize(device)
    elapsed = time.perf_counter() - start
    row = {
        "phase": phase,
        "run_index": run_index,
        "elapsed_seconds": elapsed,
        "baseline_allocated_bytes": baseline_allocated,
        "baseline_reserved_bytes": baseline_reserved,
        "end_allocated_bytes": torch.cuda.memory_allocated(device),
        "end_reserved_bytes": torch.cuda.memory_reserved(device),
        "peak_allocated_bytes": torch.cuda.max_memory_allocated(device),
        "peak_reserved_bytes": torch.cuda.max_memory_reserved(device),
    }
    return result, row


def measure(device="cuda:0", batch=2, context=128, width=1024, dtype="float32", warmup=3,
            repeats=5, seed=1729):
    index = parse_device(device)
    for name, value in (("batch", batch), ("context", context), ("width", width),
                        ("warmup", warmup), ("repeats", repeats)):
        positive_integer(value, name)
    if dtype not in ("float32", "float16", "bfloat16"):
        raise ValueError("dtype doit valoir float32, float16 ou bfloat16.")
    if isinstance(seed, bool) or not isinstance(seed, int) or not 0 <= seed < 2**63:
        raise ValueError("seed doit être un entier compris entre 0 et 2**63 - 1.")
    torch = load_torch()
    env = environment_report(device)
    if not env["selected_device_available"]:
        raise MeasurementUnavailable(
            f"GPU {device} indisponible : vérifier GPU visible, build PyTorch et pilote. "
            "Aucune mesure effectuée ; aucun repli vers le CPU."
        )
    selected = torch.device(device)
    rows = []
    with torch.cuda.device(index):
        if dtype == "bfloat16" and not torch.cuda.is_bf16_supported():
            raise MeasurementUnavailable("bfloat16 non déclaré pris en charge sur le GPU choisi.")
        scalar_type = getattr(torch, dtype)
        cpu_generator = torch.Generator(device="cpu").manual_seed(seed)
        gpu_generator = torch.Generator(device=selected).manual_seed(seed)
        previous_precision = torch.get_float32_matmul_precision()
        torch.set_float32_matmul_precision("highest")
        try:
            def make_model():
                # Initialisation CPU incluse dans model_load ; poids ensuite transférés au GPU choisi.
                model = torch.nn.Sequential(torch.nn.Linear(width, width), torch.nn.GELU(),
                                            torch.nn.Linear(width, width))
                with torch.no_grad():
                    for parameter in model.parameters():
                        if parameter.ndim > 1:
                            parameter.normal_(0, 0.02, generator=cpu_generator)
                        else:
                            parameter.zero_()
                return model.to(device=selected, dtype=scalar_type).eval()

            model, row = record_phase(torch, selected, "model_load", 0, make_model)
            rows.append(row)
            inputs, row = record_phase(
                torch, selected, "inputs", 0,
                lambda: torch.randn(batch, context, width, device=selected,
                                    dtype=scalar_type, generator=gpu_generator))
            rows.append(row)
            with torch.inference_mode():
                output, row = record_phase(torch, selected, "cold_forward", 0, lambda: model(inputs))
                rows.append(row)
                del output
                gc.collect()
                torch.cuda.synchronize(selected)

                def warmup_passes():
                    for _ in range(warmup):
                        temporary_output = model(inputs)
                        del temporary_output

                _, row = record_phase(torch, selected, "warmup", 0, warmup_passes)
                row["iterations"] = warmup
                rows.append(row)
                for repetition in range(1, repeats + 1):
                    output, row = record_phase(torch, selected, "warm_forward", repetition,
                                               lambda: model(inputs))
                    rows.append(row)
                    del output
                    gc.collect()
                    torch.cuda.synchronize(selected)
            parameter_count = sum(parameter.numel() for parameter in model.parameters())
            parameter_bytes = sum(parameter.numel() * parameter.element_size()
                                  for parameter in model.parameters())
            del inputs, model
            gc.collect()
            torch.cuda.synchronize(selected)
        finally:
            torch.set_float32_matmul_precision(previous_precision)
    return {
        "schema_version": 1,
        "artifact": "iteragpu-lab-v1",
        "measured_at_utc": datetime.now(timezone.utc).isoformat(),
        "environment": env,
        "configuration": {"device": device, "batch": batch, "context": context, "width": width,
                          "dtype": dtype, "warmup": warmup, "repeats": repeats, "seed": seed,
                          "float32_matmul_precision": "highest", "autocast": False,
                          "training": False},
        "synthetic_model": {"definition": "Linear(width,width), GELU, Linear(width,width)",
                            "parameter_count": parameter_count, "parameter_bytes": parameter_bytes},
        "scope": "Allocateur PyTorch de ce processus sur le seul device choisi. Ni mémoire totale "
                 "de la machine, ni autre processus, ni allocation externe à PyTorch, ni CPU. "
                 "cold_forward est le premier passage après initialisation du device et du modèle. "
                 "Le MLP synthétique ne mesure ni LLM, ni entraînement, ni qualité métier.",
        "interpretation": "allocated fait partie de reserved : ne pas les additionner. Les maxima "
                          "peuvent venir d'instants différents : ne pas soustraire les deux pics. "
                          "Les baselines et pics sont absolus ; le cache de l'allocateur est conservé.",
        "phases": rows,
    }


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    estimate = subparsers.add_parser("estimate", help="Arithmétique sans PyTorch, pas une mesure")
    estimate.add_argument("--parameters", type=int, required=True)
    estimate.add_argument("--bits", type=int, choices=(4, 8, 16, 32), required=True)
    estimate.add_argument("--reserve-gib", default="0")
    environment = subparsers.add_parser("environment", help="Disponibilité réelle, sans installation")
    environment.add_argument("--device", default="cuda:0")
    gpu = subparsers.add_parser("measure", help="MLP synthétique sur un seul GPU, aucun téléchargement")
    gpu.add_argument("--device", default="cuda:0")
    for name, default in (("batch", 2), ("context", 128), ("width", 1024),
                          ("warmup", 3), ("repeats", 5), ("seed", 1729)):
        gpu.add_argument("--" + name, type=int, default=default)
    gpu.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
    gpu.add_argument("--output", type=Path, help="Nouveau fichier JSON ; refuse d'écraser un fichier")
    args = parser.parse_args(argv)
    try:
        if args.command == "estimate":
            result = estimate_weights(args.parameters, args.bits, args.reserve_gib)
        elif args.command == "environment":
            result = environment_report(args.device)
        else:
            if args.output and args.output.exists():
                raise ValueError("Le fichier de sortie existe déjà : choisir un nouveau nom.")
            result = measure(args.device, args.batch, args.context, args.width, args.dtype,
                             args.warmup, args.repeats, args.seed)
        serialized = json.dumps(result, ensure_ascii=False, indent=2)
        if args.command == "measure" and args.output:
            with args.output.open("x", encoding="utf-8") as handle:
                handle.write(serialized + "\n")
        print(serialized)
        return 0
    except (ValueError, MeasurementUnavailable, OSError, RuntimeError) as error:
        print(f"Mesure non produite : {error}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
