"""
示例 7: 五传感器网格显示

本示例演示如何扩展到 5 个传感器：
- 使用异步多进程模式运行所有传感器
- 将 5 个传感器画面拼接成 2x3 网格显示
- 实时显示每个传感器的 FPS 和力数据

网格布局：
    ┌─────────┬─────────┬─────────┐
    │ 传感器0 │ 传感器1 │ 传感器2 │
    ├─────────┼─────────┼─────────┤
    │ 传感器3 │ 传感器4 │  (空)   │
    └─────────┴─────────┴─────────┘

硬件要求估算：
- CPU: 建议 8 核以上（每个传感器约 2-3 线程）
- GPU: 建议 8GB 显存以上（每个传感器约 500MB-1GB）
- 内存: 建议 16GB 以上

注意：本示例使用同一个视频文件模拟 5 个传感器，
     实际使用时请替换为不同的视频源或摄像头索引。
"""

import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"  # Windows 兼容性设置

import cv2
import orisys
import multiprocessing as mp
import numpy as np


def main():
    # =========================================================================
    # 1. 配置 5 个传感器
    # =========================================================================
    # 注意：这里使用同一个视频文件作为演示
    # 实际应用中应替换为不同的视频源：
    #   - 摄像头: 0, 1, 2, 3, 4
    #   - 视频文件: './sensor0.mp4', './sensor1.mp4', ...
    configs = [
        {'vid_src': './s1.mp4', 'isstitch': True, 'cuda': True, 'cal_path': './config/cal.npy'},
        {'vid_src': './s1.mp4', 'isstitch': True, 'cuda': True, 'cal_path': './config/cal.npy'},
        {'vid_src': './s1.mp4', 'isstitch': True, 'cuda': True, 'cal_path': './config/cal.npy'},
        {'vid_src': './s1.mp4', 'isstitch': True, 'cuda': True, 'cal_path': './config/cal.npy'},
        {'vid_src': './s1.mp4', 'isstitch': True, 'cuda': True, 'cal_path': './config/cal.npy'},
    ]
    
    # =========================================================================
    # 2. 启动多传感器异步模式
    # =========================================================================
    with orisys.MultiSensor.async_context(configs) as multi:
        print(f"\n{len(configs)} 个传感器已启动，按 'q' 退出")
        
        # =====================================================================
        # 3. 主循环
        # =====================================================================
        for all_data in multi:
            # 过滤有效数据（有图像的传感器）
            valid_data = [d for d in all_data if d and d['img'] is not None]
            
            # -----------------------------------------------------------------
            # 3.1 打印所有传感器状态（单行显示）
            # -----------------------------------------------------------------
            for data in valid_data:
                print(f"S{data['sensor_id']}: "
                      f"FPS={data['fps']:4.1f} "
                      f"Fn={data['fnormal']:7.1f}", end=" | ")
            if valid_data:
                print()  # 换行
            
            # -----------------------------------------------------------------
            # 3.2 准备网格显示的图像
            # -----------------------------------------------------------------
            imgs = []
            for data in valid_data:
                # 绘制光流箭头
                if data['flow'] is not None:
                    img = orisys.util.draw_arrows(
                        data['img'].copy(),
                        data['flow'],
                        threshold=5
                    )
                else:
                    img = data['img']
                
                # 在图像上标注传感器编号
                cv2.putText(
                    img,
                    f"S{data['sensor_id']}",
                    (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.8,
                    (0, 255, 0),
                    2
                )
                imgs.append(img)
            
            # -----------------------------------------------------------------
            # 3.3 创建 2x3 网格布局
            # -----------------------------------------------------------------
            if len(imgs) >= 5:
                h, w = imgs[0].shape[:2]
                
                # 创建黑色占位图像（用于第6格）
                if len(imgs[0].shape) == 3:
                    black = np.zeros((h, w, 3), dtype=np.uint8)
                else:
                    black = np.zeros((h, w), dtype=np.uint8)
                
                # 拼接网格
                # 第一行: 传感器 0, 1, 2
                row1 = np.hstack([imgs[0], imgs[1], imgs[2]])
                # 第二行: 传感器 3, 4, 空
                row2 = np.hstack([imgs[3], imgs[4], black])
                # 合并两行
                grid = np.vstack([row1, row2])
                
                # 缩放到适合显示的尺寸
                grid = cv2.resize(grid, (1200, 800))
                cv2.imshow("5 传感器网格视图", grid)
            
            # 检查退出键
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # Windows 下多进程必需的初始化
    mp.freeze_support()
    main()