#!/usr/bin/env python3
"""
优化的多摄像头 TensorRT 推理脚本
支持：
1. 多路摄像头并发推理
2. 动态输入尺寸（320~640）
3. 批量推理优化
4. 详细性能统计
5. 高GPU利用率
"""

import cv2
import numpy as np
import yaml
import time
import datetime
import threading
import queue
import json
import os
from collections import defaultdict
from ultralytics import YOLO


class PerformanceStats:
    """性能统计类"""
    def __init__(self, cam_id):
        self.cam_id = cam_id
        self.frame_count = 0
        self.inference_times = []
        self.start_time = None
        self.lock = threading.Lock()
    
    def start(self):
        self.start_time = time.time()
    
    def record_inference(self, inference_time_ms):
        """记录推理时间（毫秒）"""
        with self.lock:
            self.inference_times.append(inference_time_ms)
            self.frame_count += 1
    
    def get_stats(self):
        """获取统计信息"""
        with self.lock:
            if not self.start_time or self.frame_count == 0:
                return None
            
            elapsed = time.time() - self.start_time
            avg_fps = self.frame_count / elapsed if elapsed > 0 else 0
            
            stats = {
                'cam_id': self.cam_id,
                'total_frames': self.frame_count,
                'elapsed_time': elapsed,
                'avg_fps': avg_fps,
                'avg_inference_ms': np.mean(self.inference_times) if self.inference_times else 0,
                'p50_inference_ms': np.percentile(self.inference_times, 50) if self.inference_times else 0,
                'p95_inference_ms': np.percentile(self.inference_times, 95) if self.inference_times else 0,
                'p99_inference_ms': np.percentile(self.inference_times, 99) if self.inference_times else 0,
            }
            return stats


class CameraReader:
    """摄像头读取器 - 独立线程读取帧"""
    def __init__(self, cam_id, rtsp_url, target_size=640):
        self.cam_id = cam_id
        self.rtsp_url = rtsp_url
        self.target_size = target_size
        self.frame_queue = queue.Queue(maxsize=2)
        self.running = True
        self.cap = None
        self.thread = None
        
        # 性能统计
        self.stats = PerformanceStats(cam_id)
    
    def start(self):
        """启动读取线程"""
        self.thread = threading.Thread(target=self._read_loop, daemon=True)
        self.thread.start()
        self.stats.start()
    
    def _read_loop(self):
        """读取循环"""
        try:
            # 打开视频流
            self.cap = cv2.VideoCapture(self.rtsp_url, cv2.CAP_FFMPEG)
            if not self.cap.isOpened():
                print(f"[{self.cam_id}] ⚠️ 无法打开视频流")
                return
            
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            print(f"[{self.cam_id}] ✅ 视频流已连接")
            
            while self.running:
                ret, frame = self.cap.read()
                if not ret:
                    time.sleep(0.01)
                    continue
                
                # Resize到目标尺寸
                if frame.shape[0] != self.target_size or frame.shape[1] != self.target_size:
                    frame = cv2.resize(frame, (self.target_size, self.target_size))
                
                # 放入队列（如果队列满，丢弃旧帧）
                if self.frame_queue.full():
                    try:
                        self.frame_queue.get_nowait()
                    except queue.Empty:
                        pass
                
                self.frame_queue.put(frame)
        
        except Exception as e:
            print(f"[{self.cam_id}] ❌ 读取线程异常: {e}")
        
        finally:
            if self.cap is not None:
                self.cap.release()
    
    def get_frame(self):
        """获取最新帧（非阻塞）"""
        try:
            return self.frame_queue.get_nowait()
        except queue.Empty:
            return None
    
    def stop(self):
        """停止读取"""
        self.running = False
        if self.thread is not None:
            self.thread.join(timeout=3.0)


