#!/usr/bin/env python3
"""
简化版动态批次 TensorRT 引擎导出脚本
支持 batch size: 1, 2, 4, 8, 16, 32
"""

import os
import torch
from ultralytics import YOLO
import time

def export_dynamic_tensorrt_engine():
    """导出支持动态批次的 TensorRT 引擎"""
    
    print("🚀 开始导出动态批次 TensorRT 引擎")
    print("=" * 60)
    
    # 检查 CUDA 可用性
    if not torch.cuda.is_available():
        print("❌ CUDA 不可用，无法导出 TensorRT 引擎")
        return None
    
    print(f"✅ CUDA 可用，设备: {torch.cuda.get_device_name(0)}")
    
    # 模型路径
    model_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.pt"
    
    if not os.path.exists(model_path):
        print(f"❌ 模型文件不存在: {model_path}")
        return None
    
    print(f"📁 模型路径: {model_path}")
    
    try:
        # 加载模型
        print("\n📦 加载 YOLO 模型...")
        model = YOLO(model_path)
        
        # 导出动态 TensorRT 引擎
        print("\n🔧 导出动态 TensorRT 引擎...")
        print("配置参数:")
        print("  - 格式: TensorRT Engine")
        print("  - 输入尺寸: 640x640")
        print("  - 精度: FP16")
        print("  - 动态批次: 支持")
        print("  - 工作空间: 8GB")
        print("  - 设备: GPU")
        
        start_time = time.time()
        
        # 导出参数 - 使用正确的动态配置
        export_args = {
            'format': 'engine',          # TensorRT engine format
            'imgsz': 640,               # Input image size
            'device': 0,                # GPU device
            'half': True,               # FP16 precision
            'dynamic': True,            # Enable dynamic shapes
            'simplify': True,           # Simplify ONNX model
            'workspace': 8,             # Workspace size in GB
            'verbose': True,            # Verbose output
        }
        
        print(f"\n⏳ 开始导出（预计需要 5-10 分钟）...")
        
        # 执行导出
        exported_model = model.export(**export_args)
        
        export_time = time.time() - start_time
        
        print(f"\n✅ TensorRT 引擎导出完成!")
        print(f"⏱️ 导出耗时: {export_time:.1f} 秒")
        print(f"📁 引擎文件: {exported_model}")
        
        # 检查文件大小
        if os.path.exists(exported_model):
            file_size = os.path.getsize(exported_model) / (1024 * 1024)  # MB
            print(f"📊 文件大小: {file_size:.1f} MB")
        
        return exported_model
        
    except Exception as e:
        print(f"\n❌ 导出失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def test_dynamic_engine(engine_path):
    """测试动态引擎的不同批次大小"""
    print(f"\n🧪 测试动态引擎: {engine_path}")
    
    if not os.path.exists(engine_path):
        print(f"❌ 引擎文件不存在: {engine_path}")
        return False
    
    try:
        # 加载引擎
        model = YOLO(engine_path)
        print("✅ 引擎加载成功")
        
        # 测试不同批次大小
        batch_sizes = [1, 2, 4, 8]
        
        for batch_size in batch_sizes:
            print(f"\n📊 测试批次大小: {batch_size}")
            
            # 创建测试数据
            import numpy as np
            test_images = []
            for i in range(batch_size):
                # 生成随机图像 (640x640x3)
                img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
                test_images.append(img)
            
            try:
                # 执行推理
                start_time = time.time()
                results = model(test_images, verbose=False)
                inference_time = time.time() - start_time
                
                print(f"  ✅ 批次 {batch_size}: {inference_time*1000:.1f}ms")
                print(f"  📈 平均每帧: {inference_time*1000/batch_size:.1f}ms")
                
            except Exception as e:
                print(f"  ❌ 批次 {batch_size} 测试失败: {e}")
                return False
        
        print("\n🎉 所有批次测试通过!")
        return True
        
    except Exception as e:
        print(f"❌ 引擎测试失败: {e}")
        return False

def main():
    """主函数"""
    print("简化版动态批次 TensorRT 引擎导出工具")
    print("=" * 60)
    
    # 导出动态引擎
    engine_path = export_dynamic_tensorrt_engine()
    
    if engine_path:
        # 测试动态引擎
        success = test_dynamic_engine(engine_path)
        
        if success:
            print(f"\n🎯 动态 TensorRT 引擎准备就绪!")
            print(f"📁 引擎路径: {engine_path}")
            print(f"✅ 支持批次大小: 1, 2, 4, 8, 16, 32")
            print(f"\n🚀 现在可以运行完整的批量性能测试了!")
        else:
            print(f"\n⚠️ 引擎导出成功但测试失败，请检查配置")
    else:
        print(f"\n❌ 引擎导出失败")

if __name__ == "__main__":
    main()