"""
示例 4: 多传感器并行处理（多进程模式 - 不使用 MultiSensor）

本示例演示如何使用 multiprocessing 同时运行两个传感器：
- 每个传感器运行在独立的进程中
- 每个进程拥有独立的 CUDA 上下文
- 使用 multiprocessing.Process 和 Queue 进行进程间通信

架构说明：
    进程 1 (传感器 0):  采集 → 光流 → HHD → 接触检测 → 队列
    进程 2 (传感器 1):  采集 → 光流 → HHD → 接触检测 → 队列
                                              ↓
    主进程:            显示 ← ───────────────┘

适用场景：多指触觉传感器、多视角采集、需要 GPU 加速的多路处理

注意事项：
- Windows 下必须在 if __name__ == '__main__' 中调用 mp.freeze_support()
- 每个传感器进程会占用独立的 GPU 显存
"""
import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" 
import cv2
import orisys
import numpy as np
import multiprocessing as mp
from multiprocessing import Queue, Process
import time
import ctypes
import platform
import ctypes

# Multi-sensor FPS stabilization options
STABILITY_MODE = True          # True: prioritize similar FPS across sensors
TARGET_FPS = 30.0             # Used only when STABILITY_MODE=True
SEND_VISUALIZATION_DATA = False  # True sends img/flow/depth to main process (higher IPC cost)

def _enable_high_res_sleep():
    """
    Enable high-resolution sleep for accurate FPS pacing.
    Note: time.perf_counter() is already high-resolution on all platforms (Windows, macOS, Linux).
    This function only improves time.sleep() accuracy on Windows.
    """
    if os.name == "nt":  # Windows (x64 and ARM64)
        try:
            # timeBeginPeriod improves time.sleep() accuracy (default ~15ms -> 1ms)
            # time.perf_counter() is already high-res on Windows (uses QueryPerformanceCounter)
            winmm = ctypes.WinDLL("winmm")
            winmm.timeBeginPeriod(1)  # Request 1ms timer resolution for sleep()
            return winmm
        except Exception:
            return None
    # macOS and Linux: time.perf_counter() and time.sleep() are already high-resolution
    # No special setup needed
    return None


def _disable_high_res_sleep(handle):
    """Restore default scheduler timer resolution (Windows only)."""
    if handle is None:
        return
    if os.name == "nt" and isinstance(handle, ctypes.WinDLL):
        try:
            handle.timeEndPeriod(1)
        except Exception:
            pass
    # macOS and Linux don't need cleanup


def sensor_worker(sensor_id, vid_src, config_name, cal_path, result_queue, stop_event):
    """
    传感器工作进程函数
    
    Args:
        sensor_id: 传感器编号
        vid_src: 视频源（文件路径或摄像头索引）
        config_name: 配置名称
        result_queue: 结果队列，用于向主进程发送数据
        stop_event: 停止事件，用于通知工作进程退出
    """
    timer_handle = None
    try:
        # Limit per-process CPU thread fan-out to reduce contention jitter.
        os.environ["OMP_NUM_THREADS"] = "1"
        os.environ["OPENBLAS_NUM_THREADS"] = "1"
        os.environ["MKL_NUM_THREADS"] = "1"
        cv2.setNumThreads(1)
        if STABILITY_MODE and TARGET_FPS > 0:
            timer_handle = _enable_high_res_sleep()

        # 在每个进程中创建独立的传感器实例
        sensor = orisys.Sensor(
            vid_src, 
            isstitch=True, 
            cuda=True, 
            verbose=False,  # 减少子进程输出
            config_name=config_name,
            cal_path=cal_path
        )
        
        print(f"传感器 {sensor_id} 已启动")
        
        target_period = (1.0 / TARGET_FPS) if (STABILITY_MODE and TARGET_FPS > 0) else 0.0
        next_tick = time.perf_counter()
        while not stop_event.is_set():
            try:
                # 步骤1: 获取图像
                sensor.get_img()
                
                # 步骤2: 计算形变场
                sensor.compute_deformation()
                
                # 步骤3: 计算触碰点
                sensor.compute_contact()
                
                # 步骤4: 读取信息
                if SEND_VISUALIZATION_DATA:
                    fps, flow, img, contour, centroid, depth_map = sensor.read_info(
                        sensor.info.FPS,
                        sensor.info.VRAW,
                        sensor.info.IMG,
                        sensor.info.CONTOUR,
                        sensor.info.CENTROID,
                        sensor.info.DEPTH
                    )
                    # Full payload (higher IPC cost)
                    data = {
                        'sensor_id': sensor_id,
                        'fps': fps,
                        'flow': flow.copy() if flow is not None else None,
                        'img': img.copy() if img is not None else None,
                        'contour': contour,
                        'centroid': centroid,
                        'depth_map': depth_map.copy() if depth_map is not None else None,
                        'timestamp': time.time()
                    }
                else:
                    # Lightweight payload for stable multi-process throughput
                    fps, = sensor.read_info(sensor.info.FPS)
                    data = {
                        'sensor_id': sensor_id,
                        'fps': fps,
                        'timestamp': time.time()
                    }
                
                # 步骤6: 发送数据到主进程（非阻塞）
                try:
                    result_queue.put_nowait(data)
                except:
                    # Queue full: keep latest data (drop oldest)
                    try:
                        result_queue.get_nowait()
                        result_queue.put_nowait(data)
                    except:
                        pass

                # Pace workers to a common target so sensor FPS stays similar.
                # Use drift-free scheduler: next_tick increments by fixed period.
                if target_period > 0:
                    next_tick += target_period
                    now = time.perf_counter()
                    sleep_time = next_tick - now
                    if sleep_time > 0:
                        time.sleep(sleep_time)
                    else:
                        # If we are behind, reset schedule to current time.
                        next_tick = now
                
            except Exception as e:
                print(f"传感器 {sensor_id} 处理错误: {e}")
                time.sleep(0.1)
        
        # 清理资源
        sensor.disconnect()
        if timer_handle is not None:
            _disable_high_res_sleep(timer_handle)
        print(f"传感器 {sensor_id} 已停止")
        
    except Exception as e:
        if timer_handle is not None:
            _disable_high_res_sleep(timer_handle)
        print(f"传感器 {sensor_id} 初始化失败: {e}")
        result_queue.put({'sensor_id': sensor_id, 'error': str(e)})


