#!/usr/bin/env python3
"""
GPU Throughput Benchmark for AWS EC2 Spot Instances (T4 GPU)

Tests PyTorch GPU utilization with two benchmark modes:
1. CNN on MNIST - lightweight, fast iteration
2. Transformer on synthetic data - more compute-intensive

Reports: iterations/sec, samples/sec, GPU memory usage, and utilization metrics.

Supports multiple precision modes with automatic fallback:
- FP16 AMP (default for Turing/Ampere+)
- FP32 (fallback if AMP fails)
- TF32 (Ampere+ only)
"""

from __future__ import annotations
import argparse
import os
import sys
import time
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets, transforms
from tqdm import tqdm


if TYPE_CHECKING:
    from collections.abc import Generator


# -----------------------------------------------------------------------------
# Diagnostic Functions
# -----------------------------------------------------------------------------


def run_cuda_diagnostics(device: torch.device) -> dict[str, bool]:
    """
    Run diagnostic tests to verify CUDA/cuBLAS functionality.
    Returns dict of test_name -> passed.
    """
    results: dict[str, bool] = {}

    if device.type != "cuda":
        print("  Skipping CUDA diagnostics (CPU mode)")
        return results

    print("\n" + "-" * 40)
    print("Running CUDA Diagnostics")
    print("-" * 40)

    # Test 1: Basic FP32 matmul
    try:
        a = torch.randn(256, 256, device=device)
        b = torch.randn(256, 256, device=device)
        _c = torch.mm(a, b)
        torch.cuda.synchronize()
        results["fp32_matmul"] = True
        print("  ✓ FP32 matmul: PASSED")
    except Exception as e:
        results["fp32_matmul"] = False
        print(f"  ✗ FP32 matmul: FAILED - {e}")

    # Test 2: FP16 matmul (no autocast)
    try:
        a = torch.randn(256, 256, device=device, dtype=torch.float16)
        b = torch.randn(256, 256, device=device, dtype=torch.float16)
        _c = torch.mm(a, b)
        torch.cuda.synchronize()
        results["fp16_matmul"] = True
        print("  ✓ FP16 matmul: PASSED")
    except Exception as e:
        results["fp16_matmul"] = False
        print(f"  ✗ FP16 matmul: FAILED - {e}")

    # Test 3: FP16 matmul with autocast
    try:
        a = torch.randn(256, 256, device=device)
        b = torch.randn(256, 256, device=device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
            _c = torch.mm(a, b)
        torch.cuda.synchronize()
        results["fp16_autocast"] = True
        print("  ✓ FP16 autocast matmul: PASSED")
    except Exception as e:
        results["fp16_autocast"] = False
        print(f"  ✗ FP16 autocast matmul: FAILED - {e}")

    # Test 4: Linear layer with autocast (common GEMM pattern)
    try:
        linear = nn.Linear(512, 512).to(device)
        x = torch.randn(64, 512, device=device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
            _y = linear(x)
        torch.cuda.synchronize()
        results["fp16_linear"] = True
        print("  ✓ FP16 linear layer: PASSED")
    except Exception as e:
        results["fp16_linear"] = False
        print(f"  ✗ FP16 linear layer: FAILED - {e}")

    # Test 5: Conv2d with autocast
    try:
        conv = nn.Conv2d(64, 128, 3, padding=1).to(device)
        x = torch.randn(16, 64, 32, 32, device=device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
            _y = conv(x)
        torch.cuda.synchronize()
        results["fp16_conv2d"] = True
        print("  ✓ FP16 conv2d: PASSED")
    except Exception as e:
        results["fp16_conv2d"] = False
        print(f"  ✗ FP16 conv2d: FAILED - {e}")

    # Test 6: Batched matmul (transformer attention pattern)
    try:
        # Simulates attention: (batch, heads, seq, dim) @ (batch, heads, dim, seq)
        a = torch.randn(8, 8, 128, 64, device=device, dtype=torch.float16)
        b = torch.randn(8, 8, 64, 128, device=device, dtype=torch.float16)
        _c = torch.matmul(a, b)
        torch.cuda.synchronize()
        results["fp16_batched_matmul"] = True
        print("  ✓ FP16 batched matmul: PASSED")
    except Exception as e:
        results["fp16_batched_matmul"] = False
        print(f"  ✗ FP16 batched matmul: FAILED - {e}")

    print("-" * 40)

    # Summary
    passed = sum(results.values())
    total = len(results)
    print(f"Diagnostics: {passed}/{total} tests passed")

    if passed < total:
        failed_tests = [k for k, v in results.items() if not v]
        print(f"Failed tests: {', '.join(failed_tests)}")
        print("\nRecommendation: Use --precision fp32 to bypass FP16 issues")

    print("-" * 40 + "\n")
    return results


class PrecisionMode(Enum):
    """Supported precision modes for training."""

    FP32 = "fp32"
    FP16 = "fp16"
    BF16 = "bf16"
    TF32 = "tf32"


@dataclass(frozen=True)
class BenchmarkConfig:
    """Configuration for benchmark runs."""

    batch_size: int = 256
    num_warmup_batches: int = 10
    num_benchmark_batches: int = 100
    num_workers: int = 4
    pin_memory: bool = True
    precision: PrecisionMode = PrecisionMode.FP16


@dataclass
class BenchmarkResult:
    """Results from a benchmark run."""

    model_name: str
    total_samples: int
    total_time_sec: float
    peak_memory_mb: float
    avg_batch_time_ms: float
    precision_mode: str

    @property
    def samples_per_sec(self) -> float:
        return self.total_samples / self.total_time_sec

    @property
    def batches_per_sec(self) -> float:
        return 1000.0 / self.avg_batch_time_ms

    def __str__(self) -> str:
        return (
            f"\n{'=' * 60}\n"
            f"Benchmark Results: {self.model_name}\n"
            f"{'=' * 60}\n"
            f"  Precision mode: {self.precision_mode}\n"
            f"  Total samples processed: {self.total_samples:,}\n"
            f"  Total time: {self.total_time_sec:.2f}s\n"
            f"  Throughput: {self.samples_per_sec:,.1f} samples/sec\n"
            f"  Throughput: {self.batches_per_sec:.1f} batches/sec\n"
            f"  Avg batch time: {self.avg_batch_time_ms:.2f}ms\n"
            f"  Peak GPU memory: {self.peak_memory_mb:.1f}MB\n"
            f"{'=' * 60}\n"
        )


# -----------------------------------------------------------------------------
# CNN Model for MNIST
# -----------------------------------------------------------------------------


class MNISTConvNet(nn.Module):
    """
    Simple but non-trivial CNN for MNIST.
    ~1.2M parameters - enough to stress GPU without being excessive.
    """

    def __init__(self) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256, 256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, 10),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        return self.classifier(x)


# -----------------------------------------------------------------------------
# Transformer Model (GPT-style decoder)
# -----------------------------------------------------------------------------


class TransformerBlock(nn.Module):
    """Single transformer decoder block with pre-norm."""

    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor:
        # Pre-norm architecture
        normed = self.ln1(x)
        attn_out, _ = self.attn(normed, normed, normed, attn_mask=attn_mask, need_weights=False)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x


class MiniGPT(nn.Module):
    """
    Small GPT-style transformer for benchmarking.
    ~25M parameters - representative of real workloads.
    """

    def __init__(
        self,
        vocab_size: int = 32000,
        d_model: int = 512,
        n_heads: int = 8,
        n_layers: int = 6,
        d_ff: int = 2048,
        max_seq_len: int = 256,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.max_seq_len = max_seq_len

        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        self.dropout = nn.Dropout(dropout)

        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)])

        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

        # Weight tying
        self.head.weight = self.token_emb.weight

        self._init_weights()

    def _init_weights(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len = x.shape
        assert seq_len <= self.max_seq_len, f"Sequence length {seq_len} > max {self.max_seq_len}"

        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)

        x = self.dropout(self.token_emb(x) + self.pos_emb(positions))

        # Causal mask for autoregressive modeling
        causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool), diagonal=1)

        for block in self.blocks:
            x = block(x, attn_mask=causal_mask)

        x = self.ln_f(x)
        return self.head(x)


