---

import glob
import cv2 as cv
import numpy as np


# --------------------------------------------------
# Настройки
# --------------------------------------------------

# --------------------------------------------------
# IMAGES_FOLDER, BOARD_SIZE, SQUARE_SIZE

IMAGES_FOLDER = "calibration_images"
BOARD_SIZE = (9, 6)
SQUARE_SIZE = 0.025

# %% БЛОК 2. Реальные координаты углов шахматной доски

board_points = []

for y in range(BOARD_SIZE[1]):
    for x in range(BOARD_SIZE[0]):
        board_points.append(
            [x * SQUARE_SIZE, y * SQUARE_SIZE, 0]
        )

board_points = np.array(board_points, dtype=np.float32)


# %% БЛОК 3. Ищем углы на фотографиях

image_files = sorted(
    glob.glob(IMAGES_FOLDER + "/*.jpg")
)

object_points = []
image_points = []
image_size = None


for filename in image_files:

    image = cv.imread(filename)
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

    found, corners = cv.findChessboardCorners(
        gray,
        BOARD_SIZE
    )

    if found:
        object_points.append(board_points)
        image_points.append(corners)

        image_size = (
            gray.shape[1],
            gray.shape[0]
        )

        cv.drawChessboardCorners(
            image,
            BOARD_SIZE,
            corners,
            found
        )

        cv.imshow("Found corners", image)
        cv.waitKey(300)

        print("Углы найдены:", filename)

    else:
        print("Углы НЕ найдены:", filename)


cv.destroyAllWindows()

print("Подходящих фотографий:", len(image_points))


# %% БЛОК 4. Калибруем камеру

if len(image_points) == 0:
    print("Нет подходящих фотографий")
    raise SystemExit


result = cv.calibrateCamera(
    object_points,
    image_points,
    image_size,
    None,
    None
)

error = result[0]
camera_matrix = result[1]
distortion = result[2]


print()
print("Ошибка калибровки:", error)

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

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


# %% БЛОК 5. Сохраняем результат

np.savez(
    "camera_calibration.npz",
    camera_matrix=camera_matrix,
    distortion=distortion
)

print("Параметры сохранены")


# %% БЛОК 6. Смотрим изображение до и после исправления

test_image = cv.imread(image_files[0])

corrected_image = cv.undistort(
    test_image,
    camera_matrix,
    distortion
)

comparison = np.hstack(
    (test_image, corrected_image)
)

comparison = cv.resize(
    comparison,
    None,
    fx=0.5,
    fy=0.5
)

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

print("Нажмите Q или ESC в окне изображения")

while True:
    key = cv.waitKey(30) & 0xFF

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

cv.destroyAllWindows()

print("Программа завершена")