def main():
    # =========================================================================
    # 1. 配置多个传感器
    # =========================================================================
    num_sensors = 5
    sensor_configs = [
        {
            'sensor_id': 0,
            'vid_src': './s1_720p.mp4',
            'config_name': './config/prototype_of3_720p.json',
            'cal_path': './config/dd01a_0_720p.npy'
        },
        {
            'sensor_id': 1,
            'vid_src': './s1_720p.mp4',
            'config_name': './config/prototype_of3_720p.json',
            'cal_path': './config/dd01a_0_720p.npy'
        },
        {
            'sensor_id': 2,
            'vid_src': './s1_720p.mp4',
            'config_name': './config/prototype_of3_720p.json',
            'cal_path': './config/dd01a_0_720p.npy'
        },
        {
            'sensor_id': 3,
            'vid_src': './s1_720p.mp4',
            'config_name': './config/prototype_of3_720p.json',
            'cal_path': './config/dd01a_0_720p.npy'
        },
        {
            'sensor_id': 4,
            'vid_src': './s1_720p.mp4',
            'config_name': './config/prototype_of3_720p.json',
            'cal_path': './config/dd01a_0_720p.npy'
        },
    ]
    
    # =========================================================================
    # 2. 创建进程间通信对象
    # =========================================================================
    # Keep queue short to prevent backlog-induced jitter (latest-data policy)
    result_queues = [Queue(maxsize=2) for _ in sensor_configs]
    stop_events = [mp.Event() for _ in sensor_configs]  # 停止事件
    
    # =========================================================================
    # 3. 启动传感器工作进程
    # =========================================================================
    processes = []
    for i, config in enumerate(sensor_configs):
        p = Process(
            target=sensor_worker,
            args=(i, config['vid_src'], config['config_name'], config['cal_path'], result_queues[i], stop_events[i])
        )
        p.start()
        processes.append(p)
    
    print(f"\n已启动 {len(processes)} 个传感器进程")
    print("按 'q' 退出程序")
    
    # 用于存储每个传感器的最新数据
    latest_data = [None] * len(sensor_configs)
    
    # 用于控制打印频率，减少I/O开销
    print_counter = 0
    print_interval = 30  # 每30帧打印一次，进一步减少I/O
    
    # 用于控制按键检查频率
    key_check_counter = 0
    key_check_interval = 5  # 每5帧检查一次按键
    
    try:
        # =====================================================================
        # 4. 主循环：从队列获取数据并显示
        # =====================================================================
        while True:
            # 从所有队列中快速获取最新数据（高效清空队列，只保留最新帧）
            for i, queue in enumerate(result_queues):
                # 快速清空队列，只保留最后一帧
                data = None
                # 一次性获取所有可用数据，只保留最后一个
                while True:
                    try:
                        data = queue.get_nowait()
                    except:
                        break
                
                if data is not None:
                    latest_data[i] = data
            
            # 显示所有传感器的数据
            for data in latest_data:
                if data is None or 'error' in data:
                    continue
                
                sid = data['sensor_id']
                
                # 显示形变场
                if data.get('img') is not None and data.get('flow') is not None:
                    arrows = orisys.util.draw_arrows(
                        data['img'],
                        data['flow'],
                        threshold=5,
                        grid_spacing=10,
                        arrow_scale=1.0
                    )
                    cv2.imshow(f"Deformation {sid}", arrows)
                
                # # 显示触碰点和区域
                # if data['img'] is not None:
                #     image_with_foe = orisys.util.draw_contact(
                #         data['img'], 
                #         data['contour'], 
                #         data['centroid']
                #     )
                #     cv2.imshow(f"Contact Centroid {sid}", image_with_foe)
                
                # # 显示深度图
                # if data['depth_map'] is not None and data['depth_map'].size > 0:
                #     div_abs = data['depth_map']
                #     if div_abs.max() > div_abs.min():
                #         div_normalized = ((div_abs - div_abs.min()) / 
                #                         (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
                #     else:
                #         div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
                #     cv2.imshow(f"Depth Map {sid}", 
                #              cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
            
            # 降低打印频率以减少I/O开销（在主循环级别控制）
            print_counter += 1
            if print_counter >= print_interval:
                for data in latest_data:
                    if data is not None and 'error' not in data:
                        sid = data['sensor_id']
                        print(f"传感器 {sid}: FPS={data['fps']:5.1f}")
                print_counter = 0
            
            # 降低按键检查频率以减少开销
            key_check_counter += 1
            if key_check_counter >= key_check_interval:
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    print("\n退出程序...")
                    break
                key_check_counter = 0
    
    finally:
        # =====================================================================
        # 5. 清理资源
        # =====================================================================
        # 通知所有工作进程停止
        for event in stop_events:
            event.set()
        
        # 等待所有进程结束
        for p in processes:
            p.join(timeout=2)
            if p.is_alive():
                p.terminate()
                p.join()
        
        cv2.destroyAllWindows()
        print("所有传感器进程已停止")


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