import math
from functools import partial, wraps
import os
import torch
import timeit
import statistics
import torch_xla
import torch_xla.core.xla_model as xm
from torch_xla.amp import autocast


def wrap_with_mark_step(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        res = func(*args, **kwargs)
        xm.mark_step()
        return res
    return wrapper


@wrap_with_mark_step
def sdpa(query, key, value):
    return torch.nn.functional.scaled_dot_product_attention(query, key, value)


@wrap_with_mark_step
def standard_attention(query, key, value):
    # from sdpa's doc implementation (it's pretty straightforward)
    scale_factor = 1 / math.sqrt(query.size(-1))
    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight = torch.softmax(attn_weight, dim=-1)
    return attn_weight @ value


def time_attention(attention_fn, query, key, value):
    # Warm-up
    for _ in range(100):
        attention_fn(query, key, value)

    # Run timings with repeats
    repeats = 10
    number = 100
    times = timeit.repeat(partial(attention_fn, query, key, value), repeat=repeats, number=number)

    mean_time = statistics.mean(times) / number * 1000  # ms per call
    std_dev = statistics.stdev(times) / number * 1000  # ms per call

    print(f"    {attention_fn.__name__} average time: {mean_time:.3f} ms per call, dev: {std_dev:.3f} ms per call")


if __name__ == '__main__':
    os.environ["XLA_DISABLE_FUNCTIONALIZATION"] = "1"

    batch_size = 32
    seq_len = 256
    num_heads = 16
    head_dim = 64
    device = xm.xla_device()

    # Generate input tensors with correct shape: (batch_size, num_heads, seq_len, head_dim)
    query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)
    key = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)
    value = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device)

    assert torch.allclose(sdpa(query, key, value), standard_attention(query, key, value),
                          rtol=1e-02, atol=1e-02)
    print("Both attention functions produce about the same output, test time performance now")

    print("No mixed precision:")
    time_attention(sdpa, query, key, value)
    time_attention(standard_attention, query, key, value)

    print("With mixed precision:")
    with autocast(xm.xla_device(), enabled=True, dtype=torch.bfloat16):
        time_attention(sdpa, query, key, value)
        time_attention(standard_attention, query, key, value)

    print(f"torch version: {torch.__version__}, torch_xla version: {torch_xla.__version__}")