# -----------------------------------------------------------------------------
# Data Loading
# -----------------------------------------------------------------------------


def get_mnist_loader(config: BenchmarkConfig, device: torch.device) -> DataLoader:
    """Load MNIST dataset with standard preprocessing."""
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
        ]
    )

    dataset = datasets.MNIST(root="/tmp/data", train=True, download=True, transform=transform)

    return DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=config.num_workers,
        pin_memory=config.pin_memory and device.type == "cuda",
        persistent_workers=config.num_workers > 0,
    )


def get_synthetic_text_loader(
    config: BenchmarkConfig,
    vocab_size: int = 32000,
    seq_len: int = 256,
    num_samples: int = 50000,
) -> DataLoader:
    """Generate synthetic token sequences for transformer benchmarking."""
    # Random token IDs (simulates real tokenized text distribution)
    data = torch.randint(0, vocab_size, (num_samples, seq_len))
    # Labels are next-token shifted (standard LM objective)
    labels = torch.randint(0, vocab_size, (num_samples, seq_len))

    dataset = TensorDataset(data, labels)

    return DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=0,  # Synthetic data is fast enough
        pin_memory=True,
    )


# -----------------------------------------------------------------------------
# Benchmark Runner
# -----------------------------------------------------------------------------


@contextmanager
def cuda_timer(device: torch.device) -> Generator[dict[str, float]]:
    """Context manager for accurate CUDA timing using events."""
    result: dict[str, float] = {}

    if device.type == "cuda":
        torch.cuda.synchronize(device)
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()
        yield result
        end_event.record()
        torch.cuda.synchronize(device)
        result["elapsed_ms"] = start_event.elapsed_time(end_event)
    else:
        start = time.perf_counter()
        yield result
        result["elapsed_ms"] = (time.perf_counter() - start) * 1000


