import torch
import torch.nn as nn
import time

# Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1024, 2048)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(2048, 1024)
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# Prepare model and dummy data
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleNet().to(device)
dummy_input = torch.randn(128, 1024).to(device)
dummy_target = torch.randn(128, 1024).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
num_iterations = 50

# --- Benchmark without torch.compile ---
print(f"--- Running without torch.compile on {device} ---")
start_time = time.time()
for _ in range(num_iterations):
    optimizer.zero_grad()
    output = model(dummy_input)
    loss = criterion(output, dummy_target)
    loss.backward()
    optimizer.step()
if device == "cuda":
    torch.cuda.synchronize() # Wait for GPU ops to complete
time_uncompiled = time.time() - start_time
print(f"Time without compile: {time_uncompiled:.4f} seconds\n")

# --- Benchmark with torch.compile ---
# Apply torch.compile to the model. This happens once upfront.
# The default backend 'inductor' is typically the best performing.
compiled_model = torch.compile(model)
# Ensure optimizer is correctly set up for the compiled model's parameters
# (in this case, `compiled_model` shares parameters with `model`, so no re-init needed if parameters are the same object)

print(f"--- Running with torch.compile on {device} ---")
start_time = time.time()
for _ in range(num_iterations):
    optimizer.zero_grad()
    output = compiled_model(dummy_input) # Use the compiled model
    loss = criterion(output, dummy_target)
    loss.backward()
    optimizer.step()
if device == "cuda":
    torch.cuda.synchronize() # Wait for GPU ops to complete
time_compiled = time.time() - start_time
print(f"Time with compile: {time_compiled:.4f} seconds")

if time_uncompiled > 0:
    print(f"\nSpeedup: {time_uncompiled / time_compiled:.2f}x")