"""Test F.linear with exact shapes from Z-Image-Turbo model on GPU.
This isolates whether the GPU matmul hangs for model-specific shapes."""
import torch
import torch.nn.functional as F
import time
import os

os.environ['HSA_OVERRIDE_GFX_VERSION'] = '10.1.0'
os.environ['HSA_ENABLE_SDMA'] = '0'
os.environ['TORCHDYNAMO_DISABLE'] = '1'

print(f"Device: {torch.cuda.get_device_name(0)}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB" if hasattr(torch.cuda.get_device_properties(0), 'total_mem') else "")

# Weight shapes from the model (logged as CUDA#505-510)
shapes = [
    ("CBW505", (1024, 256)),
    ("CBW506", (256, 1024)),
    ("CBW507", (3840, 2560)),
    ("CBW508", (3840, 64)),
    ("CBW509", (11520, 3840)),
    ("CBW510", (3840, 3840)),
]

# Batch size for Lumina2 at 512x512: typically (1, seq_len, dim)
# Sequence length ~ 1024 patches
batch_seq = 1024

print(f"\n=== Test 1: Individual F.linear with model shapes ===")
for name, (out_feat, in_feat) in shapes:
    print(f"\n{name}: Linear({in_feat} -> {out_feat})")
    try:
        # Create tensors
        x = torch.randn(1, batch_seq, in_feat, dtype=torch.float16, device='cuda:0')
        w = torch.randn(out_feat, in_feat, dtype=torch.float16, device='cuda:0')
        b = torch.randn(out_feat, dtype=torch.float16, device='cuda:0')
        torch.cuda.synchronize()
        print(f"  Tensors created on GPU", flush=True)
        
        # F.linear = x @ w.T + b
        t0 = time.time()
        out = F.linear(x, w, b)
        torch.cuda.synchronize()
        dt = time.time() - t0
        print(f"  F.linear done: {dt*1000:.1f}ms, out={list(out.shape)}", flush=True)
        
        del x, w, b, out
    except Exception as e:
        print(f"  ERROR: {e}", flush=True)

print(f"\n=== Test 2: Sequential F.linear chain (simulating forward pass) ===")
# Simulate the model's sequence: transfer from CPU, compute, repeat
x = torch.randn(1, batch_seq, 256, dtype=torch.float16, device='cuda:0')
torch.cuda.synchronize()
t0 = time.time()

for i, (name, (out_feat, in_feat)) in enumerate(shapes):
    # Simulate: CPU dequant → transfer → F.linear
    w_cpu = torch.randn(out_feat, in_feat, dtype=torch.float16)
    w_gpu = w_cpu.to('cuda:0', non_blocking=False)
    
    # Adjust input dimension if needed
    if x.shape[-1] != in_feat:
        x = torch.randn(1, batch_seq, in_feat, dtype=torch.float16, device='cuda:0')
        torch.cuda.synchronize()
    
    out = F.linear(x, w_gpu)
    torch.cuda.synchronize()
    dt = time.time() - t0
    print(f"  {name}: transfer+linear done, dt_total={dt:.2f}s, out={list(out.shape)}", flush=True)
    x = out
    del w_cpu, w_gpu

print(f"\nTotal chain time: {time.time()-t0:.2f}s")

print(f"\n=== Test 3: Manual softmax on GPU ===")
x = torch.randn(1, 16, 1024, 1024, dtype=torch.float16, device='cuda:0')
torch.cuda.synchronize()
t0 = time.time()
x_max = x.max(dim=-1, keepdim=True).values
exp_x = torch.exp(x - x_max)
result = exp_x / exp_x.sum(dim=-1, keepdim=True)
torch.cuda.synchronize()
print(f"  Manual softmax: {(time.time()-t0)*1000:.1f}ms, shape={list(result.shape)}")

print(f"\n=== Test 4: Full manual SDPA on GPU ===")
Q = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
K = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
V = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
torch.cuda.synchronize()
t0 = time.time()
scale = 64 ** -0.5
attn = torch.matmul(Q, K.transpose(-2, -1)) * scale
torch.cuda.synchronize()
print(f"  Q@K^T: {(time.time()-t0)*1000:.1f}ms", flush=True)
t1 = time.time()
attn_max = attn.max(dim=-1, keepdim=True).values
exp_attn = torch.exp(attn - attn_max)
attn_weights = exp_attn / exp_attn.sum(dim=-1, keepdim=True)
torch.cuda.synchronize()
print(f"  Softmax: {(time.time()-t1)*1000:.1f}ms", flush=True)
t2 = time.time()
out = torch.matmul(attn_weights, V)
torch.cuda.synchronize()
print(f"  Attn@V: {(time.time()-t2)*1000:.1f}ms", flush=True)
print(f"  Total SDPA: {(time.time()-t0)*1000:.1f}ms, out={list(out.shape)}")

print("\n=== ALL TESTS PASSED ===")