---

import cv2
import numpy as np
import os
import json

TEMPLATES_DIR = "templates"
COLORS_FILE = "templates_colors.json"

PIECES_CONFIG = {
    'r': ('black_r.png', (0, 0)),
    'n': ('black_n.png', (0, 1)),
    'b': ('black_b.png', (0, 2)),
    'q': ('black_q.png', (0, 3)),
    'k': ('black_k.png', (0, 4)),
    'p': ('black_p.png', (1, 0)),
    'R': ('white_R.png', (7, 0)),
    'N': ('white_N.png', (7, 1)),
    'B': ('white_B.png', (7, 2)),
    'Q': ('white_Q.png', (7, 3)),
    'K': ('white_K.png', (7, 4)),
    'P': ('white_P.png', (6, 0))
}

ALL_PIECES = list(PIECES_CONFIG.keys())


# =====================================================================
#  [НАЧАЛО БЛОКА: ВИЗУАЛЬНЫЙ ВЫВОД ДОСКИ 8x8]
# =====================================================================
def print_visual_board(board_matrix):
    print("\n" + "="*43)
    print("        ВИЗУАЛЬНАЯ ДОСКА (СЕТКА 8x8)")
    print("="*43)
    print("    a   b   c   d   e   f   g   h")
    print("  +---------------------------------+")
    for idx, row in enumerate(board_matrix):
        rank = 8 - idx
        row_cells = [f"[{p}]" if p != '.' else " . " for p in row]
        row_str = " ".join(row_cells)
        print(f"{rank} | {row_str} | {rank}")
    print("  +---------------------------------+")
    print("    a   b   c   d   e   f   g   h")
    print("="*43)
# =====================================================================
#  [КОНЕЦ БЛОКА: ВИЗУАЛЬНЫЙ ВЫВОД ДОСКИ 8x8]
# =====================================================================


def crop_board(image_path):
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"Не удалось открыть файл: {image_path}")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blur, 50, 150)

    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    board_cnt = None
    max_area = 0

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 10000:
            peri = cv2.arcLength(cnt, True)
            approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
            if len(approx) == 4 and area > max_area:
                board_cnt = approx
                max_area = area

    if board_cnt is not None:
        pts = board_cnt.reshape(4, 2)
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]
        rect[2] = pts[np.argmax(s)]
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]
        rect[3] = pts[np.argmax(diff)]
        dst = np.array([[0, 0], [399, 0], [399, 399], [0, 399]], dtype="float32")
        M = cv2.getPerspectiveTransform(rect, dst)
        return cv2.warpPerspective(img, M, (400, 400))
    else:
        return cv2.resize(img, (400, 400))


def extract_dominant_piece_color(cell_bgr):
    """Определяет преобладающий RGB-цвет самой фигуры, исключая фон клетки"""
    cell_rgb = cv2.cvtColor(cell_bgr, cv2.COLOR_BGR2RGB)
    h, w, _ = cell_rgb.shape

    # Берём фон из углов
    corners = np.concatenate([
        cell_rgb[0:5, 0:5], cell_rgb[0:5, w-5:w],
        cell_rgb[h-5:h, 0:5], cell_rgb[h-5:h, w-5:w]
    ])
    bg_color = np.median(corners.reshape(-1, 3), axis=0)

    # Пиксели внутри клетки
    inner_pixels = cell_rgb[int(h*0.15):int(h*0.85), int(w*0.15):int(w*0.85)].reshape(-1, 3)

    # Фильтруем пиксели, отличающиеся от цвета фона клетки
    distances = np.linalg.norm(inner_pixels - bg_color, axis=1)
    piece_pixels = inner_pixels[distances > 25]

    if len(piece_pixels) == 0:
        return [128, 128, 128]

    # Наиболее часто встречаемый цвет (медиана тела фигуры)
    dominant_color = np.median(piece_pixels, axis=0)
    return [int(c) for c in dominant_color]


