---

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
        
        # Соответствие имен файлов символам FEN
        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'
        }
        
        self.templates = {}
        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:
                    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) if len(img_resized.shape) == 3 else img_resized

                self.templates[piece_name] = composite_gray

    def remove_highlights(self, cell_img):
        """
        Убирает подсвеченные области любого типа:
        - Зелёные (обычные ходы / выделения)
        - Сине-фиолетовые (предходы / premoves)
        """
        if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
            return cell_img

        hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)
        
        # 1. Маска зелёного цвета (обычная подсветка)
        lower_green = np.array([30, 40, 40])
        upper_green = np.array([85, 255, 255])
        green_mask = cv2.inRange(hsv, lower_green, upper_green)

        # 2. Маска сине-фиолетового цвета (предходы / premoves)
        lower_blue_purple = np.array([90, 35, 40])
        upper_blue_purple = np.array([160, 255, 255])
        blue_purple_mask = cv2.inRange(hsv, lower_blue_purple, upper_blue_purple)

        # Объединяем маски подсветок
        highlight_mask = cv2.bitwise_or(green_mask, blue_purple_mask)

        if not np.any(highlight_mask):
            return cell_img

        h, w = cell_img.shape[:2]
        highlight_pixel_count = np.sum(highlight_mask > 0)
        total_pixels = h * w

        # --- СЦЕНАРИЙ 1: Небольшие кружки / точки подсказок (< 35% площади клетки) ---
        if highlight_pixel_count < (total_pixels * 0.35):
            # Затягиваем точки фоновой текстурой через inpaint
            clean_img = cv2.inpaint(cell_img, highlight_mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
            return clean_img

        # --- СЦЕНАРИЙ 2: Вся клетка подсвечена цветным тоном ---
        s_channel = hsv[:, :, 1]
        v_channel = hsv[:, :, 2]
        
        # Защищаем саму фигуру: у фигур низкая насыщенность или крайние значения яркости
        piece_mask = (s_channel < 50) | (v_channel < 45) | (v_channel > 215)
        
        # Выделяем только подсвеченный фон вокруг фигуры
        bg_highlight_mask = cv2.bitwise_and(highlight_mask, cv2.bitwise_not(piece_mask.astype(np.uint8) * 255))

        clean_img = cell_img.copy()
        # Заменяем подсвеченный фон на нейтральный серый
        clean_img[bg_highlight_mask != 0] = [128, 128, 128]

        return clean_img

    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 ""

        # 1. Очищаем от подсветок (зелёных, синих, фиолетовых)
        cleaned_cell = self.remove_highlights(cell_img)

        # 2. Проверяем очищенную клетку на пустоту
        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

        # 3. Сравниваем с шаблонами через matchTemplate
        for piece_name, template_img in self.templates.items():
            res = cv2.matchTemplate(gray_cell, template_img, cv2.TM_CCOEFF_NORMED)
            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):
        """Получение матрицы 8x8 со символами FEN"""
        if image_input is None:
            return []

        if isinstance(image_input, str):
            if not os.path.exists(image_input):
                print(f"[!] Файл '{image_input}' не найден!")
                return []
            board_img = cv2.imread(image_input)
        else:
            board_img = image_input

        if board_img is None:
            return []

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

        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]

                symbol = self.identify_piece(cell_crop)
                row_pieces.append(symbol)
            
            board_matrix.append(row_pieces)

        return board_matrix

    def get_fen(self, image_input, active_color='w', castling='KQkq', en_passant='-', halfmove=0, fullmove=1):
        """Преобразует изображение доски в FEN-строку"""
        board_matrix = self.get_board_state(image_input)
        if not board_matrix:
            return ""

        fen_rows = []
        for row in board_matrix:
            empty_count = 0
            row_fen = ""
            for cell in row:
                if cell == "":
                    empty_count += 1
                else:
                    if empty_count > 0:
                        row_fen += str(empty_count)
                        empty_count = 0
                    row_fen += cell
            if empty_count > 0:
                row_fen += str(empty_count)
            fen_rows.append(row_fen)

        board_fen = "/".join(fen_rows)
        return f"{board_fen} {active_color} {castling} {en_passant} {halfmove} {fullmove}"

    def analyze_initial_board(self, image_input):
        return self.get_board_state(image_input)