---

import os
import cv2
import numpy as np

# --- 1. ВСПОМОГАТЕЛЬНЫЕ ФУНКЦИИ И ОЧИСТКА ---

def extract_dominant_piece_color(cell_bgr):
    """Определение среднего цвета фигуры для разделения Белая/Черная."""
    h, w = cell_bgr.shape[:2]
    center = cell_bgr[int(h * 0.25):int(h * 0.75), int(w * 0.25):int(w * 0.75)]
    return np.mean(center, axis=(0, 1)).astype(float).tolist()


def isolate_piece_silhouette(crop_bgr):
    """
    Создает чистый силуэт фигуры с центрированием по центру масс.
    """
    gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
    
    # Повышение контраста
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)
    
    # Адаптивный порог
    thresh = cv2.adaptiveThreshold(
        enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
        cv2.THRESH_BINARY_INV, 11, 3
    )
    
    # Чистка мелкого шума
    kernel = np.ones((2, 2), np.uint8)
    clean_silhouette = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    
    # Центрирование по центру масс
    M = cv2.moments(clean_silhouette)
    if M["m00"] > 0:
        cX = int(M["m10"] / M["m00"])
        cY = int(M["m01"] / M["m00"])
        
        shift_x = 18 - cX
        shift_y = 18 - cY
        TranslationMatrix = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
        clean_silhouette = cv2.warpAffine(clean_silhouette, TranslationMatrix, (36, 36))

    return clean_silhouette


def get_silhouette_metrics(silhouette):
    """
    Возвращает геометрию фигуры: площадь, высоту и ширину.
    """
    area = np.count_nonzero(silhouette)
    y_indices, x_indices = np.where(silhouette > 0)
    
    if len(y_indices) > 0:
        height = np.max(y_indices) - np.min(y_indices) + 1
        width = np.max(x_indices) - np.min(x_indices) + 1
    else:
        height, width = 0, 0
        
    return area, height, width


def print_board(board_matrix, fen):
    """Печать красивой шахматной доски в консоль."""
    divider = "  +" + "---+" * 8
    print("\n    a   b   c   d   e   f   g   h")
    print(divider)
    for idx, row in enumerate(board_matrix):
        rank = 8 - idx
        row_str = " | ".join([f"{p}" if p != '.' else " " for p in row])
        print(f"{rank} | {row_str} | {rank}")
        print(divider)
    print("    a   b   c   d   e   f   g   h")
    print(f"\nFEN: {fen}\n")


# --- 2. СКАНИРОВАНИЕ ДОСКИ ---