def generate_templates_and_colors(start_image_path="start_board.png"):
    """Создает графические шаблоны и файл templates_colors.json"""
    os.makedirs(TEMPLATES_DIR, exist_ok=True)
    board = crop_board(start_image_path)
    cell_size = 50

    color_database = {}

    for piece_symbol, (filename, (row, col)) in PIECES_CONFIG.items():
        y1, y2 = row * cell_size, (row + 1) * cell_size
        x1, x2 = col * cell_size, (col + 1) * cell_size
        cell_img = board[y1:y2, x1:x2]
        
        # Сохранение картинки-шаблона
        save_path = os.path.join(TEMPLATES_DIR, filename)
        cv2.imwrite(save_path, cell_img)

        # Вычисление и сохранение преобладающего цвета
        dom_color = extract_dominant_piece_color(cell_img)
        color_database[piece_symbol] = dom_color

    # Сохраняем в JSON
    with open(COLORS_FILE, 'w', encoding='utf-8') as f:
        json.dump(color_database, f, indent=4)

    print("✓ Папка шаблонов и файл 'templates_colors.json' созданы!")


def scan_board_with_templates(target_image_path="board.png"):
    # Пересоздаем базы при отсутствии хотя бы одного файла
    if not os.path.exists(TEMPLATES_DIR) or not os.path.exists(COLORS_FILE):
        generate_templates_and_colors("start_board.png")

    # Загружаем базу цветов фигур из JSON
    with open(COLORS_FILE, 'r', encoding='utf-8') as f:
        color_database = json.load(f)

    # Подготавливаем контурные шаблоны
    templates_edges = {}
    for piece_symbol, (filename, _) in PIECES_CONFIG.items():
        t_path = os.path.join(TEMPLATES_DIR, filename)
        t_img = cv2.imread(t_path)
        gray_t = cv2.cvtColor(t_img, cv2.COLOR_BGR2GRAY)
        h, w = gray_t.shape
        inner_t = gray_t[int(h*0.15):int(h*0.85), int(w*0.15):int(w*0.85)]
        templates_edges[piece_symbol] = cv2.Canny(inner_t, 50, 150)

    board = crop_board(target_image_path)
    cell_size = 50
    board_matrix = []

    for row in range(8):
        row_pieces = []
        for col in range(8):
            y1, y2 = row * cell_size, (row + 1) * cell_size
            x1, x2 = col * cell_size, (col + 1) * cell_size
            cell = board[y1:y2, x1:x2]

            # 1. Проверка на пустую клетку по наличию контуров
            gray_cell = cv2.cvtColor(cell, cv2.COLOR_BGR2GRAY)
            h, w = gray_cell.shape
            inner_cell = gray_cell[int(h*0.15):int(h*0.85), int(w*0.15):int(w*0.85)]
            cell_edge = cv2.Canny(inner_cell, 50, 150)

            if np.count_nonzero(cell_edge) < 25:
                row_pieces.append('.')
                continue

            # 2. Извлекаем преобладающий цвет фигуры на текущей клетке
            cell_color = np.array(extract_dominant_piece_color(cell))

            best_match = '.'
            max_combined_score = -999.0

            # 3. Комплексная проверка: ФОРМА + ЦВЕТ
            for piece_symbol in ALL_PIECES:
                # Оценка формы (Canny match)
                t_edge = templates_edges[piece_symbol]
                res = cv2.matchTemplate(cell_edge, t_edge, cv2.TM_CCOEFF_NORMED)
                _, shape_score, _, _ = cv2.minMaxLoc(res)

                # Оценка цвета (дистанция до цвета из шаблона)
                ref_color = np.array(color_database[piece_symbol])
                color_dist = np.linalg.norm(cell_color - ref_color)

                # Итоговый балл: совпадение формы со штрафом за разницу в цвете
                combined_score = shape_score - (color_dist / 100.0)

                if combined_score > max_combined_score:
                    max_combined_score = combined_score
                    best_match = piece_symbol

            row_pieces.append(best_match)

        board_matrix.append(row_pieces)

    # 1. Вывод сетки 8x8
    print_visual_board(board_matrix)

    # 2. Генерация FEN
    fen_rows = []
    for r in board_matrix:
        empty_count = 0
        row_str = ""
        for p in r:
            if p == '.':
                empty_count += 1
            else:
                if empty_count > 0:
                    row_str += str(empty_count)
                    empty_count = 0
                row_str += p
        if empty_count > 0:
            row_str += str(empty_count)
        fen_rows.append(row_str)

    final_fen = "/".join(fen_rows) + " w - - 0 1"
    
    print("FEN-строка:")
    print(final_fen)


if __name__ == "__main__":
    try:
        scan_board_with_templates("board.png")
    except Exception as e:
        print("Ошибка:", e)
    
    input("\nНажми Enter, чтобы закрыть...")