import time
import torch
import torch.nn as nn
from torch_surgeon import Surgeon


def make_model():
    return nn.Sequential(
        nn.Linear(128, 64), nn.ReLU(),
        nn.Linear(64, 32),  nn.ReLU(),
        nn.Linear(32, 1),
    )


def time_loop(model, n=100):
    opt = torch.optim.SGD(model.parameters(), lr=0.01)
    crit = nn.MSELoss()
    start = time.perf_counter()
    for _ in range(n):
        x = torch.randn(32, 128)
        y = torch.randn(32, 1)
        loss = crit(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return time.perf_counter() - start


def test_overhead_under_2_percent():
    """Surgeon overhead must be <2% on 100 training steps."""
    N = 100
    REPS = 3

    # Baseline (no surgeon)
    base_times = []
    for _ in range(REPS):
        m = make_model()
        base_times.append(time_loop(m, N))
    base = min(base_times)

    # With surgeon
    surg_times = []
    for _ in range(REPS):
        m = make_model()
        s = Surgeon(m, rules={'verbose': False, 'log_every': 1000}).attach()
        surg_times.append(time_loop(m, N))
        s.detach()
    surg = min(surg_times)

    overhead_pct = (surg - base) / base * 100
    print('
Overhead: ' + format(overhead_pct, '.2f') + '% (base=' + format(base, '.3f') + 's, with_surgeon=' + format(surg, '.3f') + 's)')
    assert overhead_pct < 2.0, 'Overhead ' + format(overhead_pct, '.2f') + '% exceeds 2% threshold'