---

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.template_items = []
        self.load_templates()

        # Кэш последнего анализа
        self.last_board_state = None
        self.last_cell_small_cache = None
        self.last_board_shape = None
        self.last_active_color = None
        self.last_board_image = None
        self.last_board_info = None

        # Настройки скорости/чувствительности
        self.compare_size = (16, 16)
        self.cell_change_threshold = 7.0
        self.highlight_green_low = np.array([22, 20, 25])
        self.highlight_green_high = np.array([95, 255, 255])
        self.highlight_threshold = 0.12
        self.empty_stddev_threshold = 12.0
        self.piece_match_threshold = 0.35

    def load_templates(self):
        """Загрузка PNG-шаблонов"""
        self.templates.clear()
        self.template_items = []

        if not os.path.exists(self.templates_dir):
            print(f"[!] Ошибка: Папка '{self.templates_dir}' не найдена.")
            return

        for filename in os.listdir(self.templates_dir):
            if not filename.lower().endswith('.png'):
                continue

            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, interpolation=cv2.INTER_AREA)

            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.0 - alpha)).astype(np.uint8)
            else:
                if len(img_resized.shape) == 3:
                    composite_gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY)
                else:
                    composite_gray = img_resized

            self.templates[piece_name] = composite_gray
            self.template_items.append((piece_name, composite_gray))

    def _normalize_input(self, image_input):
        """
        Поддерживает:
        - numpy.ndarray
        - str (путь к файлу)
        - dict от crop.get_board_info()
        """
        if image_input is None:
            return None, {}

        if isinstance(image_input, dict):
            board_img = image_input.get("image")
            return board_img, image_input

        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, {}

    def _get_bounds(self, height, width):
        """Границы клеток без накопления ошибки округления."""
        row_bounds = np.linspace(0, height, 9, dtype=int)
        col_bounds = np.linspace(0, width, 9, dtype=int)
        return row_bounds, col_bounds

    def _cell_to_small_gray(self, cell_img):
        """Быстрый уменьшенный отпечаток клетки для сравнения с кэшем."""
        if cell_img is None or cell_img.size == 0:
            return None

        if len(cell_img.shape) == 3:
            gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY)
        else:
            gray = cell_img

        small = cv2.resize(gray, self.compare_size, interpolation=cv2.INTER_AREA)
        return small

    def _cell_change_score(self, prev_small, curr_small):
        """Средняя абсолютная разница двух маленьких отпечатков клетки."""
        if prev_small is None or curr_small is None:
            return float("inf")
        if prev_small.shape != curr_small.shape:
            return float("inf")
        diff = cv2.absdiff(prev_small, curr_small)
        return float(np.mean(diff))

    def _highlight_score(self, cell_img):
        """
        Возвращает долю зелёной подсветки в клетке.
        Используется и для поиска подсветки хода, и для определения активного цвета.
        """
        if cell_img is None or cell_img.size == 0 or len(cell_img.shape) < 3:
            return 0.0

        hsv = cv2.cvtColor(cell_img, cv2.COLOR_BGR2HSV)
        mask_green = cv2.inRange(hsv, self.highlight_green_low, self.highlight_green_high)

        h, w = cell_img.shape[:2]
        if h == 0 or w == 0:
            return 0.0

        # Исключаем центр, чтобы точки возможных ходов не мешали
        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)

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

        full_ratio = green_pixels / float(h * w)
        border_ratio = green_border_pixels / border_pixels if border_pixels > 0 else 0.0

        # Комбинированная оценка: либо подсвечен фон, либо рамка клетки
        return max(full_ratio, border_ratio)

    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)
        mean_val = float(mean_val[0][0])

        low = max(0, int(mean_val - 50))
        high = min(255, int(mean_val + 50))

        dot_mask = cv2.inRange(gray, low, high)
        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 = int(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)
        piece_mask_u8 = (piece_mask.astype(np.uint8) * 255)
        bg_mask = cv2.bitwise_and(full_mask, cv2.bitwise_not(piece_mask_u8))

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

    def is_highlighted_cell(self, cell_img):
        """Проверяет, подсвечена ли клетка зелёным цветом."""
        return self._highlight_score(cell_img) > self.highlight_threshold

    def is_empty_cell(self, cell_img, stddev_threshold=None):
        """Проверка клетки на пустоту."""
        if stddev_threshold is None:
            stddev_threshold = self.empty_stddev_threshold

        if cell_img is None or cell_img.size == 0:
            return True

        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)]

        if inner_crop.size == 0:
            return True

        _, stddev = cv2.meanStdDev(inner_crop)
        return float(stddev[0][0]) < stddev_threshold

    def identify_piece(self, cell_img):
        """Распознавание фигуры."""
        if cell_img is None or cell_img.size == 0 or not self.template_items:
            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, interpolation=cv2.INTER_AREA)
        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 in self.template_items:
            res = cv2.matchTemplate(gray_cell, template_img, cv2.TM_CCOEFF_NORMED)
            max_val = float(res[0][0])

            if max_val > max_corr:
                max_corr = max_val
                best_piece = piece_name

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

        return ""

    def _analyze_active_color(self, board_img, board_matrix, row_bounds, col_bounds):
        """
        Определяет чей сейчас ход по подсветке предыдущего хода.
        Если подсветка не найдена или она похожа на простое выделение фигуры,
        возвращает None.
        """
        if board_img is None or board_matrix is None:
            return None

        highlighted = []

        for row in range(8):
            for col in range(8):
                y1, y2 = row_bounds[row], row_bounds[row + 1]
                x1, x2 = col_bounds[col], col_bounds[col + 1]
                cell_crop = board_img[y1:y2, x1:x2]
                score = self._highlight_score(cell_crop)
                if score > self.highlight_threshold:
                    highlighted.append((score, row, col))

        # Для обычного "выбора фигуры" часто подсвечена только 1 клетка
        if len(highlighted) < 2:
            return None

        highlighted.sort(reverse=True)

        # Берём несколько самых вероятных клеток подсветки.
        # Для рокировки их может быть 4.
        for _, row, col in highlighted[:4]:
            symbol = board_matrix[row][col]
            if not symbol:
                continue

            if symbol.isupper():
                return 'b'  # Сходили белые -> сейчас ход чёрных
            if symbol.islower():
                return 'w'  # Сходили чёрные -> сейчас ход белых

        return None

    def get_board_state(self, image_input):
        """Получение матрицы 8x8 с символами FEN."""
        board_img, meta = self._normalize_input(image_input)
        if board_img is None:
            return None

        h, w = board_img.shape[:2]
        if h < 8 or w < 8:
            return None

        row_bounds, col_bounds = self._get_bounds(h, w)
        board_shape = (h, w)

        # Первый запуск или смена размера доски — полный анализ
        full_rebuild = (
            self.last_board_state is None
            or self.last_cell_small_cache is None
            or self.last_board_shape != board_shape
        )

        current_board = []
        current_small_cache = []

        for row in range(8):
            row_pieces = []
            for col in range(8):
                y1, y2 = row_bounds[row], row_bounds[row + 1]
                x1, x2 = col_bounds[col], col_bounds[col + 1]
                cell_crop = board_img[y1:y2, x1:x2]

                curr_small = self._cell_to_small_gray(cell_crop)
                current_small_cache.append(curr_small)

                if not full_rebuild:
                    idx = row * 8 + col
                    prev_small = self.last_cell_small_cache[idx]
                    prev_symbol = self.last_board_state[row][col] if self.last_board_state else ""

                    change_score = self._cell_change_score(prev_small, curr_small)

                    # Если клетка почти не изменилась, берём старый результат
                    if change_score < self.cell_change_threshold:
                        row_pieces.append(prev_symbol)
                        continue

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

            current_board.append(row_pieces)

        # Определяем активный цвет по подсветке текущей позиции
        detected_active_color = self._analyze_active_color(board_img, current_board, row_bounds, col_bounds)

        # Обновляем кэш
        self.last_board_state = current_board
        self.last_cell_small_cache = current_small_cache
        self.last_board_shape = board_shape
        self.last_active_color = detected_active_color
        self.last_board_image = board_img
        self.last_board_info = meta if isinstance(meta, dict) else None

        return current_board

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

        return board_matrix, self.last_active_color

    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)