---

import os
import cv2
import numpy as np

os.makedirs("Template_pieces", exist_ok=True)

def prepare_templates():
    if not os.path.exists("Start_board.png"):
        return
    img = cv2.imread("Start_board.png")
    if img is None:
        return
    h, w = img.shape[:2]
    
    board_setup = [
        ["r", "n", "b", "q", "k", "b", "n", "r"],
        ["p", "p", "p", "p", "p", "p", "p", "p"],
        ["", "", "", "", "", "", "", ""],
        ["", "", "", "", "", "", "", ""],
        ["", "", "", "", "", "", "", ""],
        ["", "", "", "", "", "", "", ""],
        ["P", "P", "P", "P", "P", "P", "P", "P"],
        ["R", "N", "B", "Q", "K", "B", "N", "R"]
    ]
    
    for r in range(8):
        for c in range(8):
            piece = board_setup[r][c]
            if piece:
                y1 = int(round(r * h / 8))
                y2 = int(round((r + 1) * h / 8))
                x1 = int(round(c * w / 8))
                x2 = int(round((c + 1) * w / 8))
                
                cell = img[y1:y2, x1:x2]
                color_prefix = "white_" if piece.isupper() else "black_"
                filename = f"Template_pieces/{color_prefix}{piece.upper()}.png"
                if not os.path.exists(filename):
                    cv2.imwrite(filename, cell)

prepare_templates()

def capture_board():
    return cv2.imread("board.png")

def recognize_fen():
    board_img = capture_board()
    if board_img is None:
        print("Файл board.png не найден!")
        return

    h, w = board_img.shape[:2]
    
    templates = {}
    for filename in os.listdir("Template_pieces"):
        if filename.endswith(".png"):
            path = os.path.join("Template_pieces", filename)
            templ = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
            templates[filename.split(".")[0]] = templ

    fen_rows = []
    ascii_board = []

    for r in range(8):
        empty_count = 0
        row_fen = ""
        ascii_row = []
        
        y1 = int(round(r * h / 8))
        y2 = int(round((r + 1) * h / 8))
        
        for c in range(8):
            x1 = int(round(c * w / 8))
            x2 = int(round((c + 1) * w / 8))
            
            cell = board_img[y1:y2, x1:x2]
            cell_gray = cv2.cvtColor(cell, cv2.COLOR_BGR2GRAY) if len(cell.shape) == 3 else cell
            
            best_match = None
            best_score = -1.0 # Используем метод корреляции ( TM_CCOEFF_NORMED )
            
            for name, templ in templates.items():
                resized_templ = cv2.resize(templ, (cell_gray.shape[1], cell_gray.shape[0]))
                res = cv2.matchTemplate(cell_gray, resized_templ, cv2.TM_CCOEFF_NORMED)
                _, score, _, _ = cv2.minMaxLoc(res)
                
                if score > best_score:
                    best_score = score
                    best_match = name
            
            # Порог схожести для корреляции (обычно > 0.5-0.6)
            if best_match and best_score > 0.45:
                piece_char = best_match.split("_")[1]
                if best_match.startswith("black"):
                    piece_char = piece_char.lower()
                
                if empty_count > 0:
                    row_fen += str(empty_count)
                    empty_count = 0
                row_fen += piece_char
                ascii_row.append(piece_char)
            else:
                empty_count += 1
                ascii_row.append(".")
                
        if empty_count > 0:
            row_fen += str(empty_count)
        fen_rows.append(row_fen)
        ascii_board.append(ascii_row)

    fen = "/".join(fen_rows) + " w - - 0 1"
    
    print("\n--- Распознанная доска ---")
    for row in ascii_board:
        print(" ".join(row))
    print(f"\nFEN: {fen}")

recognize_fen()

input("\nНажмите Enter для выхода...")