---

import cv2
import numpy as np
import os

class ChessAnalyzer:
    def __init__(self, templates_dir='templates', win_size=(64, 64)):
        self.templates_dir = templates_dir
        self.win_size = win_size
        
        # Настройка HOG
        self.hog = cv2.HOGDescriptor(
            _winSize=self.win_size,
            _blockSize=(16, 16),
            _blockStride=(8, 8),
            _cellSize=(8, 8),
            _nbins=9
        )
        
        # Словарь коротких обозначений для красоты в таблице
        self.piece_symbols = {
            'white_pawn': 'wP', 'white_knight': 'wN', 'white_bishop': 'wB',
            'white_rook': 'wR', 'white_queen': 'wQ', 'white_king': 'wK',
            'black_pawn': 'bP', 'black_knight': 'bN', 'black_bishop': 'bB',
            'black_rook': 'bR', 'black_queen': 'bQ', 'black_king': 'bK'
        }
        
        self.templates = {}
        self.load_templates()

    def load_templates(self):
        """Загрузка PNG-шаблонов фигур из папки templates"""
        if not os.path.exists(self.templates_dir):
            print(f"[!] Ошибка: Папка '{self.templates_dir}' не найдена.")
            return

        for filename in os.listdir(self.templates_dir):
            if filename.lower().endswith('.png'):
                piece_name = os.path.splitext(filename)[0]
                path = os.path.join(self.templates_dir, filename)
                
                img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
                if img is None:
                    continue

                img_resized = cv2.resize(img, self.win_size)

                # Удаление альфа-канала с заполнением серым фоном
                if img_resized.shape[2] == 4:
                    alpha = img_resized[:, :, 3] / 255.0
                    gray_piece = cv2.cvtColor(img_resized[:, :, :3], cv2.COLOR_BGR2GRAY)
                    bg = np.ones_like(gray_piece, dtype=np.uint8) * 128
                    composite_gray = (gray_piece * alpha + bg * (1 - alpha)).astype(np.uint8)
                else:
                    composite_gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY)

                hog_feat = self.hog.compute(composite_gray)
                self.templates[piece_name] = hog_feat

    def identify_piece(self, cell_img):
        """Определяет фигуру на одной клетке через HOG"""
        if cell_img is None or cell_img.size == 0 or not self.templates:
            return "  "

        cell_resized = cv2.resize(cell_img, self.win_size)
        gray = cv2.cvtColor(cell_resized, cv2.COLOR_BGR2GRAY) if len(cell_resized.shape) == 3 else cell_resized
        cell_hog = self.hog.compute(gray)
        
        best_piece = "  "
        min_dist = float('inf')

        for piece_name, template_hog in self.templates.items():
            dist = np.linalg.norm(cell_hog - template_hog)
            if dist < min_dist:
                min_dist = dist
                best_piece = piece_name

        # Если уверенность хорошая — возвращаем двухбуквенный символ, иначе пустую клетку
        # Порог ( threshold ) можно подкорректировать при необходимости
        if min_dist < 12.0:
            return self.piece_symbols.get(best_piece, "  ")
        return "  "

    def process_board_image(self, image_path='aboard1.png', output_txt='test.txt'):
        """Нарезает изображение доски на 64 клетки, распознает их и сохраняет в test.txt"""
        if not os.path.exists(image_path):
            print(f"[!] Файл {image_path} не найден!")
            return

        board_img = cv2.imread(image_path)
        h, w = board_img.shape[:2]
        
        # Высота и ширина одной клетки
        cell_h = h // 8
        cell_w = w // 8

        rows = ["8", "7", "6", "5", "4", "3", "2", "1"]
        cols = ["a", "b", "c", "d", "e", "f", "g", "h"]
        
        grid_data = []

        # Проходим по всем 8 строкам и 8 столбцам
        for row in range(8):
            row_pieces = []
            for col in range(8):
                # Вырезаем клетку из доски
                y1, y2 = row * cell_h, (row + 1) * cell_h
                x1, x2 = col * cell_w, (col + 1) * cell_w
                cell_crop = board_img[y1:y2, x1:x2]

                # Распознаем фигуру
                symbol = self.identify_piece(cell_crop)
                row_pieces.append(symbol)
            
            grid_data.append(row_pieces)

        # Записываем результат в test.txt
        with open(output_txt, 'w', encoding='utf-8') as f:
            f.write("+---+----+----+----+----+----+----+----+----+\n")
            f.write("|   |  a |  b |  c |  d |  e |  f |  g |  h |\n")
            f.write("+---+----+----+----+----+----+----+----+----+\n")
            for r_label, row_list in zip(rows, grid_data):
                row_str = f"| {r_label} | " + " | ".join([f"{p:2s}" for p in row_list]) + " |\n"
                f.write(row_str)
                f.write("+---+----+----+----+----+----+----+----+----+\n")

        print(f"[✓] Успешно! Результат сохранен в {output_txt}")


# Если запускаем analysis.py напрямую:
if __name__ == '__main__':
    analyzer = ChessAnalyzer()
    analyzer.process_board_image('aboard1.png', 'test.txt')