def run_benchmark(
    model: nn.Module,
    loader: DataLoader,
    config: BenchmarkConfig,
    device: torch.device,
    model_name: str,
    precision: PrecisionMode,
    is_lm: bool = False,
) -> BenchmarkResult:
    """
    Run training benchmark with warmup phase.

    Args:
        model: PyTorch model to benchmark
        loader: DataLoader providing batches
        config: Benchmark configuration
        device: Target device
        model_name: Name for reporting
        precision: Precision mode to use
        is_lm: If True, use language modeling loss (ignore_index=-100)
    """
    model = model.to(device)
    model.train()

    # Configure precision-specific settings
    use_amp = precision in (PrecisionMode.FP16, PrecisionMode.BF16)
    amp_dtype = torch.float16 if precision == PrecisionMode.FP16 else torch.bfloat16

    # GradScaler is only needed for FP16 (BF16 has sufficient dynamic range)
    use_scaler = precision == PrecisionMode.FP16 and device.type == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=use_scaler)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)

    data_iter = iter(loader)
    batch_times: list[float] = []

    total_batches = config.num_warmup_batches + config.num_benchmark_batches

    print(f"\nRunning {model_name} benchmark...")
    print(f"  Precision: {precision.value}")
    print(f"  AMP enabled: {use_amp}")
    print(f"  GradScaler enabled: {use_scaler}")
    print(f"  Warmup batches: {config.num_warmup_batches}")
    print(f"  Benchmark batches: {config.num_benchmark_batches}")
    print(f"  Batch size: {config.batch_size}")

    pbar = tqdm(range(total_batches), desc=model_name, unit="batch")
    for batch_idx in pbar:
        # Get next batch, cycling if needed
        try:
            batch = next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            batch = next(data_iter)

        inputs = batch[0].to(device, non_blocking=True)
        targets = batch[1].to(device, non_blocking=True)

        with cuda_timer(device) as timer:
            optimizer.zero_grad(set_to_none=True)

            # Use autocast only when AMP is enabled
            with torch.amp.autocast(
                device_type=device.type,
                dtype=amp_dtype,
                enabled=use_amp,
            ):
                outputs = model(inputs)

                if is_lm:
                    # Reshape for cross-entropy: (batch * seq_len, vocab_size)
                    loss = F.cross_entropy(
                        outputs.view(-1, outputs.size(-1)),
                        targets.view(-1),
                    )
                else:
                    loss = F.cross_entropy(outputs, targets)

            # Backward pass with optional gradient scaling
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

        # Only record times after warmup
        is_benchmark = batch_idx >= config.num_warmup_batches
        if is_benchmark:
            batch_times.append(timer["elapsed_ms"])

        # Update progress bar
        phase = "bench" if is_benchmark else "warmup"
        postfix: dict[str, str] = {"phase": phase, "loss": f"{loss.item():.4f}"}
        if is_benchmark and batch_times:
            sps = config.batch_size / (batch_times[-1] / 1000)
            postfix["samples/s"] = f"{sps:,.0f}"
        if device.type == "cuda":
            mem_mb = torch.cuda.memory_allocated(device) / (1024**2)
            postfix["gpu_mem"] = f"{mem_mb:.0f}MB"
        pbar.set_postfix(postfix)

    # Compute statistics
    avg_batch_time = sum(batch_times) / len(batch_times)
    total_time = sum(batch_times) / 1000  # Convert to seconds
    total_samples = config.num_benchmark_batches * config.batch_size

    peak_memory = 0.0
    if device.type == "cuda":
        peak_memory = torch.cuda.max_memory_allocated(device) / (1024 * 1024)

    return BenchmarkResult(
        model_name=model_name,
        total_samples=total_samples,
        total_time_sec=total_time,
        peak_memory_mb=peak_memory,
        avg_batch_time_ms=avg_batch_time,
        precision_mode=precision.value,
    )


