Загрузка данных


import os
import time
import cv2
import numpy as np
import mss
import keyboard
from colorama import init, Fore, Style

from board_detector import crop_chess_board
from hog_scanner import get_hog_features, extract_dominant_piece_color, scan_frame
from engine_adapter import ChessEngine

# Инициализация colorama для корректной работы цветов в Windows 7
init(autoreset=True)


def select_board_region(sct, monitor):
    """Открывает окно для выделения области мышкой."""
    print("\n-----------------------------------------------------------")
    print(" ИНСТРУКЦИЯ: Выделите мышкой шахматную доску на экране.")
    print(" После выделения нажмите ENTER или ПРОБЕЛ.")
    print("-----------------------------------------------------------\n")

    sct_img = sct.grab(monitor)
    frame = cv2.cvtColor(np.array(sct_img), cv2.COLOR_BGRA2BGR)

    window_name = "Выделите шахматную доску (Enter - подтвердить)"
    r = cv2.selectROI(window_name, frame, showCrosshair=True, fromCenter=False)
    cv2.destroyWindow(window_name)

    x, y, w, h = map(int, r)
    
    if w == 0 or h == 0:
        return monitor

    return {
        "top": monitor["top"] + y,
        "left": monitor["left"] + x,
        "width": w,
        "height": h
    }


def parse_move_to_coords(move_uci, is_black=False):
    """Преобразует 'e2e4' в координаты экрана с учётом поворота доски."""
    if not move_uci or len(move_uci) < 4:
        return None, None
    
    col_map = {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4, 'f': 5, 'g': 6, 'h': 7}
    
    col_from = col_map.get(move_uci[0])
    row_from = 8 - int(move_uci[1])
    
    col_to = col_map.get(move_uci[2])
    row_to = 8 - int(move_uci[3])
    
    # Если играем за черных (доска на экране перевернута), отражаем координаты под экран
    if is_black:
        row_from = 7 - row_from
        col_from = 7 - col_from
        row_to = 7 - row_to
        col_to = 7 - col_to
        
    return (row_from, col_from), (row_to, col_to)


def flip_matrix_180(matrix):
    """Разворачивает матрицу 8x8 на 180 градусов (для игры за черных)."""
    return [row[::-1] for row in matrix[::-1]]


def matrix_to_fen(matrix):
    """Преобразует матрицу 8x8 в первую часть FEN-строки."""
    fen_rows = []
    for row in matrix:
        empty_count = 0
        row_str = ""
        for cell in row:
            if cell == '.' or cell == ' ':
                empty_count += 1
            else:
                if empty_count > 0:
                    row_str += str(empty_count)
                    empty_count = 0
                row_str += cell
        if empty_count > 0:
            row_str += str(empty_count)
        fen_rows.append(row_str)
    return "/".join(fen_rows)


def print_board(board_matrix, fen, best_move=None, evaluation="", engine_status="ВЫКЛ", player_color="Белые (w)", is_black=False):
    os.system('cls' if os.name == 'nt' else 'clear')

    from_square, to_square = parse_move_to_coords(best_move, is_black=is_black)

    # Обозначение колонок в зависимости от ориентации доски
    files_hdr = "    h   g   f   e   d   c   b   a" if is_black else "    a   b   c   d   e   f   g   h"
    divider = "  +" + "---+" * 8
    print(f"\n{files_hdr}")
    print(divider)

    for r_idx, row in enumerate(board_matrix):
        rank = r_idx + 1 if is_black else 8 - r_idx
        row_cells = []
        
        for c_idx, symbol in enumerate(row):
            cell_char = symbol if symbol != '.' else ' '
            
            # Подсветка ходов
            if from_square and (r_idx, c_idx) == from_square:
                cell_str = f"{Fore.RED}{Style.BRIGHT}[{cell_char}]{Style.RESET_ALL}"
            elif to_square and (r_idx, c_idx) == to_square:
                cell_str = f"{Fore.RED}{Style.BRIGHT}*{cell_char}*{Style.RESET_ALL}" if cell_char != ' ' else f"{Fore.RED}{Style.BRIGHT} X {Style.RESET_ALL}"
            else:
                cell_str = f" {cell_char} "
                
            row_cells.append(cell_str)

        print(f"{rank} |" + "|".join(row_cells) + f"| {rank}")
        print(divider)

    print(files_hdr)
    print(f"\nFEN: {fen}")
    print(f"Статус Stockfish: [{engine_status}]")
    print(f"Вы играете за: {Fore.YELLOW}{Style.BRIGHT}{player_color}{Style.RESET_ALL}")
    print("\nГорячие клавиши (работают в ЛЮБОМ окне):")
    print("  • [ S ]  — Включить / Выключить бот")
    print("  • [ C ]  — Сменить цвет (Белые <-> Чёрные)")
    print("  • [ H ]  — Выделить доску мышкой")
    print(f"\nОценка: {evaluation}")
    
    best_move_str = f"{Fore.RED}{Style.BRIGHT}{best_move}{Style.RESET_ALL}" if best_move else "Ожидание..."
    print(f"ЛУЧШИЙ ХОД: {best_move_str}\n")


