---

import cv2 as cv
import numpy as np


# --------------------------------------------------
# Настройки
# --------------------------------------------------
# CAMERA_INDEX, , FRAME_HEIGHT, CALIBRATION_FILE
# ----------------------------FRAME_WIDTH----------------------
FRAME_WIDTH = 640
FRAME_HEIGHT = 480
CAMERA_INDEX = 1
CALIBRATION_FILE = "camera_calibration.npz"
# Загружаем параметры калибровки
# --------------------------------------------------

calibration = np.load(CALIBRATION_FILE)

camera_matrix = calibration["camera_matrix"]
distortion = calibration["distortion"]


print("Матрица камеры:")
print(camera_matrix)

print()
print("Коэффициенты искажения:")
print(distortion)


# --------------------------------------------------
# Открываем камеру
# --------------------------------------------------

camera = cv.VideoCapture(CAMERA_INDEX)

if not camera.isOpened():
    print("Не удалось открыть камеру")
    raise SystemExit


camera.set(cv.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
camera.set(cv.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)


width = int(camera.get(cv.CAP_PROP_FRAME_WIDTH))
height = int(camera.get(cv.CAP_PROP_FRAME_HEIGHT))


print()
print("Камера открыта")
print("Разрешение:", width, "x", height)
print("Q или ESC — завершить работу")


# --------------------------------------------------
# Получаем и исправляем кадры
# --------------------------------------------------

try:
    while True:

        success, frame = camera.read()

        if not success:
            print("Не удалось получить кадр")
            break


        # Исправляем искажения объектива
        corrected_frame = cv.undistort(
            frame,
            camera_matrix,
            distortion
        )


        # Создаём копии для надписей
        original_view = frame.copy()
        corrected_view = corrected_frame.copy()


        # Подпись исходного изображения
        cv.putText(
            original_view,
            "ORIGINAL",
            (20, 40),
            cv.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 0, 255),
            2,
            cv.LINE_AA
        )


        # Подпись исправленного изображения
        cv.putText(
            corrected_view,
            "CORRECTED",
            (20, 40),
            cv.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 0),
            2,
            cv.LINE_AA
        )


        # Соединяем два изображения горизонтально
        comparison = np.hstack(
            (original_view, corrected_view)
        )


        # Уменьшаем объединённое изображение,
        # чтобы оно помещалось на экране
        comparison = cv.resize(
            comparison,
            None,
            fx=0.5,
            fy=0.5
        )


        cv.imshow(
            "Original | Corrected",
            comparison
        )


        key = cv.waitKey(1) & 0xFF

        if key == ord("q") or key == 27:
            break


finally:
    camera.release()
    cv.destroyAllWindows()

    print("Камера закрыта")
    print("Программа завершена")