#!/usr/bin/env python3
"""
昇腾 310B4 NPU 推理性能基准测试
测试项目：
1. 确认 NPU 推理（非 CPU fallback）
2. 单次推理延迟
3. 连续推理吞吐 (FPS)
4. 多 worker 并发
5. 不同分辨率影响
"""
import os
import sys
import time
import socket
import struct
import json
import base64
import threading
import numpy as np
import cv2
from concurrent.futures import ThreadPoolExecutor

SOCK_PATH = "/tmp/edge-infer.sock"

def make_test_image(w=640, h=640):
    """生成随机测试图片"""
    img = np.random.randint(0, 255, (h, w, 3), dtype=np.uint8)
    _, buf = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 90])
    return buf.tobytes()

def infer_one(jpeg_bytes):
    """发送一次推理请求，返回响应时间(ms)"""
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect(SOCK_PATH)
    
    msg = {
        "stream_id": 0,
        "device_id": "bench",
        "url": "bench://test",
        "ts": 1234567890.0,
        "jpeg_b64": base64.b64encode(jpeg_bytes).decode("utf-8")
    }
    
    data = json.dumps(msg).encode("utf-8")
    
    t0 = time.time()
    sock.sendall(struct.pack(">I", len(data)) + data)
    hdr = sock.recv(4)
    length = struct.unpack(">I", hdr)[0]
    result = json.loads(sock.recv(length).decode("utf-8"))
    elapsed = (time.time() - t0) * 1000  # ms
    
    sock.close()
    return elapsed, result

def test_single_infer():
    """单次推理延迟"""
    print("\n=== 单次推理延迟测试 ===")
    jpeg = make_test_image()
    times = []
    for i in range(5):
        t, _ = infer_one(jpeg)
        times.append(t)
        print(f"  第 {i+1} 次: {t:.1f} ms")
    
    # 跳过第一次预热
    times = times[1:]
    avg = sum(times) / len(times)
    print(f"\n  平均延迟 (排除预热): {avg:.1f} ms")
    print(f"  理论 FPS: {1000/avg:.1f}")
    return avg

def test_throughput():
    """连续推理吞吐"""
    print("\n=== 连续推理吞吐测试 ===")
    jpeg = make_test_image()
    count = 30
    t0 = time.time()
    
    for i in range(count):
        _, _ = infer_one(jpeg)
        if (i+1) % 10 == 0:
            elapsed = time.time() - t0
            print(f"  {i+1}/{count} 完成, 累计: {elapsed:.1f}s, FPS: {(i+1)/elapsed:.1f}")
    
    total = time.time() - t0
    fps = count / total
    print(f"\n  总计 {count} 帧: {total:.2f}s")
    print(f"  吞吐: {fps:.1f} FPS")
    print(f"  每帧延迟: {total/count*1000:.1f} ms")
    return fps

def test_concurrent(workers=4):
    """多 worker 并发推理"""
    print(f"\n=== {workers} 路并发测试 ===")
    jpeg = make_test_image()
    results_per_worker = []
    
    def worker_task(worker_id, n_frames=20):
        times = []
        for _ in range(n_frames):
            t, _ = infer_one(jpeg)
            times.append(t)
        return worker_id, times
    
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = [pool.submit(worker_task, i, 15) for i in range(workers)]
        for f in futures:
            wid, times = f.result()
            results_per_worker.append((wid, sum(times)/len(times), max(times), min(times)))
    
    total = time.time() - t0
    total_frames = workers * 15
    total_fps = total_frames / total
    
    for wid, avg, mx, mn in results_per_worker:
        print(f"  Worker {wid}: avg={avg:.1f}ms, max={mx:.1f}ms, min={mn:.1f}ms")
    
    print(f"\n  并发 {workers} 路: 总计 {total_frames} 帧, {total:.2f}s")
    print(f"  总吞吐: {total_fps:.1f} FPS")
    print(f"  单路等效 FPS: {total_fps/workers:.1f}")
    return total_fps

def test_power():
    """读取 NPU 功耗信息"""
    import subprocess
    try:
        result = subprocess.run(['npu-smi', 'info'], capture_output=True, text=True, timeout=5)
        print("\n=== NPU 状态 ===")
        print(result.stdout)
    except Exception as e:
        print(f"无法读取 NPU 状态: {e}")

if __name__ == "__main__":
    print("=" * 60)
    print("  昇腾 310B4 + ACL 原生推理 性能基准测试")
    print("=" * 60)
    
    test_power()
    test_single_infer()
    test_throughput()
    test_concurrent(2)
    test_concurrent(4)
    test_concurrent(6)
    
    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)