"""
Example 4: Multi-Sensor with Multiprocessing

Demonstrates the simplified MultiSensor API for running multiple sensors
in parallel with full CUDA support.
"""

import os
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"

import cv2
import orisys
import multiprocessing as mp


def main():
    # Configure sensors
    configs = [
        {
            'vid_src': './s1.mp4',
            'isstitch': True,
            'cuda': True,
            'cal_path': './config/cal.npy'
        },
        {
            'vid_src': './s1_copy.mp4',
            'isstitch': True,
            'cuda': True,
            'cal_path': './config/cal.npy'
        },
    ]
    
    with orisys.MultiSensor(configs) as multi:
        print("\nPress 'q' to quit")
        
        for all_data in multi:
            for data in all_data:
                if data is None:
                    continue
                
                sensor_id = data['sensor_id']
                print(f"Sensor {sensor_id}: FPS={data['fps']:5.1f}, "
                      f"Fn={data['fnormal']:8.2f}")
                
                # Draw arrows
                if data['img'] is not None and data['flow'] is not None:
                    arrows = orisys.util.draw_arrows(
                        data['img'].copy(),
                        data['flow'],
                        threshold=5
                    )
                    cv2.imshow(f"Sensor {sensor_id}", arrows)
                
                # Show depth
                if data['depth'] is not None:
                    depth_img = cv2.applyColorMap(
                        (data['depth'] * 255).astype('uint8'),
                        cv2.COLORMAP_HOT
                    )
                    cv2.imshow(f"Sensor {sensor_id} - Depth", depth_img)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    
    cv2.destroyAllWindows()


def main_manual():
    """Alternative: Manual start/stop control"""
    configs = [
        {'vid_src': './s1.mp4', 'isstitch': True, 'cal_path': './config/cal.npy'},
        {'vid_src': './s1_copy.mp4', 'isstitch': True, 'cal_path': './config/cal.npy'},
    ]
    
    multi = orisys.MultiSensor(configs)
    multi.start()
    
    try:
        while True:
            # Get data from specific sensor
            data0 = multi.get(0)  # or multi[0]
            data1 = multi.get(1)  # or multi[1]
            
            # Or get all at once
            # all_data = multi.get_all()
            
            for data in [data0, data1]:
                if data and data['img'] is not None:
                    cv2.imshow(f"Sensor {data['sensor_id']}", data['img'])
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        multi.stop()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    mp.freeze_support()  # Required for Windows
    main()