#!/usr/bin/env python3
"""
测试 TensorRT 引擎加载和推理
使用随机图像测试
"""

import numpy as np
import time
from ultralytics import YOLO


def test_tensorrt_engine():
    """测试 TensorRT 引擎"""
    print("TensorRT 引擎测试")
    print("=" * 60)
    
    engine_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.engine"
    
    # 1. 加载引擎
    print(f"🚀 加载 TensorRT 引擎: {engine_path}")
    try:
        model = YOLO(engine_path, task='detect')
        print("✅ 引擎加载成功")
    except Exception as e:
        print(f"❌ 引擎加载失败: {e}")
        return
    
    # 2. 测试单帧推理
    print(f"\n🔄 测试单帧推理...")
    test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
    
    try:
        start_time = time.time()
        results = model.predict(
            test_image,
            imgsz=640,
            conf=0.45,
            verbose=False,
            device=0
        )
        elapsed = (time.time() - start_time) * 1000
        print(f"✅ 单帧推理成功，耗时: {elapsed:.1f}ms")
    except Exception as e:
        print(f"❌ 单帧推理失败: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # 3. 测试批量推理
    print(f"\n🔄 测试批量推理（batch=4）...")
    test_images = [np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8) for _ in range(4)]
    
    try:
        start_time = time.time()
        results = model.predict(
            test_images,
            imgsz=640,
            conf=0.45,
            verbose=False,
            device=0
        )
        elapsed = (time.time() - start_time) * 1000
        per_frame = elapsed / 4
        print(f"✅ 批量推理成功，总耗时: {elapsed:.1f}ms，每帧: {per_frame:.1f}ms")
    except Exception as e:
        print(f"❌ 批量推理失败: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # 4. 性能测试
    print(f"\n🔄 性能测试（100帧）...")
    inference_times = []
    
    try:
        for i in range(100):
            test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
            
            start_time = time.time()
            results = model.predict(
                test_image,
                imgsz=640,
                conf=0.45,
                verbose=False,
                device=0
            )
            elapsed = (time.time() - start_time) * 1000
            inference_times.append(elapsed)
            
            if (i + 1) % 20 == 0:
                print(f"  进度: {i+1}/100")
        
        print(f"\n{'='*60}")
        print("性能统计")
        print(f"{'='*60}")
        print(f"平均推理延迟: {np.mean(inference_times):.1f}ms")
        print(f"P50推理延迟: {np.percentile(inference_times, 50):.1f}ms")
        print(f"P95推理延迟: {np.percentile(inference_times, 95):.1f}ms")
        print(f"P99推理延迟: {np.percentile(inference_times, 99):.1f}ms")
        print(f"最小延迟: {np.min(inference_times):.1f}ms")
        print(f"最大延迟: {np.max(inference_times):.1f}ms")
        print(f"{'='*60}")
        
        print("\n✅ 所有测试通过！TensorRT 引擎工作正常")
    
    except Exception as e:
        print(f"❌ 性能测试失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    try:
        test_tensorrt_engine()
    except Exception as e:
        print(f"\n❌ 测试异常: {e}")
        import traceback
        traceback.print_exc()