class BatchInferenceEngine:
    """批量推理引擎"""
    def __init__(self, model_path, batch_size=4, imgsz=640, conf_thresh=0.45):
        self.model_path = model_path
        self.batch_size = batch_size
        self.imgsz = imgsz
        self.conf_thresh = conf_thresh
        
        # 加载模型
        print(f"🚀 加载 TensorRT 引擎: {model_path}")
        self.model = YOLO(model_path, task='detect')
        # TensorRT引擎不需要.to('cuda')，它已经是GPU模型
        print(f"✅ 引擎加载成功，批次大小: {batch_size}")
        
        # 批量缓冲区
        self.batch_buffer = []
        self.batch_cam_ids = []
        self.lock = threading.Lock()
    
    def add_to_batch(self, cam_id, frame):
        """添加帧到批次缓冲区"""
        with self.lock:
            self.batch_buffer.append(frame)
            self.batch_cam_ids.append(cam_id)
            
            # 如果达到批次大小，返回True
            return len(self.batch_buffer) >= self.batch_size
    
    def infer_batch(self):
        """批量推理"""
        with self.lock:
            if not self.batch_buffer:
                return []
            
            # 获取当前批次
            frames = self.batch_buffer[:self.batch_size]
            cam_ids = self.batch_cam_ids[:self.batch_size]
            
            # 清空已处理的
            self.batch_buffer = self.batch_buffer[self.batch_size:]
            self.batch_cam_ids = self.batch_cam_ids[self.batch_size:]
        
        # 批量推理
        start_time = time.time()
        
        try:
            results = self.model.predict(
                frames,
                imgsz=self.imgsz,
                conf=self.conf_thresh,
                verbose=False,
                device=0,  # 使用GPU 0
                half=True,
                classes=[0]  # person only
            )
            
            inference_time = (time.time() - start_time) * 1000  # 转换为毫秒
            
            # 计算每帧的推理时间
            per_frame_time = inference_time / len(frames)
            
            # 返回结果
            return [(cam_ids[i], results[i], per_frame_time) for i in range(len(frames))]
        
        except Exception as e:
            print(f"❌ 批量推理失败: {e}")
            return []
    
    def get_remaining_batch(self):
        """获取剩余的批次（用于测试结束时）"""
        with self.lock:
            if not self.batch_buffer:
                return []
            
            frames = self.batch_buffer
            cam_ids = self.batch_cam_ids
            
            self.batch_buffer = []
            self.batch_cam_ids = []
        
        # 推理剩余帧
        start_time = time.time()
        
        try:
            results = self.model.predict(
                frames,
                imgsz=self.imgsz,
                conf=self.conf_thresh,
                verbose=False,
                device=0,
                half=True,
                classes=[0]
            )
            
            inference_time = (time.time() - start_time) * 1000
            per_frame_time = inference_time / len(frames)
            
            return [(cam_ids[i], results[i], per_frame_time) for i in range(len(frames))]
        
        except Exception as e:
            print(f"❌ 剩余批次推理失败: {e}")
            return []


