import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import perf_counter
import blaze as bl
from conftest import COMPILE_BACKENDS


def test_performance():
    """Test that the performance is reasonable compared to regular PyTorch."""
    n_layers = 100
    n_repeats = 1000

    def forward(x):
        for _ in range(n_layers):
            x = bl.Linear(10, 10)(x)
            x = F.relu(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))
    model = torch.jit.trace(model, torch.randn(2, 10))

    start_time = perf_counter()
    for _ in range(n_repeats):
        _ = model(torch.randn(2, 10))
    elapsed_blaze = perf_counter() - start_time

    # Compare to regular PyTorch
    class RegularModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.linears = nn.ModuleList([nn.Linear(10, 10) for _ in range(n_layers)])

        def forward(self, x):
            for linear in self.linears:
                x = linear(x)
                x = F.relu(x)
            return x

    regular_model = RegularModel()
    regular_model = torch.jit.trace(regular_model, torch.randn(2, 10))

    start_time = perf_counter()
    for _ in range(n_repeats):
        regular_model(torch.randn(2, 10))
    elapsed_regular = perf_counter() - start_time

    print(f"Blaze elapsed: {elapsed_blaze:.4f} seconds")
    print(f"Regular elapsed: {elapsed_regular:.4f} seconds")

    assert elapsed_blaze < 1.05 * elapsed_regular, "Blaze is too slow compared to regular PyTorch"


@pytest.mark.parametrize("backend", COMPILE_BACKENDS)
def test_performance_compile(backend):
    """torch.compile should not be slower than eager blaze."""
    n_layers = 100
    n_repeats = 1000

    def forward(x):
        for _ in range(n_layers):
            x = bl.Linear(10, 10)(x)
            x = F.relu(x)
        return x

    model = bl.transform(forward)
    model.init(torch.randn(2, 10))

    # Benchmark eager (uncompiled)
    for _ in range(3):
        model(torch.randn(2, 10))

    start_time = perf_counter()
    for _ in range(n_repeats):
        model(torch.randn(2, 10))
    elapsed_eager = perf_counter() - start_time

    # Benchmark compiled
    compiled = torch.compile(model, backend=backend)

    for _ in range(3):
        compiled(torch.randn(2, 10))

    start_time = perf_counter()
    for _ in range(n_repeats):
        compiled(torch.randn(2, 10))
    elapsed_compiled = perf_counter() - start_time

    print(f"\n[{backend}] Eager: {elapsed_eager:.4f}s, Compiled: {elapsed_compiled:.4f}s, "
          f"Speedup: {elapsed_eager/elapsed_compiled:.2f}x")

    assert elapsed_compiled <= elapsed_eager * 1.10, (
        f"[{backend}] Compiled slower than eager: {elapsed_compiled:.4f}s vs {elapsed_eager:.4f}s "
        f"({elapsed_compiled/elapsed_eager:.2f}x)")
