---

import os
import cv2
import numpy as np


def process_piece_contour(square_img):
    """Очищает фон вокруг фигуры, заливает внутренности и возвращает PNG с прозрачным фоном."""
    # Переводим в формат с альфа-каналом (BGRA)
    bgra = cv2.cvtColor(square_img, cv2.COLOR_BGR2BGRA)

    # Переводим в оттенки серого
    gray = cv2.cvtColor(square_img, cv2.COLOR_BGR2GRAY)

    # Определяем цвет фона по углам клетки
    h, w = gray.shape
    corners = [
        gray[0, 0],
        gray[0, w - 1],
        gray[h - 1, 0],
        gray[h - 1, w - 1],
    ]
    bg_val = int(np.mean(corners))

    # Вычисляем разницу между пикселями и фоном клетки
    diff = cv2.absdiff(gray, bg_val)

    # Применяем порог: всё, что сильно отличается от цвета фона клетки, считается фигурой
    _, thresh = cv2.threshold(diff, 25, 255, cv2.THRESH_BINARY)

    # Убираем мелкий шум и сглаживаем края
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

    # Находим контуры
    contours, _ = cv2.findContours(
        thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    if contours:
        # Ищем самый крупный контур (саму фигуру)
        c = max(contours, key=cv2.contourArea)

        # Создаем маску и заливаем фигуру СПЛОШНЫМ белым цветом (чтобы не было дыр внутри)
        mask = np.zeros_like(gray)
        cv2.drawContours(mask, [c], -1, 255, thickness=cv2.FILLED)

        # Небольшое размытие маски для гладких краев
        mask = cv2.GaussianBlur(mask, (3, 3), 0)

        # Применяем маску к альфа-каналу
        bgra[:, :, 3] = mask
    else:
        bgra[:, :, 3] = 0

    return bgra


def extract_unique_pieces(image_path, output_dir="templates"):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    img = cv2.imread(image_path)
    if img is None:
        print("Ошибка: Не удалось загрузить картинку!")
        return

    height, width, _ = img.shape
    square_h = height / 8.0
    square_w = width / 8.0

    # Словарь только для УНИКАЛЬНЫХ фигур (по 1 позиции на каждый тип)
    # Формат: "Название": (строка, колонка)
    unique_positions = {
        "black_rook": (0, 0),
        "black_knight": (0, 1),
        "black_bishop": (0, 2),
        "black_queen": (0, 3),
        "black_king": (0, 4),
        "black_pawn": (1, 0),
        "white_pawn": (6, 0),
        "white_rook": (7, 0),
        "white_knight": (7, 1),
        "white_bishop": (7, 2),
        "white_queen": (7, 3),
        "white_king": (7, 4),
    }

    print("Начинаем вырезание уникальных фигур...")

    for name, (row, col) in unique_positions.items():
        # Координаты клетки
        y1 = int(row * square_h)
        y2 = int((row + 1) * square_h)
        x1 = int(col * square_w)
        x2 = int((col + 1) * square_w)

        # Кроп клетки
        square_crop = img[y1:y2, x1:x2]

        # Обработка фона и дыр
        clean_piece = process_piece_contour(square_crop)

        # Сохранение (например, templates/white_king.png)
        save_path = os.path.join(output_dir, f"{name}.png")
        cv2.imwrite(save_path, clean_piece)
        print(f"Сохранено: {save_path}")


if __name__ == "__main__":
    # Укажите точное имя вашей картинки
    IMAGE_PATH = "figure.png"
    extract_unique_pieces(IMAGE_PATH)