"""
RTX 3050 优化压力测试启动脚本

目标：
- 使用原生 TensorRT API 而不是 Ultralytics 封装
- 达到 GPU 利用率 70%+ 而不是 30%
- 测试真正的 GPU 并发性能极限

测试内容：
1. 不同分辨率下最大每秒处理帧数
2. 最大接入摄像头路数
3. 不同摄像头数量下单路最大帧数
4. 不同抽帧率下最大摄像头数

优化策略：
- 多流并行推理 (2-8个CUDA流)
- 大批次处理 (4-32批次)
- 原生 TensorRT API
- 优化内存管理

运行方式：
    python run_optimized_stress_test.py
"""

import sys
import os

# 添加当前目录到路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

def check_requirements():
    """检查必要的依赖"""
    missing = []
    tensorrt_available = False
    
    try:
        import tensorrt
        print(f"✅ TensorRT: {tensorrt.__version__}")
        tensorrt_available = True
    except (ImportError, FileNotFoundError) as e:
        print(f"⚠️ TensorRT: 不可用 ({str(e)[:50]}...)")
        print("  将使用 Ultralytics 优化模式")
    
    try:
        import torch
        print(f"✅ PyTorch: {torch.__version__}")
        if torch.cuda.is_available():
            print(f"✅ CUDA: {torch.version.cuda}")
            print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
        else:
            missing.append("CUDA")
    except ImportError:
        missing.append("torch")
    
    try:
        import ultralytics
        print(f"✅ Ultralytics: {ultralytics.__version__}")
    except ImportError:
        missing.append("ultralytics")
    
    if missing:
        print(f"\n❌ 缺少关键依赖: {', '.join(missing)}")
        return False
    
    if not tensorrt_available:
        print(f"\n⚠️ 注意: 将使用 Ultralytics 优化模式而不是原生 TensorRT")
        print(f"  仍然可以通过以下方式提升性能:")
        print(f"  • 多线程并行推理")
        print(f"  • 大批次处理")
        print(f"  • GPU 预处理")
        print(f"  • 优化内存管理")
    
    return True

def main():
    print("=" * 60)
    print("RTX 3050 优化压力测试")
    print("=" * 60)
    print("目标: GPU 利用率 70%+ (vs 之前的 30%)")
    print("策略: 原生 TensorRT + 多流并行 + 大批次")
    print("=" * 60)
    
    # 检查依赖
    print("\n🔍 检查依赖...")
    if not check_requirements():
        return
    
    # 模型路径
    MODEL_PATH = "C:/Users/16337/PycharmProjects/Security_project/yolov8n.pt"
    OUTPUT_DIR = "./optimized_stress_results"
    
    print(f"\n📁 配置:")
    print(f"  模型: {MODEL_PATH}")
    print(f"  输出: {OUTPUT_DIR}")
    
    if not os.path.exists(MODEL_PATH):
        print(f"\n❌ 模型文件不存在: {MODEL_PATH}")
        return
    
    print(f"\n🚀 优化策略:")
    print(f"  • 原生 TensorRT API (避免 Ultralytics 封装损失)")
    print(f"  • 多 CUDA 流并行 (2-8个流)")
    print(f"  • 大批次处理 (4-32批次)")
    print(f"  • 优化内存管理")
    print(f"  • 2GB TensorRT 工作空间")
    
    print(f"\n📊 测试内容:")
    print(f"  1. 最大处理帧数 (320x320, 480x480)")
    print(f"  2. 最大摄像头接入数")
    print(f"  3. 不同摄像头数量下单路最大帧数")
    print(f"  4. 不同抽帧率下最大摄像头数")
    
    print(f"\n⏱️ 预计时间: 15-25 分钟")
    print(f"💾 结果自动保存，Ctrl+C 中断不丢数据")
    print("=" * 60)
    
    # 确认开始
    try:
        input("\n按 Enter 开始测试，或 Ctrl+C 取消...")
    except KeyboardInterrupt:
        print("\n❌ 测试取消")
        return
    
    # 运行优化测试
    try:
        # 检查是否可以使用原生 TensorRT
        try:
            import tensorrt
            from benchmark.optimized_stress_test import run_optimized_stress_test
            print(f"\n🚀 使用原生 TensorRT 优化测试...")
            run_optimized_stress_test(MODEL_PATH, OUTPUT_DIR)
        except (ImportError, FileNotFoundError):
            # 使用 Ultralytics 优化模式
            from benchmark.ultralytics_optimized_stress import run_ultralytics_optimized_test
            print(f"\n🚀 使用 Ultralytics 优化测试...")
            run_ultralytics_optimized_test(MODEL_PATH, OUTPUT_DIR)
        
        print(f"\n🎉 测试完成!")
        print(f"📊 查看结果:")
        print(f"  - JSON 数据: {OUTPUT_DIR}/*_results_*.json")
        print(f"  - 测试报告: {OUTPUT_DIR}/*_report_*.md")
        
        # 生成可视化
        print(f"\n🎨 生成可视化图表...")
        try:
            from benchmark.optimized_visualizer import generate_optimized_charts
            chart_files = generate_optimized_charts(OUTPUT_DIR)
            print(f"✅ 生成了 {len(chart_files)} 个图表")
        except Exception as e:
            print(f"⚠️ 可视化生成失败: {e}")
            print(f"💡 可以手动运行: python generate_optimized_charts.py")
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()