#!/usr/bin/env python3
"""
BC-250 GGUF Load Test - isolated from ComfyUI
Loads z_image_turbo GGUF and times each step.
"""
import os, sys, time

os.environ["HSA_OVERRIDE_GFX_VERSION"] = "10.1.0"
os.environ["HSA_ENABLE_SDMA"] = "0"
os.environ["HIP_VISIBLE_DEVICES"] = "0"

print("[TEST] Starting GGUF load test...", flush=True)

t0 = time.time()
import torch
print(f"[TEST] torch imported in {time.time()-t0:.1f}s", flush=True)

# Add ComfyUI paths
sys.path.insert(0, "/home/fabian/ComfyUI")
sys.path.insert(0, "/home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF")

t1 = time.time()
import gguf
print(f"[TEST] gguf imported in {time.time()-t1:.1f}s", flush=True)

GGUF_PATH = "/home/fabian/ComfyUI/models/unet/z_image_turbo-Q5_K_S.gguf"

# Step 1: Open and read the file
print(f"\n[TEST] Step 1: Reading GGUF file...", flush=True)
t2 = time.time()
reader = gguf.GGUFReader(GGUF_PATH)
print(f"[TEST] Reader created in {time.time()-t2:.1f}s", flush=True)
print(f"[TEST] Tensors: {len(reader.tensors)}", flush=True)

# Check architecture
arch = None
field = reader.get_field("general.architecture")
if field:
    arch = str(field.parts[field.data[-1]], "utf-8")
print(f"[TEST] Architecture: {arch}", flush=True)

# Step 2: List first few tensors with types
print(f"\n[TEST] Step 2: Tensor info (first 10)...", flush=True)
qtype_counts = {}
for i, tensor in enumerate(reader.tensors):
    ttype = tensor.tensor_type
    tname = getattr(ttype, 'name', repr(ttype))
    qtype_counts[tname] = qtype_counts.get(tname, 0) + 1
    if i < 10:
        print(f"  {tensor.name}: type={tname} shape={list(reversed(tensor.shape))} size={tensor.data.nbytes}", flush=True)

print(f"\n[TEST] Quant type distribution:", flush=True)
for k, v in sorted(qtype_counts.items()):
    print(f"  {k}: {v}", flush=True)

# Step 3: Try loading one tensor to CPU
print(f"\n[TEST] Step 3: Load first Q5_K tensor to CPU...", flush=True)
import warnings
from dequant import dequantize_tensor, is_quantized

for tensor in reader.tensors:
    if tensor.tensor_type == gguf.GGMLQuantizationType.Q5_K:
        name = tensor.name
        print(f"  Tensor: {name}", flush=True)
        print(f"  Raw shape: {list(reversed(tensor.shape))}", flush=True)
        print(f"  Raw bytes: {tensor.data.nbytes}", flush=True)
        
        t3 = time.time()
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
            torch_tensor = torch.from_numpy(tensor.data)
        
        # Try to get orig shape
        field_key = f"comfy.gguf.orig_shape.{name}"
        field = reader.get_field(field_key)
        if field:
            shape = torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
        else:
            shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
        
        print(f"  Original shape: {shape}", flush=True)
        print(f"  Tensor dtype: {torch_tensor.dtype}, device: {torch_tensor.device}", flush=True)
        
        # Create a GGMLTensor-like wrapper
        torch_tensor.tensor_type = tensor.tensor_type
        torch_tensor.tensor_shape = shape
        torch_tensor.patches = []
        
        # Dequantize on CPU
        print(f"  Dequantizing on CPU...", flush=True)
        t4 = time.time()
        result = dequantize_tensor(torch_tensor, dtype=torch.float16)
        dt = time.time() - t4
        print(f"  Dequantized in {dt:.3f}s -> shape={result.shape} dtype={result.dtype}", flush=True)
        
        # Try moving to GPU
        print(f"  Moving to GPU...", flush=True)
        t5 = time.time()
        gpu_result = result.to("cuda:0")
        torch.cuda.synchronize()
        dt2 = time.time() - t5
        print(f"  GPU transfer in {dt2:.3f}s", flush=True)
        print(f"  GPU tensor: shape={gpu_result.shape} dtype={gpu_result.dtype} device={gpu_result.device}", flush=True)
        
        del gpu_result, result, torch_tensor
        torch.cuda.empty_cache()
        break

# Step 4: Try dequantizing ALL tensors on CPU
print(f"\n[TEST] Step 4: Dequantize all Q5_K tensors on CPU...", flush=True)
t_all = time.time()
count = 0
for tensor in reader.tensors:
    if tensor.tensor_type == gguf.GGMLQuantizationType.Q5_K:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="The given NumPy array is not writable")
            tt = torch.from_numpy(tensor.data)
        
        field_key = f"comfy.gguf.orig_shape.{tensor.name}"
        field = reader.get_field(field_key)
        if field:
            shape = torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data))
        else:
            shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape)))
        
        tt.tensor_type = tensor.tensor_type
        tt.tensor_shape = shape
        tt.patches = []
        
        result = dequantize_tensor(tt, dtype=torch.float16)
        count += 1
        if count % 20 == 0:
            print(f"  Dequantized {count} tensors ({time.time()-t_all:.1f}s)...", flush=True)

dt_all = time.time() - t_all
print(f"\n[TEST] Dequantized {count} Q5_K tensors in {dt_all:.1f}s total", flush=True)
print(f"[TEST] Average: {dt_all/max(count,1):.3f}s per tensor", flush=True)

print(f"\n[TEST] ALL DONE in {time.time()-t0:.1f}s", flush=True)
os._exit(0)