---

import cv2
import numpy as np

# --- НАСТРОЙКИ ---
MODEL_FILE = "chess_pieces.onnx"
IMAGE_FILE = "board.png"

# Карта классов фигур (0 = пустая клетка '.', 1..6 = белые, 7..12 = черные)
CLASSES = ['.', 'P', 'N', 'B', 'R', 'Q', 'K', 'p', 'n', 'b', 'r', 'q', 'k']


def order_points(pts):
    """Сортировка углов доски: верх-лево, верх-право, низ-право, низ-лево"""
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect


def crop_board(image_path):
    """Находит контур шахматной доски и выравнивает ее в квадрат 400x400"""
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"Не удалось открыть файл изображения: {image_path}")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blur, 50, 150)

    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    board_cnt = None
    max_area = 0

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 10000:
            peri = cv2.arcLength(cnt, True)
            approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
            if len(approx) == 4 and area > max_area:
                board_cnt = approx
                max_area = area

    if board_cnt is not None:
        pts = board_cnt.reshape(4, 2)
        rect = order_points(pts)
        dst = np.array([[0, 0], [399, 0], [399, 399], [0, 399]], dtype="float32")
        M = cv2.getPerspectiveTransform(rect, dst)
        return cv2.warpPerspective(img, M, (400, 400))
    else:
        # Если четкий контур не найден, масштабируем изображение целиком
        return cv2.resize(img, (400, 400))


def scan_chess():
    print("Загрузка ONNX модели через OpenCV DNN (совместимо с Win 7)...")
    net = cv2.dnn.readNetFromONNX(MODEL_FILE)
    print("✓ Модель успешно загружена!")

    # 1. Вырезаем и выравниваем доску
    board = crop_board(IMAGE_FILE)

    # 2. Нарезаем доску на 64 клетки (по 50x50 пикселей)
    cells = []
    cell_size = 50
    for row in range(8):
        for col in range(8):
            y1, y2 = row * cell_size, (row + 1) * cell_size
            x1, x2 = col * cell_size, (col + 1) * cell_size
            cell = board[y1:y2, x1:x2]
            cells.append(cell)

    # 3. Автоматически подбираем подходящий размер входа для модели
    possible_sizes = [(64, 64), (128, 128), (224, 224), (32, 32)]
    outputs = None

    for size in possible_sizes:
        try:
            blob = cv2.dnn.blobFromImages(
                cells, 
                scalefactor=1.0/255.0, 
                size=size, 
                swapRB=True, 
                crop=False
            )
            net.setInput(blob)
            outputs = net.forward()
            print(f"✓ Подходящий размер клетки найден: {size[0]}x{size[1]}")
            break
        except cv2.error:
            continue

    if outputs is None:
        raise ValueError("Не удалось подобрать размер входа для вашей ONNX-модели.")

    # 4. Распознавание: определяем индексы наиболее вероятных фигур
    predictions = np.argmax(outputs, axis=1)

    # 5. Сборка FEN-строки
    fen_rows = []
    for r in range(8):
        row_pieces = [CLASSES[idx] for idx in predictions[r*8:(r+1)*8]]
        empty_count = 0
        row_str = ""
        for piece in row_pieces:
            if piece == '.':
                empty_count += 1
            else:
                if empty_count > 0:
                    row_str += str(empty_count)
                    empty_count = 0
                row_str += piece
        if empty_count > 0:
            row_str += str(empty_count)
        fen_rows.append(row_str)

    final_fen = "/".join(fen_rows) + " w - - 0 1"
    
    print("\n-------------------------------------------")
    print("УСПЕХ! Сгенерированный FEN-код позиции:")
    print(final_fen)
    print("-------------------------------------------")


if __name__ == "__main__":
    try:
        scan_chess()
    except Exception as e:
        print("\nПроизошла ошибка:", e)