# -----------------------------------------------------------------------------
# System Information and GPU Configuration
# -----------------------------------------------------------------------------


def get_gpu_architecture(device: torch.device) -> tuple[int, int]:
    """Get GPU compute capability (major, minor)."""
    if device.type != "cuda":
        return (0, 0)
    props = torch.cuda.get_device_properties(device)
    return (props.major, props.minor)


def configure_precision(device: torch.device, requested: PrecisionMode) -> PrecisionMode:
    """
    Configure and validate precision mode based on GPU capabilities.

    GPU Architecture Reference:
    - Turing (T4): sm_75 - Supports FP16 tensor cores, NO native BF16, NO TF32
    - Ampere (A100, A10, 3090): sm_80/86 - Supports FP16, BF16, TF32
    - Hopper (H100): sm_90 - Full support for all modes

    Returns the actual precision mode that will be used.
    """
    if device.type != "cuda":
        print("  CPU mode: Using FP32")
        return PrecisionMode.FP32

    major, minor = get_gpu_architecture(device)
    sm_version = major * 10 + minor

    print(f"  GPU compute capability: sm_{sm_version}")

    # =========================================================================
    # CRITICAL: Disable problematic cuBLAS features that can cause GEMM errors
    # These settings improve stability on older architectures like Turing (T4)
    # =========================================================================

    # Disable TF32 on non-Ampere hardware (TF32 is Ampere+ only)
    if sm_version < 80:
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        print("  TF32 disabled (requires sm_80+)")

    # Disable reduced precision reductions in FP16 GEMMs
    # This can cause overflow/execution failures on some cuBLAS versions
    # See: https://docs.pytorch.org/docs/stable/notes/cuda.html
    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
    print("  FP16 reduced precision reduction disabled (for stability)")

    # Also disable BF16 reduced precision reduction for consistency
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
    print("  BF16 reduced precision reduction disabled (for stability)")

    # TF32 requires Ampere or newer (sm_80+)
    if requested == PrecisionMode.TF32:
        if sm_version >= 80:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            print("  TF32 mode enabled for matmul and cuDNN")
            return PrecisionMode.TF32
        else:
            print(f"  WARNING: TF32 requires sm_80+, but GPU is sm_{sm_version}")
            print("  Falling back to FP16 AMP")
            requested = PrecisionMode.FP16

    # BF16 requires Ampere or newer (sm_80+) for efficient operation
    if requested == PrecisionMode.BF16:
        if sm_version >= 80 and torch.cuda.is_bf16_supported():
            print("  BF16 mode enabled")
            return PrecisionMode.BF16
        else:
            print(f"  WARNING: BF16 not efficiently supported on sm_{sm_version}")
            print("  Falling back to FP16 AMP")
            requested = PrecisionMode.FP16

    # FP16 works on Volta (sm_70) and newer
    if requested == PrecisionMode.FP16:
        if sm_version >= 70:
            print("  FP16 AMP mode enabled")
            return PrecisionMode.FP16
        else:
            print(f"  WARNING: FP16 tensor cores require sm_70+, but GPU is sm_{sm_version}")
            print("  Falling back to FP32")
            return PrecisionMode.FP32

    # FP32 always works
    print("  FP32 mode (no mixed precision)")
    return PrecisionMode.FP32