class MultiCameraInferenceSystem:
    """多摄像头推理系统"""
    def __init__(self, config_path, model_path, batch_size=4, target_size=640, max_cameras=None):
        self.config_path = config_path
        self.model_path = model_path
        self.batch_size = batch_size
        self.target_size = target_size
        
        # 加载配置
        with open(config_path, 'r', encoding='utf-8') as f:
            cfg = yaml.safe_load(f)
        
        # 获取摄像头配置
        cameras = cfg['cameras']
        if max_cameras:
            cameras = cameras[:max_cameras]
        
        # 初始化摄像头读取器
        self.camera_readers = {}
        for cam_cfg in cameras:
            cam_id = cam_cfg['id']
            rtsp_url = cam_cfg['rtsp_url']
            reader = CameraReader(cam_id, rtsp_url, target_size)
            self.camera_readers[cam_id] = reader
        
        print(f"✅ 初始化 {len(self.camera_readers)} 个摄像头")
        
        # 初始化推理引擎
        model_cfg = cfg['model']
        self.inference_engine = BatchInferenceEngine(
            model_path,
            batch_size=batch_size,
            imgsz=target_size,
            conf_thresh=model_cfg['conf_threshold']
        )
        
        self.running = False
    
    def start(self):
        """启动系统"""
        print(f"\n{'='*60}")
        print("启动多摄像头推理系统")
        print(f"{'='*60}")
        print(f"摄像头数量: {len(self.camera_readers)}")
        print(f"批次大小: {self.batch_size}")
        print(f"输入尺寸: {self.target_size}x{self.target_size}")
        print(f"{'='*60}\n")
        
        # 启动所有摄像头读取器
        for reader in self.camera_readers.values():
            reader.start()
        
        # 等待摄像头连接
        print("⏳ 等待摄像头连接...")
        time.sleep(3)
        
        self.running = True
    
    def run(self, test_duration=60):
        """运行推理"""
        print(f"🚀 开始推理，测试时长: {test_duration}秒\n")
        
        start_time = time.time()
        last_print_time = start_time
        total_frames = 0
        
        try:
            while self.running and (time.time() - start_time) < test_duration:
                # 从所有摄像头收集帧
                frames_collected = 0
                for cam_id, reader in self.camera_readers.items():
                    frame = reader.get_frame()
                    if frame is not None:
                        # 添加到批次缓冲区
                        batch_ready = self.inference_engine.add_to_batch(cam_id, frame)
                        frames_collected += 1
                        
                        # 如果批次准备好，执行推理
                        if batch_ready:
                            results = self.inference_engine.infer_batch()
                            
                            # 记录统计
                            for cam_id, result, inference_time in results:
                                self.camera_readers[cam_id].stats.record_inference(inference_time)
                                total_frames += 1
                
                # 如果没有收集到帧，短暂休眠
                if frames_collected == 0:
                    time.sleep(0.001)
                
                # 每5秒打印一次进度
                current_time = time.time()
                if current_time - last_print_time >= 5.0:
                    elapsed = current_time - start_time
                    avg_fps = total_frames / elapsed if elapsed > 0 else 0
                    print(f"⏱️ {elapsed:.0f}s | 总帧数: {total_frames} | 平均FPS: {avg_fps:.1f}")
                    last_print_time = current_time
        
        except KeyboardInterrupt:
            print("\n⏹️ 测试被用户中断")
        
        finally:
            # 处理剩余的批次
            remaining_results = self.inference_engine.get_remaining_batch()
            for cam_id, result, inference_time in remaining_results:
                self.camera_readers[cam_id].stats.record_inference(inference_time)
                total_frames += 1
        
        # 生成报告
        self.generate_report(total_frames, time.time() - start_time)
    
    def generate_report(self, total_frames, elapsed_time):
        """生成性能报告"""
        print(f"\n{'='*60}")
        print("性能测试报告")
        print(f"{'='*60}\n")
        
        # 收集所有摄像头的统计
        all_stats = []
        all_inference_times = []
        
        for cam_id, reader in self.camera_readers.items():
            stats = reader.stats.get_stats()
            if stats:
                all_stats.append(stats)
                all_inference_times.extend(reader.stats.inference_times)
        
        # 总体统计
        avg_fps = total_frames / elapsed_time if elapsed_time > 0 else 0
        
        print(f"总体性能:")
        print(f"  总帧数: {total_frames}")
        print(f"  测试时长: {elapsed_time:.1f}秒")
        print(f"  平均FPS: {avg_fps:.1f}")
        
        if all_inference_times:
            print(f"  平均推理延迟: {np.mean(all_inference_times):.1f}ms")
            print(f"  P50推理延迟: {np.percentile(all_inference_times, 50):.1f}ms")
            print(f"  P95推理延迟: {np.percentile(all_inference_times, 95):.1f}ms")
            print(f"  P99推理延迟: {np.percentile(all_inference_times, 99):.1f}ms")
        
        print(f"\n各摄像头性能:")
        print(f"{'摄像头ID':<15} {'帧数':<10} {'FPS':<10} {'平均延迟(ms)':<15} {'P95延迟(ms)':<15}")
        print(f"{'-'*70}")
        
        for stats in sorted(all_stats, key=lambda x: x['cam_id']):
            print(f"{stats['cam_id']:<15} {stats['total_frames']:<10} "
                  f"{stats['avg_fps']:<10.1f} {stats['avg_inference_ms']:<15.1f} "
                  f"{stats['p95_inference_ms']:<15.1f}")
        
        # 保存结果
        output_dir = "multi_camera_results"
        os.makedirs(output_dir, exist_ok=True)
        
        results_data = {
            'total_frames': total_frames,
            'elapsed_time': elapsed_time,
            'avg_fps': avg_fps,
            'avg_inference_ms': np.mean(all_inference_times) if all_inference_times else 0,
            'p50_inference_ms': np.percentile(all_inference_times, 50) if all_inference_times else 0,
            'p95_inference_ms': np.percentile(all_inference_times, 95) if all_inference_times else 0,
            'p99_inference_ms': np.percentile(all_inference_times, 99) if all_inference_times else 0,
            'camera_stats': all_stats,
            'batch_size': self.batch_size,
            'target_size': self.target_size,
            'num_cameras': len(self.camera_readers),
            'timestamp': datetime.datetime.now().isoformat()
        }
        
        json_file = os.path.join(output_dir, f"results_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
        with open(json_file, 'w', encoding='utf-8') as f:
            json.dump(results_data, f, indent=2, ensure_ascii=False)
        
        print(f"\n✅ 结果已保存: {json_file}")
    
    def stop(self):
        """停止系统"""
        print("\n正在停止系统...")
        self.running = False
        
        # 停止所有摄像头读取器
        for reader in self.camera_readers.values():
            reader.stop()
        
        print("系统已停止")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='多摄像头TensorRT推理系统')
    parser.add_argument('--config', default='config.yaml', help='配置文件路径')
    parser.add_argument('--model', default='C:/Users/16337/PycharmProjects/Security/yolo11n.engine', 
                       help='TensorRT引擎路径')
    parser.add_argument('--batch-size', type=int, default=4, help='批次大小')
    parser.add_argument('--target-size', type=int, default=640, help='输入尺寸')
    parser.add_argument('--duration', type=int, default=60, help='测试时长（秒）')
    parser.add_argument('--max-cameras', type=int, default=None, help='最大摄像头数量')
    
    args = parser.parse_args()
    
    print("多摄像头 TensorRT 推理系统")
    print("=" * 60)
    
    # 检查文件
    if not os.path.exists(args.config):
        print(f"❌ 配置文件不存在: {args.config}")
        return
    
    if not os.path.exists(args.model):
        print(f"❌ TensorRT引擎不存在: {args.model}")
        return
    
    # 创建系统
    try:
        system = MultiCameraInferenceSystem(
            config_path=args.config,
            model_path=args.model,
            batch_size=args.batch_size,
            target_size=args.target_size,
            max_cameras=args.max_cameras
        )
        
        # 启动系统
        system.start()
        
        # 运行推理
        system.run(test_duration=args.duration)
        
        # 停止系统
        system.stop()
        
        print("\n🎉 测试完成！")
    
    except Exception as e:
        print(f"\n❌ 系统异常: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⏹️ 程序被用户中断")
    except Exception as e:
        print(f"\n❌ 程序异常: {e}")
        import traceback
        traceback.print_exc()