import torch
import time
import os
import sys

print(f'Device: {torch.cuda.get_device_name(0)}')
print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
print()

def timed_op(name, fn, timeout=120):
    """Run operation with timeout detection"""
    print(f'=== {name} ===', flush=True)
    try:
        torch.cuda.synchronize()
        t = time.time()
        result = fn()
        torch.cuda.synchronize()
        elapsed = time.time() - t
        if elapsed > timeout:
            print(f'  SLOW: {elapsed:.1f}s (>{timeout}s)', flush=True)
        else:
            print(f'  OK: {elapsed:.3f}s', flush=True)
        return True
    except Exception as e:
        print(f'  FAIL: {e}', flush=True)
        return False

# Lumina2-like dimensions
# Hidden dim ~3072, heads ~24, head_dim ~128
# Latent 64x64 = 4096 tokens for 512x512 image

# Test 1: QKV projection (Linear 3072 -> 3*3072)
def test_qkv():
    lin = torch.nn.Linear(3072, 9216).cuda().float()  # 3*3072
    x = torch.randn(1, 4096, 3072, device='cuda', dtype=torch.float32)
    y = lin(x)
    del lin, x, y
timed_op('QKV projection (1x4096x3072 -> 9216)', test_qkv)

# Test 2: Attention scores (24 heads, 4096x4096)
def test_attn_scores():
    q = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
    k = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
    scores = torch.matmul(q, k.transpose(-2, -1))  # 24x4096x4096 = 1.5GB fp32!
    del q, k, scores
timed_op('Attention scores 24x4096x4096', test_attn_scores)

# Test 2b: Smaller attention (fewer heads)
def test_attn_small():
    q = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    k = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    scores = torch.matmul(q, k.transpose(-2, -1))
    del q, k, scores
timed_op('Attention scores 8x1024x1024 (smaller)', test_attn_small)

# Test 3: Attention value multiply
def test_attn_values():
    attn = torch.randn(1, 8, 1024, 1024, device='cuda', dtype=torch.float32)
    v = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
    out = torch.matmul(attn, v)
    del attn, v, out
timed_op('Attn * values 8x1024x1024 @ 8x1024x128', test_attn_values)

# Test 4: FFN (3072 -> 12288 -> 3072)
def test_ffn():
    up = torch.nn.Linear(3072, 12288).cuda().float()
    down = torch.nn.Linear(12288, 3072).cuda().float()
    x = torch.randn(1, 4096, 3072, device='cuda', dtype=torch.float32)
    h = torch.nn.functional.gelu(up(x))
    y = down(h)
    del up, down, x, h, y
timed_op('FFN 3072->12288->3072 (4096 tokens)', test_ffn)

# Test 5: Conv2d 128ch (VAE-like)
def test_vae_conv():
    conv = torch.nn.Conv2d(128, 128, 3, padding=1).cuda().float()
    x = torch.randn(1, 128, 64, 64, device='cuda', dtype=torch.float32)
    y = conv(x)
    del conv, x, y
timed_op('VAE Conv2d 128ch 64x64', test_vae_conv)

# Test 6: Memory pressure test
def test_mem():
    tensors = []
    total = 0
    for i in range(10):
        t = torch.randn(256, 256, 256, device='cuda', dtype=torch.float32)  # 64MB each
        tensors.append(t)
        total += 64
    print(f'  Allocated {total}MB on GPU', flush=True)
    for t in tensors:
        del t
    del tensors
    torch.cuda.empty_cache()
timed_op('Memory: allocate 640MB', test_mem)

print()
print('=== ALL LUMINA2 TESTS DONE ===', flush=True)
torch.cuda.synchronize()
os._exit(0)