def print_system_info(
    requested_precision: PrecisionMode, force_cpu: bool = False
) -> tuple[torch.device, PrecisionMode]:
    """Print system and CUDA information, return device and actual precision mode."""
    print("\n" + "=" * 60)
    print("System Information")
    print("=" * 60)
    print(f"PyTorch version: {torch.__version__}")
    print(f"Python version: {sys.version.split()[0]}")

    if torch.cuda.is_available() and not force_cpu:
        device = torch.device("cuda")
        print("CUDA available: Yes")
        print(f"CUDA version: {torch.version.cuda}")

        cudnn_version = torch.backends.cudnn.version()
        if cudnn_version:
            print(f"cuDNN version: {cudnn_version}")

        print(f"Device count: {torch.cuda.device_count()}")

        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            print(f"\nGPU {i}: {props.name}")
            print(f"  Compute capability: {props.major}.{props.minor}")
            print(f"  Total memory: {props.total_memory / (1024**3):.1f}GB")
            print(f"  SM count: {props.multi_processor_count}")

        print("\nPrecision Configuration:")
        actual_precision = configure_precision(device, requested_precision)

        # Set deterministic cuBLAS workspace config for stability
        # This can help avoid sporadic GEMM failures
        if "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
            print("  Set CUBLAS_WORKSPACE_CONFIG=:4096:8 for stability")

    else:
        device = torch.device("cpu")
        actual_precision = PrecisionMode.FP32
        if force_cpu:
            print("CPU-only mode requested (--cpu flag)")
        else:
            print("CUDA available: No (running on CPU)")
        print("Running on CPU for benchmarking")

    print("=" * 60)
    return device, actual_precision