def main():
    TEMPLATES_DIR = "extracted_pieces"
    ALL_PIECES = ['K', 'Q', 'R', 'B', 'N', 'P', 'k', 'q', 'r', 'b', 'n', 'p']
    
    PIECES_CONFIG = {
        "K": "white_K.png", "Q": "white_Q.png", "R": "white_R.png",
        "B": "white_B.png", "N": "white_N.png", "P": "white_P.png",
        "k": "black_k.png", "q": "black_q.png", "r": "black_r.png",
        "b": "black_b.png", "n": "black_n.png", "p": "black_p.png"
    }

    if not os.path.exists(TEMPLATES_DIR):
        print(f"[ОШИБКА] Папка '{TEMPLATES_DIR}' не найдена!")
        return

    print("[1/2] Загрузка шаблонов...")
    templates_hog = {}
    color_database = {}

    for sym, filename in PIECES_CONFIG.items():
        path = os.path.join(TEMPLATES_DIR, filename)
        if not os.path.exists(path):
            continue

        img = cv2.imread(path)
        if img is None:
            continue

        img_50 = cv2.resize(img, (50, 50))
        crop_center = img_50[7:43, 7:43]
        
        templates_hog[sym] = get_hog_features(crop_center)
        color_database[sym] = extract_dominant_piece_color(img_50)

    print("[2/2] Запуск Stockfish...")
    engine = ChessEngine(engine_path="stockfish.exe", time_limit=0.1)

    sct = mss.mss()
    full_monitor = sct.monitors[1]
    current_roi = full_monitor

    pending_fen = ""
    change_timestamp = 0.0
    stockfish_sent = False
    
    active_color_code = 'w'
    player_color_str = "Белые (w)"

    current_matrix = []
    best_move = None
    evaluation = ""

    print("\n[ЗАПУСК] Нажмите 'S' для включения бота, 'C' для смены цвета.")

    try:
        while True:
            start_time = time.time()

            # 1. Клавиша 'S' (Вкл/Выкл бота)
            if keyboard.is_pressed('s'):
                is_active = engine.toggle()
                stockfish_sent = False
                status_str = "ВКЛ" if is_active else "ВЫКЛ"
                print_board(current_matrix, pending_fen, best_move, evaluation, status_str, player_color_str, is_black=(active_color_code == 'b'))
                time.sleep(0.3)

            # 2. Клавиша 'C' (Смена цвета)
            if keyboard.is_pressed('c'):
                if active_color_code == 'w':
                    active_color_code = 'b'
                    player_color_str = "Чёрные (b)"
                else:
                    active_color_code = 'w'
                    player_color_str = "Белые (w)"
                
                stockfish_sent = False
                status_str = "ВКЛ" if engine.enabled else "ВЫКЛ"
                print_board(current_matrix, pending_fen, best_move, evaluation, status_str, player_color_str, is_black=(active_color_code == 'b'))
                time.sleep(0.3)

            # 3. Клавиша 'H' (Ручная обрезка)
            if keyboard.is_pressed('h'):
                current_roi = select_board_region(sct, full_monitor)
                time.sleep(0.3)

            # 4. Захват кадра и сканирование
            sct_img = sct.grab(current_roi)
            frame = cv2.cvtColor(np.array(sct_img), cv2.COLOR_BGRA2BGR)

            cropped_board = crop_chess_board(frame)
            matrix, _ = scan_frame(cropped_board, color_database, templates_hog, ALL_PIECES)

            is_black = (active_color_code == 'b')

            # Если играем за чёрных — разворачиваем сканированную матрицу под нормальный FEN
            if is_black:
                standard_matrix = flip_matrix_180(matrix)
            else:
                standard_matrix = matrix

            # Формируем правильный FEN
            base_fen = matrix_to_fen(standard_matrix)
            current_fen = f"{base_fen} {active_color_code} - - 0 1"

            current_matrix = matrix  # Матрица для отрисовки на экране

            # 5. Фиксация изменений
            if current_fen != pending_fen:
                pending_fen = current_fen
                change_timestamp = time.time()
                stockfish_sent = False
                status_str = "ВКЛ" if engine.enabled else "ВЫКЛ"
                print_board(current_matrix, current_fen, None, "Ожидание 333 мс...", status_str, player_color_str, is_black=is_black)

            # 6. Отправка FEN в Stockfish
            if not stockfish_sent and (time.time() - change_timestamp) >= 0.333:
                if engine.enabled:
                    best_move, evaluation = engine.analyze_fen(pending_fen)
                else:
                    best_move, evaluation = None, "Движок отключен [Жми S]"
                    
                status_str = "ВКЛ" if engine.enabled else "ВЫКЛ"
                print_board(current_matrix, pending_fen, best_move, evaluation, status_str, player_color_str, is_black=is_black)
                stockfish_sent = True

            elapsed = time.time() - start_time
            time.sleep(max(0.001, 0.1 - elapsed))

    finally:
        engine.close()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n[ИНФО] Работа программы завершена.")
    except Exception as e:
        print(f"\n[КРИТИЧЕСКАЯ ОШИБКА]: {e}")
        import traceback
        traceback.print_exc()
        input("\nНажмите Enter, чтобы закрыть...")