---

import os
import cv2
import numpy as np

def to_binary(gray_img):
    """Преобразует изображение клетки в чистый силуэт фигуры."""
    blurred = cv2.GaussianBlur(gray_img, (3, 3), 0)
    _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    if np.mean(thresh) > 127:
        thresh = cv2.bitwise_not(thresh)
        
    return thresh


def get_piece_color(cell_gray, cell_binary):
    """
    Определяет цвет фигуры (white или black) по средней яркости
    пикселей, которые входят в силуэт фигуры.
    """
    # Берем только те пиксели, где находится сама фигура (маска)
    piece_pixels = cell_gray[cell_binary > 0]
    
    if len(piece_pixels) == 0:
        return "white"

    mean_brightness = np.mean(piece_pixels)
    
    # Порог яркости: 128 — середина между 0 (черный) и 255 (белый)
    return "white" if mean_brightness > 128 else "black"


def load_templates(templates_folder="extracted_pieces"):
    """Загрузка и классификация шаблонов по цвету"""
    symbol_map = {
        "white_K": "K", "white_Q": "Q", "white_R": "R", "white_B": "B", "white_N": "N", "white_P": "P",
        "black_k": "k", "black_q": "q", "black_r": "r", "black_b": "b", "black_n": "n", "black_p": "p",
    }

    templates = {"white": {}, "black": {}}
    if not os.path.exists(templates_folder):
        print(f"[!] Ошибка: Папка '{templates_folder}' не найдена!")
        return templates

    for filename in os.listdir(templates_folder):
        if filename.endswith(".png"):
            name_without_ext = os.path.splitext(filename)[0]
            symbol = symbol_map.get(name_without_ext)
            if not symbol:
                continue

            file_path = os.path.join(templates_folder, filename)
            tpl_img = cv2.imread(file_path)
            
            if tpl_img is not None:
                gray = cv2.cvtColor(tpl_img, cv2.COLOR_BGR2GRAY)
                binary_tpl = to_binary(gray)
                
                # Разделяем шаблоны на белые и чёрные
                color_key = "white" if name_without_ext.startswith("white") else "black"
                templates[color_key][symbol] = binary_tpl

    return templates


def is_square_occupied(cell_gray, std_threshold=15.0):
    """Проверка наличия фигуры на клетке"""
    h, w = cell_gray.shape
    margin_h, margin_w = int(h * 0.2), int(w * 0.2)
    center_crop = cell_gray[margin_h:h-margin_h, margin_w:w-margin_w]
    return np.std(center_crop) > std_threshold


def generate_fen(board_matrix):
    """Генерация строки FEN"""
    fen_rows = []
    for row in board_matrix:
        empty_count = 0
        row_str = ""
        for cell in row:
            if cell == ".":
                empty_count += 1
            else:
                if empty_count > 0:
                    row_str += str(empty_count)
                    empty_count = 0
                row_str += cell
        if empty_count > 0:
            row_str += str(empty_count)
        fen_rows.append(row_str)

    return "/".join(fen_rows) + " w - - 0 1"


def scan_board(
    board_path="board.png",
    templates_folder="extracted_pieces",
    empty_std_threshold=15.0,
    inner_margin_pct=0.06
):
    templates = load_templates(templates_folder)
    if not templates["white"] and not templates["black"]:
        print("[!] Ошибка: Нет шаблонов для сравнения.")
        return

    board_img = cv2.imread(board_path)
    if board_img is None:
        print(f"[!] Ошибка: Файл '{board_path}' не найден.")
        return

    img_h, img_w, _ = board_img.shape
    cell_w = img_w / 8.0
    cell_h = img_h / 8.0

    board_matrix = []

    print("Сканирование фигур и цвета...\n")

    for row in range(8):
        row_symbols = []
        for col in range(8):
            # Координаты клетки
            x1 = int(round(col * cell_w + cell_w * inner_margin_pct))
            y1 = int(round(row * cell_h + cell_h * inner_margin_pct))
            x2 = int(round((col + 1) * cell_w - cell_w * inner_margin_pct))
            y2 = int(round((row + 1) * cell_h - cell_h * inner_margin_pct))

            cell_crop = board_img[y1:y2, x1:x2]
            cell_gray = cv2.cvtColor(cell_crop, cv2.COLOR_BGR2GRAY)

            # 1. Проверяем, есть ли фигура
            if not is_square_occupied(cell_gray, std_threshold=empty_std_threshold):
                row_symbols.append(".")
            else:
                # 2. Получаем силуэт фигуры
                cell_binary = to_binary(cell_gray)

                # 3. Определяем цвет фигуры (white / black)
                detected_color = get_piece_color(cell_gray, cell_binary)

                best_match_symbol = "."
                best_score = -1.0

                # 4. Сравниваем ТОЛЬКО с шаблонами того же цвета!
                target_templates = templates[detected_color]

                for symbol, tpl_binary in target_templates.items():
                    tpl_h, tpl_w = tpl_binary.shape
                    cell_binary_resized = cv2.resize(cell_binary, (tpl_w, tpl_h))

                    res = cv2.matchTemplate(cell_binary_resized, tpl_binary, cv2.TM_CCOEFF_NORMED)
                    _, max_val, _, _ = cv2.minMaxLoc(res)

                    if max_val > best_score:
                        best_score = max_val
                        best_match_symbol = symbol

                row_symbols.append(best_match_symbol)

        board_matrix.append(row_symbols)

    # Вывод доски
    divider = "  +" + "---+" * 8
    print(divider)
    for i, row in enumerate(board_matrix):
        rank_label = 8 - i
        row_content = " | ".join(row)
        print(f"{rank_label} | {row_content} |")
        print(divider)
    print("    a   b   c   d   e   f   g   h\n")

    # Вывод FEN
    fen_string = generate_fen(board_matrix)
    print(f"FEN: {fen_string}\n")


if __name__ == "__main__":
    try:
        scan_board()
    except Exception as e:
        print(f"\n[!] Ошибка: {e}")
    
    input("\nНажмите Enter, чтобы закрыть программу...")