---

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

        self.piece_symbols = {
            'white_pawn': 'P', 'white_knight': 'N', 'white_bishop': 'B',
            'white_rook': 'R', 'white_queen': 'Q', 'white_king': 'K',
            'black_pawn': 'p', 'black_knight': 'n', 'black_bishop': 'b',
            'black_rook': 'r', 'black_queen': 'q', 'black_king': 'k'
        }

        # Словарная структура: {'piece_name': (gray_template, alpha_mask)}
        self.templates = {}
        
        # Память для диффинга (сравнения изменений)
        self.last_board_img = None
        self.last_board_matrix = None

        self.load_templates()

    def load_templates(self):
        """Загрузка PNG-шаблонов с сохранением альфа-маски"""
        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 len(img_resized.shape) == 3 and img_resized.shape[2] == 4:
                    gray_template = cv2.cvtColor(img_resized[:, :, :3], cv2.COLOR_BGR2GRAY)
                    alpha_mask = img_resized[:, :, 3]  # Берем маску прозрачности
                else:
                    gray_template = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY) if len(img_resized.shape) == 3 else img_resized
                    alpha_mask = np.ones_like(gray_template, dtype=np.uint8) * 255  # Сплошная маска

                self.templates[piece_name] = (gray_template, alpha_mask)

    def remove_highlights(self, cell_img):
        """Убирает подсвеченные области, точки ходов и шах"""
        if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
            return cell_img

        clean_img = cell_img.copy()
        h, w = cell_img.shape[:2]

        # 1. ТОЧКИ ХОДОВ В ЦЕНТРЕ КЛЕТКИ
        center_mask = np.zeros((h, w), dtype=np.uint8)
        cv2.circle(center_mask, (w // 2, h // 2), int(min(h, w) * 0.28), 255, -1)

        gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY)
        mean_val, _ = cv2.meanStdDev(gray)

        dot_mask = cv2.inRange(gray, int(mean_val[0][0] - 50), int(mean_val[0][0] + 50))
        dot_mask = cv2.bitwise_and(dot_mask, center_mask)

        # 2. HSV ЦВЕТОВАЯ МАСКА
        hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)

        mask_green_blue = cv2.inRange(hsv, np.array([22, 10, 30]), np.array([170, 255, 255]))
        mask_red1 = cv2.inRange(hsv, np.array([0, 15, 40]), np.array([21, 255, 255]))
        mask_red2 = cv2.inRange(hsv, np.array([171, 15, 40]), np.array([180, 255, 255]))

        color_mask = mask_green_blue | mask_red1 | mask_red2
        full_mask = cv2.bitwise_or(dot_mask, color_mask)

        if not np.any(full_mask):
            return cell_img

        highlight_pixels = np.sum(full_mask > 0)
        total_pixels = h * w

        if highlight_pixels < (total_pixels * 0.40):
            return cv2.inpaint(cell_img, full_mask, inpaintRadius=5, flags=cv2.INPAINT_TELEA)

        s_channel = hsv[:, :, 1]
        v_channel = hsv[:, :, 2]

        piece_mask = (s_channel < 10) | (v_channel < 40)
        bg_mask = cv2.bitwise_and(full_mask, cv2.bitwise_not(piece_mask.astype(np.uint8) * 255))

        clean_img[bg_mask != 0] = [128, 128, 128]
        return clean_img

    def is_highlighted_cell(self, cell_img):
        """Проверяет, подсвечена ли вся клетка зелёным цветом (фон совершенного хода)"""
        if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
            return False

        hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)
        mask_green = cv2.inRange(hsv, np.array([25, 30, 30]), np.array([85, 255, 255]))

        h, w = cell_img.shape[:2]
        center_mask = np.zeros((h, w), dtype=np.uint8)
        cv2.circle(center_mask, (w // 2, h // 2), int(min(h, w) * 0.25), 255, -1)

        border_mask = cv2.bitwise_not(center_mask)
        green_in_border = cv2.bitwise_and(mask_green, border_mask)

        border_pixels = np.sum(border_mask > 0)
        green_border_pixels = np.sum(green_in_border > 0)

        return (green_border_pixels / border_pixels) > 0.15

    def is_empty_cell(self, cell_img, stddev_threshold=12.0):
        """Проверка клетки на пустоту"""
        gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY) if len(cell_img.shape) == 3 else cell_img
        h, w = gray.shape
        inner_crop = gray[int(h * 0.25):int(h * 0.75), int(w * 0.25):int(w * 0.75)]
        stddev = cv2.meanStdDev(inner_crop)
        return stddev[0][0] < stddev_threshold

    def identify_piece(self, cell_img):
        """Распознавание фигуры с использованием АЛЬФА-МАСОК"""
        if cell_img is None or cell_img.size == 0 or not self.templates:
            return ""

        cleaned_cell = self.remove_highlights(cell_img)

        if self.is_empty_cell(cleaned_cell):
            return ""

        cell_resized = cv2.resize(cleaned_cell, self.win_size)
        gray_cell = cv2.cvtColor(cell_resized, cv2.COLOR_BGR2GRAY) if len(cell_resized.shape) == 3 else cell_resized

        best_piece = ""
        max_corr = -1.0

        for piece_name, (template_img, alpha_mask) in self.templates.items():
            # Метод TM_CCORR_NORMED идеален при передаче альфа-маски!
            res = cv2.matchTemplate(gray_cell, template_img, cv2.TM_CCORR_NORMED, mask=alpha_mask)
            max_val = res[0][0]

            if max_val > max_corr:
                max_corr = max_val
                best_piece = piece_name

        if max_corr > 0.35 and best_piece:
            return self.piece_symbols.get(best_piece, "")

        return ""

    def get_board_state(self, image_input, diff_threshold=6.0, force_full_scan=False):
        """
        Получение матрицы 8х8 с FEN-символами.
        
        :param diff_threshold: Порог изменения пикселей. Чем ниже, тем чувствительнее диффинг.
        :param force_full_scan: Заставить сканировать все 64 клетки (игнорируя память).
        """
        if isinstance(image_input, str):
            if not os.path.exists(image_input):
                print(f"[!] Файл '{image_input}' не найден!")
                return None
            board_img = cv2.imread(image_input)
        else:
            board_img = image_input

        if board_img is None:
            return None

        h, w = board_img.shape[:2]
        cell_h, cell_w = h // 8, w // 8

        # Проверяем, можем ли мы использовать диффинг (запомнен ли прошлый кадр)
        can_use_diff = (
            not force_full_scan and
            self.last_board_img is not None and
            self.last_board_matrix is not None and
            self.last_board_img.shape == board_img.shape
        )

        new_board_matrix = []

        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]

                # Если включен диффинг — сравниваем клетку с прошлым кадром
                if can_use_diff:
                    prev_cell_crop = self.last_board_img[y1:y2, x1:x2]
                    # Вычисляем абсолютную разницу пикселей
                    diff = cv2.absdiff(cell_crop, prev_cell_crop)
                    mean_diff = np.mean(diff)

                    # Если клетка НЕ изменилась, берем прошлое значение!
                    if mean_diff < diff_threshold:
                        row_pieces.append(self.last_board_matrix[row][col])
                        continue

                # Если клетка изменилась (или это первый запуск) — распознаем заново
                symbol = self.identify_piece(cell_crop)
                row_pieces.append(symbol)

            new_board_matrix.append(row_pieces)

        # Сохраняем состояние для следующего кадра
        self.last_board_img = board_img.copy()
        self.last_board_matrix = new_board_matrix

        return new_board_matrix

    def get_board_state_with_color(self, image_input, diff_threshold=6.0, force_full_scan=False):
        """Возвращает матрицу 8х8 и определенный цвет следующего хода"""
        board_matrix = self.get_board_state(image_input, diff_threshold, force_full_scan)
        if not board_matrix:
            return None, None

        if isinstance(image_input, str):
            board_img = cv2.imread(image_input)
        else:
            board_img = image_input

        if board_img is None:
            return board_matrix, None

        h, w = board_img.shape[:2]
        cell_h, cell_w = h // 8, w // 8

        highlighted_cells = []
        for row in range(8):
            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]

                if self.is_highlighted_cell(cell_crop):
                    highlighted_cells.append((row, col))

        detected_next_turn = None

        if len(highlighted_cells) >= 2:
            for (r, c) in highlighted_cells:
                symbol = board_matrix[r][c]
                if symbol:
                    if symbol.isupper():
                        detected_next_turn = 'b'  # Сходили белые -> следующий ход Чёрных
                    elif symbol.islower():
                        detected_next_turn = 'w'  # Сходили чёрные -> следующий ход Белых
                    break

        return board_matrix, detected_next_turn

def _load_image(self, image_input):
        """Вспомогательный метод загрузки изображения."""
        if image_input is None:
            return None
        if isinstance(image_input, str):
            if not os.path.exists(image_input):
                print(f"[!] Файл '{image_input}' не найден!")
                return None
            return cv2.imread(image_input)
        return image_input


# Пример использования
if __name__ == "__main__":
    analyzer = ChessAnalyzer(templates_dir='templates')

    # Первый кадр: полное сканирование всех 64 клеток
    matrix, next_turn = analyzer.get_board_state_with_color("board_step1.png")
    print("Ход 1:", next_turn)

    # Второй кадр (следующий ход): сканируются только изменившиеся клетки
    matrix_step2, next_turn2 = analyzer.get_board_state_with_color("board_step2.png")
    print("Ход 2:", next_turn2)