# -----------------------------------------------------------------------------
# Main Entry Point
# -----------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="GPU Throughput Benchmark",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--mode",
        choices=["cnn", "transformer", "both"],
        default="both",
        help="Benchmark mode: cnn (MNIST), transformer (synthetic LM), or both",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=256,
        help="Batch size for CNN training",
    )
    parser.add_argument(
        "--transformer-batch-size",
        type=int,
        default=32,
        help="Batch size for transformer training (smaller due to large vocab logits)",
    )
    parser.add_argument(
        "--warmup-batches",
        type=int,
        default=10,
        help="Number of warmup batches (not timed)",
    )
    parser.add_argument(
        "--benchmark-batches",
        type=int,
        default=100,
        help="Number of batches to benchmark",
    )
    parser.add_argument(
        "--precision",
        choices=["fp32", "fp16", "bf16", "tf32"],
        default="fp16",
        help="Precision mode: fp32 (full), fp16 (AMP), bf16 (AMP), tf32 (Ampere+)",
    )
    parser.add_argument(
        "--diagnose",
        action="store_true",
        help="Run CUDA/cuBLAS diagnostic tests before benchmarking",
    )
    parser.add_argument(
        "--cpu",
        action="store_true",
        help="Force CPU-only execution (for CPU vs GPU comparison)",
    )
    args = parser.parse_args()

    requested_precision = PrecisionMode(args.precision)
    device, actual_precision = print_system_info(requested_precision, force_cpu=args.cpu)

    # Run diagnostics if requested
    if args.diagnose:
        diag_results = run_cuda_diagnostics(device)
        # If FP16 tests fail, suggest using FP32
        fp16_tests = ["fp16_matmul", "fp16_autocast", "fp16_linear", "fp16_batched_matmul"]
        fp16_failures = [t for t in fp16_tests if t in diag_results and not diag_results[t]]
        if fp16_failures and actual_precision == PrecisionMode.FP16:
            print("WARNING: FP16 diagnostic tests failed. Switching to FP32.")
            actual_precision = PrecisionMode.FP32

    config = BenchmarkConfig(
        batch_size=args.batch_size,
        num_warmup_batches=args.warmup_batches,
        num_benchmark_batches=args.benchmark_batches,
        precision=actual_precision,
    )

    results: list[BenchmarkResult] = []

    # CNN Benchmark
    if args.mode in ("cnn", "both"):
        model = MNISTConvNet()
        param_count = sum(p.numel() for p in model.parameters())
        print(f"\nMNIST CNN parameters: {param_count:,}")

        loader = get_mnist_loader(config, device)

        try:
            result = run_benchmark(model, loader, config, device, "MNIST CNN", actual_precision, is_lm=False)
            results.append(result)
            print(result)
        except RuntimeError as e:
            if "CUBLAS" in str(e) or "cuBLAS" in str(e):
                print(f"\n*** cuBLAS error encountered with {actual_precision.value} ***")
                print(f"Error: {e}")
                print("\nRetrying with FP32 (no AMP)...")

                # Cleanup and retry with FP32
                del model
                if device.type == "cuda":
                    torch.cuda.empty_cache()

                model = MNISTConvNet()
                result = run_benchmark(model, loader, config, device, "MNIST CNN", PrecisionMode.FP32, is_lm=False)
                results.append(result)
                print(result)
            else:
                raise

        # Cleanup
        del model, loader
        if device.type == "cuda":
            torch.cuda.empty_cache()

    # Transformer Benchmark
    if args.mode in ("transformer", "both"):
        transformer_config = BenchmarkConfig(
            batch_size=args.transformer_batch_size,
            num_warmup_batches=args.warmup_batches,
            num_benchmark_batches=args.benchmark_batches,
            precision=actual_precision,
        )

        model = MiniGPT()
        param_count = sum(p.numel() for p in model.parameters())
        print(f"\nMiniGPT parameters: {param_count:,}")

        loader = get_synthetic_text_loader(transformer_config)

        try:
            result = run_benchmark(
                model, loader, transformer_config, device, "MiniGPT Transformer", actual_precision, is_lm=True
            )
            results.append(result)
            print(result)
        except RuntimeError as e:
            if "CUBLAS" in str(e) or "cuBLAS" in str(e):
                print(f"\n*** cuBLAS error encountered with {actual_precision.value} ***")
                print(f"Error: {e}")
                print("\nRetrying with FP32 (no AMP)...")

                # Cleanup and retry with FP32
                del model
                if device.type == "cuda":
                    torch.cuda.empty_cache()

                model = MiniGPT()
                result = run_benchmark(
                    model, loader, transformer_config, device, "MiniGPT Transformer", PrecisionMode.FP32, is_lm=True
                )
                results.append(result)
                print(result)
            else:
                raise

    # Summary
    if len(results) > 1:
        print("\n" + "=" * 60)
        print("BENCHMARK SUMMARY")
        print("=" * 60)
        for r in results:
            print(f"{r.model_name} ({r.precision_mode}):")
            print(f"  {r.samples_per_sec:,.1f} samples/sec | {r.peak_memory_mb:.0f}MB peak")
        print("=" * 60)


if __name__ == "__main__":
    main()
