import time
import torch
import torch.nn as nn
import torch.optim as optim

# Конфигурация параметров нагрузочного тестирования
NUM_SAMPLES = 300_000
DIM_IN = 4096
DIM_OUT = 10
BATCH_SIZE = 16_384
NUM_STEPS = 1000

def run_gpu_benchmark():
    # Проверка доступности аппаратного ускорения
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Инициализация стресс-теста на устройстве: {device}")

    # Генерация синтетического датасета для аллокации VRAM
    x_data = torch.randn(NUM_SAMPLES, DIM_IN, device=device, dtype=torch.float32)
    y_data = torch.randn(NUM_SAMPLES, DIM_OUT, device=device, dtype=torch.float32)

    # Архитектура тестовой полносвязной сети
    model = nn.Sequential(
        nn.Linear(DIM_IN, 2048),
        nn.ReLU(),
        nn.Linear(2048, 1024),
        nn.ReLU(),
        nn.Linear(1024, 512),
        nn.ReLU(),
        nn.Linear(512, DIM_OUT)
    ).to(device)

    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    print("Начало прогрева GPU и симуляции цикла обучения...")
    start_time = time.time()

    for step in range(NUM_STEPS):
        # Сэмплирование случайного батча
        idx = torch.randint(0, NUM_SAMPLES, (BATCH_SIZE,), device=device)
        x_batch = x_data[idx]
        y_batch = y_data[idx]

        optimizer.zero_grad()
        predictions = model(x_batch)
        loss = loss_fn(predictions, y_batch)
        
        loss.backward()
        optimizer.step()

        # Логирование статуса (каждые 100 итераций для снижения I/O overhead)
        if step % 100 == 0:
            print(f"Итерация {step}/{NUM_STEPS} | Текущий loss: {loss.item():.4f}")

    end_time = time.time()
    print(f"Стресс-тест завершен. Общее время: {end_time - start_time:.2f} сек.")

if __name__ == "__main__":
    run_gpu_benchmark()