def scan_frame(cell_board, color_database, templates_silhouettes, ALL_PIECES, variance_threshold=12.0):
    board = cv2.resize(cell_board, (400, 400))
    board_matrix = []
    
    for row in range(8):
        row_pieces = []
        for col in range(8):
            cell = board[row * 50:(row + 1) * 50, col * 50:(col + 1) * 50]
            gray_cell = cv2.cvtColor(cell, cv2.COLOR_BGR2GRAY)
            center_crop = cell[7:43, 7:43]
            center_gray = gray_cell[7:43, 7:43]

            # 1. Отсекаем пустую клетку
            if np.std(center_gray) < variance_threshold:
                row_pieces.append('.')
                continue

            # 2. Получаем силуэт и его геометрические параметры
            cell_silhouette = isolate_piece_silhouette(center_crop)
            cell_area, cell_height, cell_width = get_silhouette_metrics(cell_silhouette)

            # 3. Разделение фигуры по цвету (Белая / Черная)
            cell_color_vec = np.array(extract_dominant_piece_color(cell))
            brightness = np.mean(cell_color_vec)
            
            is_white_piece = brightness > 110
            candidate_pieces = [p for p in ALL_PIECES if p.isupper() == is_white_piece]

            best_match, min_difference = '.', float('inf')

            # 4. Сравнение формы, площади и высоты
            for piece_symbol in candidate_pieces:
                if piece_symbol not in templates_silhouettes or piece_symbol not in color_database:
                    continue
                
                tpl_silhouette = templates_silhouettes[piece_symbol]
                tpl_area, tpl_height, tpl_width = get_silhouette_metrics(tpl_silhouette)

                # Попиксельное совпадение формы
                diff = cv2.absdiff(cell_silhouette, tpl_silhouette)
                shape_diff = np.mean(diff)

                # Штраф за разницу в площади (Ладья vs Пешка)
                area_diff = abs(cell_area - tpl_area) / 10.0

                # Штраф за разницу в высоте (Слон vs Пешка)
                height_diff = abs(cell_height - tpl_height) * 2.5

                # Штраф за разницу в цвете
                tpl_color = np.array(color_database[piece_symbol])
                color_dist = np.linalg.norm(cell_color_vec - tpl_color)
                
                # Итоговая ошибка (чем меньше, тем точнее совпадение)
                total_error = shape_diff + area_diff + height_diff + (color_dist / 10.0)

                if total_error < min_difference:
                    min_difference = total_error
                    best_match = piece_symbol
                    
            row_pieces.append(best_match)
        board_matrix.append(row_pieces)

    # 5. Генерация FEN-строки
    fen_rows = []
    for r in board_matrix:
        empty, row_str = 0, ""
        for p in r:
            if p == '.': 
                empty += 1
            else:
                if empty > 0: 
                    row_str += str(empty)
                    empty = 0
                row_str += p
        if empty > 0: 
            row_str += str(empty)
        fen_rows.append(row_str)
        
    return board_matrix, "/".join(fen_rows) + " w - - 0 1"


# --- 3. ГЛАВНЫЙ ИСПОЛНЯЕМЫЙ БЛОК ---

def main():
    TEMPLATES_DIR = "extracted_pieces"
    BOARD_IMAGE_PATH = "board.png"
    
    ALL_PIECES = ['K', 'Q', 'R', 'B', 'N', 'P', 'k', 'q', 'r', 'b', 'n', 'p']
    
    PIECES_CONFIG = {
        "K": "white_K.png", "Q": "white_Q.png", "R": "white_R.png",
        "B": "white_B.png", "N": "white_N.png", "P": "white_P.png",
        "k": "black_k.png", "q": "black_q.png", "r": "black_r.png",
        "b": "black_b.png", "n": "black_n.png", "p": "black_p.png"
    }

    print("[1/3] Загрузка и подготовка шаблонов...")
    if not os.path.exists(TEMPLATES_DIR):
        print(f"[ОШИБКА] Папка '{TEMPLATES_DIR}' не найдена!")
        return

    templates_silhouettes = {}
    color_database = {}

    for sym, filename in PIECES_CONFIG.items():
        path = os.path.join(TEMPLATES_DIR, filename)
        if not os.path.exists(path):
            continue

        img = cv2.imread(path)
        if img is None:
            continue

        img_50 = cv2.resize(img, (50, 50))
        crop_center = img_50[7:43, 7:43]
        
        templates_silhouettes[sym] = isolate_piece_silhouette(crop_center)
        color_database[sym] = extract_dominant_piece_color(img_50)

    print(f"[ОК] Загружено шаблонов: {len(templates_silhouettes)}/12")

    print(f"[2/3] Загрузка картинки доски '{BOARD_IMAGE_PATH}'...")
    if not os.path.exists(BOARD_IMAGE_PATH):
        print(f"[ОШИБКА] Файл '{BOARD_IMAGE_PATH}' не найден!")
        return

    board_img = cv2.imread(BOARD_IMAGE_PATH)
    if board_img is None:
        print("[ОШИБКА] Не удалось прочитать картинку!")
        return

    print("[3/3] Выполнение точного распознавания...")
    matrix, fen = scan_frame(board_img, color_database, templates_silhouettes, ALL_PIECES)
    
    print_board(matrix, fen)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n[КРИТИЧЕСКАЯ ОШИБКА]: {e}")
        import traceback
        traceback.print_exc()
        
    input("\nНажмите Enter, чтобы закрыть окно...")