# rpi_cam_universal.py
# Две камеры: либо CSI (picamera2), либо Daheng USB/Gige
# Один файл — два режима — идеальный интерфейс — ничего не ломается

import os
import cv2
import numpy as np
import time
import pprint

# ====================== ВЫБОР РЕЖИМА ======================
USE_RPI_CAMERAS = False   # True = CSI камеры (IMX219/IMX708 и т.д.), False = Daheng USB
# =========================================================

w, h = 640, 480

# ====================== ОБЩАЯ ФУНКЦИЯ КОНТРАСТА ======================
def linear_contrast_stretch(img):
    """Правильный авто-контраст для монохромных кадров"""
    if img.ndim != 2 or img.size == 0:
        return img
    mn, mx = img.min(), img.max()
    if mx == mn:
        return img.copy()
    stretched = (img.astype(np.float32) - mn) * (255.0 / (mx - mn))
    return np.clip(stretched, 0, 255).astype(np.uint8)


# ====================== RPI CSI КАМЕРЫ (picamera2) ======================
if USE_RPI_CAMERAS:
    from picamera2 import Picamera2
    from libcamera import Transform

    os.environ["LIBCAMERA_LOG_LEVELS"] = "3"

    _TRANSFORMS = {
        0: Transform(vflip=1),   # подбери под свою установку
        1: Transform(),
    }

    picams = None

    def init_rpi():
        global picams
        print("Запуск Raspberry Pi CSI камер...")
        picams = [Picamera2(i) for i in range(2)]

        config_controls = {
            "FrameDurationLimits": (16667, 16667),  # 60 fps
            "Saturation": 0,
        }

        for i, cam in enumerate(picams):
            config = cam.create_preview_configuration(
                main={"size": (w, h), "format": "YUV420"},
                transform=_TRANSFORMS[i],
                controls=config_controls
            )
            cam.configure(config)
            time.sleep(0.3)

        for cam in picams:
            cam.start()
        time.sleep(1.2)
        print("RPi камеры готовы")

    def capture_rpi():
        reqs = [cam.capture_request() for cam in picams]
        frames = []
        for req in reqs:
            arr = req.make_array("main")
            y_h = arr.shape[0] * 2 // 3
            y = arr[:y_h, :]
            y = linear_contrast_stretch(y)
            frame = cv2.cvtColor(y, cv2.COLOR_GRAY2BGR)
            frames.append(frame)
            req.release()
        return frames

    def close_rpi():
        if picams:
            for cam in picams:
                try: cam.stop(); cam.close()
                except: pass


# ====================== DAHENG USB КАМЕРЫ ======================
else:
    import gxipy as gx

    device_manager = None
    cam1 = cam2 = None

    def init_daheng():
        global device_manager, cam1, cam2
        print("Запуск Daheng USB/Gige камер...")
        device_manager = gx.DeviceManager()
        dev_num, _ = device_manager.update_all_device_list()

        if dev_num < 2:
            raise RuntimeError(f"Найдено только {dev_num} Daheng камер!")

        cam1 = device_manager.open_device_by_index(1)
        cam2 = device_manager.open_device_by_index(2)

        for cam in (cam1, cam2):
            cam.TriggerMode.set(0)
            try:
                cam.PixelFormat.set(gx.GxPixelFormatEntry.MONO8)
            except:
                pass
            cam.Width.set(w)
            cam.Height.set(h)

        cam1.stream_on()
        cam2.stream_on()
        time.sleep(1.0)
        print("Daheng камеры готовы")

    def capture_daheng():
        frames = [None, None]

        raw1 = cam1.data_stream[0].get_image(timeout=500)
        raw2 = cam2.data_stream[0].get_image(timeout=500)

        if raw1:
            img = raw1.get_numpy_array()
            if img is not None:
                img = linear_contrast_stretch(img)
                frames[0] = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        if raw2:
            img = raw2.get_numpy_array()
            if img is not None:
                img = linear_contrast_stretch(img)
                frames[1] = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

        return frames

    def close_daheng():
        print("Закрытие Daheng камер...")
        for cam in (cam1, cam2):
            if cam:
                try: cam.stream_off()
                except: pass
                try: cam.close_device()
                except: pass


# ====================== ОСНОВНОЙ ЦИКЛ (ОДИНАКОВЫЙ ДЛЯ ВСЕХ) ======================
def main():
    cv2.namedWindow("Cam 0 | Cam 1", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("Cam 0 | Cam 1", w*2, h)

    mode_name = "Raspberry Pi CSI" if USE_RPI_CAMERAS else "Daheng USB/Gige"
    print(f"\n=== {mode_name} ===\n")

    if USE_RPI_CAMERAS:
        init_rpi()
        capture_func = capture_rpi
        cleanup_func = close_rpi
    else:
        init_daheng()
        capture_func = capture_daheng
        cleanup_func = close_daheng

    fps_counter = 0
    fps_start = time.time()

    try:
        while True:
            frames = capture_func()

            for i, frame in enumerate(frames):
                if frame is None:
                    frame = np.zeros((h, w, 3), np.uint8)
                    cv2.putText(frame, f"Cam {i+1}: NO SIGNAL", (40, 240),
                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3)
                else:
                    cv2.putText(frame, f"Cam {i+1}", (15, 40),
                                cv2.FONT_HERSHEY_SIMPLEX, 1.1, (0, 255, 0), 3)

            combined = np.hstack(frames)
            cv2.imshow("Cam 0 | Cam 1", combined)

            # FPS в консоль
            fps_counter += 1
            if time.time() - fps_start >= 1.0:
                print(f"FPS: {fps_counter / (time.time() - fps_start):.1f}")
                fps_counter = 0
                fps_start = time.time()

            if cv2.waitKey(1) == ord('q'):
                break

    except KeyboardInterrupt:
        print("\nОстановлено пользователем")
    finally:
        cleanup_func()
        cv2.destroyAllWindows()
        print("Всё отключено. До свидания!")

if __name__ == "__main__":
    main()