#!/usr/bin/env python3
"""调试: 检查模型原始输出"""
import numpy as np
import cv2
import acl

ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2

img = cv2.imread("/home/强光车灯误报.png")
print("Image shape:", img.shape)

acl.init()
acl.rt.set_device(0)
ctx, _ = acl.rt.create_context(0)

model_id, _ = acl.mdl.load_from_file("/root/AI-tianyan/model/model.om")
desc = acl.mdl.create_desc()
acl.mdl.get_desc(desc, model_id)

in_sz = acl.mdl.get_input_size_by_index(desc, 0)
print("Input size:", in_sz, "bytes")

# Preprocess
h0, w0 = img.shape[:2]
inp_h, inp_w = 640, 640
scale = min(inp_h / h0, inp_w / w0)
nh, nw = int(h0 * scale), int(w0 * scale)
resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
canvas = np.full((inp_h, inp_w, 3), 114, dtype=np.uint8)
pad_top = (inp_h - nh) // 2
pad_left = (inp_w - nw) // 2
canvas[pad_top:pad_top + nh, pad_left:pad_left + nw] = resized
rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
blob = rgb.astype(np.float32) / 255.0
blob = np.ascontiguousarray(blob.transpose(2, 0, 1)[np.newaxis])
print("Blob shape:", blob.shape, "dtype:", blob.dtype)

# Allocate device input
in_ds = acl.mdl.create_dataset()
host_buf = np.ascontiguousarray(blob)
host_addr = acl.util.bytes_to_ptr(host_buf.tobytes())
in_buf, _ = acl.rt.malloc(in_sz, 0)
acl.rt.memcpy(in_buf, in_sz, host_addr, in_sz, ACL_MEMCPY_HOST_TO_DEVICE)
acl.mdl.add_dataset_buffer(in_ds, acl.create_data_buffer(in_buf, in_sz))

# Allocate output
out_num = acl.mdl.get_num_outputs(desc)
out_sizes = [acl.mdl.get_output_size_by_index(desc, i) for i in range(out_num)]
out_ds = acl.mdl.create_dataset()
out_bufs = []
for i in range(out_num):
    buf, _ = acl.rt.malloc(out_sizes[i], 0)
    out_bufs.append(buf)
    acl.mdl.add_dataset_buffer(out_ds, acl.create_data_buffer(buf, out_sizes[i]))
    print(f"Output[{i}] size: {out_sizes[i]} bytes")

# Execute
ret = acl.mdl.execute(model_id, in_ds, out_ds)
print("Execute ret:", ret)

# Get output
for i in range(out_num):
    sz = out_sizes[i]
    host = np.zeros(sz, dtype=np.uint8)
    host_addr = acl.util.bytes_to_ptr(host.tobytes())
    acl.rt.memcpy(host_addr, sz, out_bufs[i], sz, ACL_MEMCPY_DEVICE_TO_HOST)
    
    dims_out, _ = acl.mdl.get_output_dims(desc, i)
    d = dims_out['dims']
    print(f"\nOutput[{i}] dims: {d}")
    print(f"  Size: {sz} bytes")
    
    # Try FP16
    fp16 = host.view(np.float16)
    print(f"  FP16 shape: {fp16.shape}")
    
    reshaped = fp16.astype(np.float32).reshape(d)
    print(f"  Reshaped: {reshaped.shape}")
    print(f"  Min: {reshaped.min():.4f}, Max: {reshaped.max():.4f}, Mean: {reshaped.mean():.4f}")
    
    # Check if it's all zeros or NaN
    nan_count = np.isnan(reshaped).sum()
    zero_count = (reshaped == 0).sum()
    print(f"  NaN count: {nan_count}, Zero count: {zero_count}/{reshaped.size}")
    
    # For YOLOv8 output [1, 84, 8400], check class scores
    # boxes: reshaped[:, :4, :]
    # scores: reshaped[:, 4:, :]
    scores = reshaped[0, 4:, :]
    max_scores = scores.max(axis=0)  # max class score per anchor
    print(f"\n  Max class scores per anchor:")
    print(f"    Min: {max_scores.min():.4f}, Max: {max_scores.max():.4f}, Mean: {max_scores.mean():.4f}")
    
    # Count anchors with score > 0.1
    high_conf = (max_scores > 0.1).sum()
    print(f"    Anchors with conf > 0.1: {high_conf}")
    print(f"    Anchors with conf > 0.01: {(max_scores > 0.01).sum()}")
    
    # Print top 5
    top5_idx = np.argsort(max_scores)[-5:][::-1]
    for idx in top5_idx:
        best_cls = scores[:, idx].argmax()
        print(f"    Anchor {idx}: class={best_cls} (score={max_scores[idx]:.4f})")
        box = reshaped[0, :4, idx]
        print(f"      box: {box}")

# Cleanup
for buf in out_bufs:
    acl.rt.free(buf)
acl.rt.free(in_buf)
acl.mdl.destroy_dataset(in_ds)
acl.mdl.destroy_dataset(out_ds)
acl.mdl.unload(model_id)
acl.mdl.destroy_desc(desc)
acl.rt.destroy_context(ctx)
acl.rt.reset_device(0)
acl.finalize()
print("\nDone")