OK AEC ale moc vystupu

import chess
import time
from functools import lru_cache
import traceback
from typing import Iterator
from chess import Move, BB_ALL, Bitboard, scan_reversed

# Definice nových figurek
AMAZON = 7
CYRIL = 8
EVA = 9

# Nová funkce attacks_mask
def new_attacks_mask(self, piece_type, square, occupied):
    print("new_attacks_mask")
    bb_square = chess.BB_SQUARES[square]

    if piece_type == chess.PAWN:
        color = bool(bb_square & self.occupied_co[chess.WHITE])
        return chess.BB_PAWN_ATTACKS[color][square]
    elif piece_type == chess.KNIGHT:
        return chess.BB_KNIGHT_ATTACKS[square]
    elif piece_type == chess.KING:
        return chess.BB_KING_ATTACKS[square]
    else:
        attacks = 0
        if piece_type == chess.BISHOP or piece_type == chess.QUEEN:
            attacks = chess.BB_DIAG_ATTACKS[square][chess.BB_DIAG_MASKS[square] & occupied]
        if piece_type == chess.ROOK or piece_type == chess.QUEEN:
            attacks |= (chess.BB_RANK_ATTACKS[square][chess.BB_RANK_MASKS[square] & occupied] |
                        chess.BB_FILE_ATTACKS[square][chess.BB_FILE_MASKS[square] & occupied])
        if piece_type == AMAZON:
            attacks = (chess.BB_KNIGHT_ATTACKS[square] |
                       chess.BB_DIAG_ATTACKS[square][chess.BB_DIAG_MASKS[square] & occupied] |
                       chess.BB_RANK_ATTACKS[square][chess.BB_RANK_MASKS[square] & occupied] |
                       chess.BB_FILE_ATTACKS[square][chess.BB_FILE_MASKS[square] & occupied])
        elif piece_type == CYRIL:
            attacks = (chess.BB_KNIGHT_ATTACKS[square] |
                       chess.BB_RANK_ATTACKS[square][chess.BB_RANK_MASKS[square] & occupied] |
                       chess.BB_FILE_ATTACKS[square][chess.BB_FILE_MASKS[square] & occupied])
        elif piece_type == EVA:
            attacks = (chess.BB_KNIGHT_ATTACKS[square] |
                       chess.BB_DIAG_ATTACKS[square][chess.BB_DIAG_MASKS[square] & occupied])
        return attacks

# Přiřazení nové funkce
chess.BaseBoard.attacks_mask = new_attacks_mask

class CustomBoard(chess.Board):
    def __init__(self, fen=None):
        print("__init__")
        super().__init__(fen)
        self._custom_piece_map = {}
        if fen:
            self._custom_piece_map = self._parse_custom_fen(fen)
        else:
            self._custom_piece_map = self._create_custom_piece_map()

        # Inicializace bitboardů pro nové figury
        self.amazons = chess.BB_EMPTY
        self.cyrils = chess.BB_EMPTY
        self.evas = chess.BB_EMPTY
        for square, piece in self._custom_piece_map.items():
            bb = chess.BB_SQUARES[square]
            if piece.piece_type == AMAZON:
                self.amazons |= bb
            elif piece.piece_type == CYRIL:
                self.cyrils |= bb
            elif piece.piece_type == EVA:
                self.evas |= bb
        self.occupied_co[chess.WHITE] |= self.amazons | self.cyrils | self.evas
        self.occupied_co[chess.BLACK] |= self.amazons | self.cyrils | self.evas
        self.occupied |= self.amazons | self.cyrils | self.evas

    def _parse_custom_fen(self, fen):
        print("_parse_custom_fen")
        piece_map = {}
        board_part = fen.split()[0]
        square = 0
        for char in board_part:
            if char.isdigit():
                square += int(char)
            elif char == '/':
                continue
            else:
                color = chess.WHITE if char.isupper() else chess.BLACK
                piece_type = {
                    'A': AMAZON,
                    'C': CYRIL,
                    'E': EVA,
                    'P': chess.PAWN,
                    'N': chess.KNIGHT,
                    'B': chess.BISHOP,
                    'R': chess.ROOK,
                    'Q': chess.QUEEN,
                    'K': chess.KING
                }.get(char.upper(), None)
                if piece_type:
                    piece = chess.Piece(piece_type, color)
                    piece_map[square] = piece
                square += 1
        return piece_map

    def piece_symbol(self, piece):
        print("piece_symbol")
        if piece is None:
            return "."
        if piece.piece_type == AMAZON:
            return "A" if piece.color == chess.WHITE else "a"
        if piece.piece_type == CYRIL:
            return "C" if piece.color == chess.WHITE else "c"
        if piece.piece_type == EVA:
            return "E" if piece.color == chess.WHITE else "e"
        return chess.PIECE_SYMBOLS[piece.piece_type].upper() if piece.color == chess.WHITE else chess.PIECE_SYMBOLS[piece.piece_type]

    def _create_custom_piece_map(self):
        print("_create_custom_piece_map")
        piece_map = {}
        for square, piece in self.piece_map().items():
            piece_map[square] = piece
        return piece_map

    def generate_pseudo_legal_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboard = BB_ALL) -> Iterator[Move]:
        print("generate_pseudo_legal_moves začíná")
        our_pieces = self.occupied_co[self.turn]
        print(f"Na tahu je: {'Bílý' if self.turn == chess.WHITE else 'Černý'}")
        print(f"Počet našich figur: {bin(our_pieces).count('1')}")

        # Generate piece moves.
        non_pawns = our_pieces & ~self.pawns & from_mask
        print(f"Počet figur kromě pěšců: {bin(non_pawns).count('1')}")
        for from_square in scan_reversed(non_pawns):
            piece = self.piece_at(from_square)
            if piece is None or piece.color != self.turn:
                continue
            piece_type = piece.piece_type
            print(f"Generuji tahy pro figuru {piece_type} na poli {chess.SQUARE_NAMES[from_square]}")
            if piece_type == AMAZON:
                moves = (self.attacks_mask(chess.KNIGHT, from_square, self.occupied) |
                         self.attacks_mask(chess.BISHOP, from_square, self.occupied) |
                         self.attacks_mask(chess.ROOK, from_square, self.occupied)) & ~our_pieces & to_mask
            elif piece_type == CYRIL:
                moves = (self.attacks_mask(chess.KNIGHT, from_square, self.occupied) |
                         self.attacks_mask(chess.ROOK, from_square, self.occupied)) & ~our_pieces & to_mask
            elif piece_type == EVA:
                moves = (self.attacks_mask(chess.KNIGHT, from_square, self.occupied) |
                         self.attacks_mask(chess.BISHOP, from_square, self.occupied)) & ~our_pieces & to_mask
            else:
                moves = self.attacks_mask(piece_type, from_square, self.occupied) & ~our_pieces & to_mask
            move_count = 0
            for to_square in scan_reversed(moves):
                move_count += 1
                yield Move(from_square, to_square)
            print(f"Vygenerováno {move_count} tahů pro tuto figuru")

        # Generate castling moves.
        if from_mask & self.kings:
            print("Generuji rošády")
            castling_moves = list(self.generate_castling_moves(from_mask, to_mask))
            print(f"Vygenerováno {len(castling_moves)} rošád")
            yield from castling_moves

        # The remaining moves are all pawn moves.
        pawns = self.pawns & self.occupied_co[self.turn] & from_mask
        print(f"Počet pěšců: {bin(pawns).count('1')}")
        if not pawns:
            print("Žádní pěšci k tahu")
            return

        # Generate pawn moves.
        print("Generuji tahy pěšců")
        pawn_moves = list(self.generate_pseudo_legal_pawn_moves(pawns, to_mask))
        print(f"Vygenerováno {len(pawn_moves)} tahů pěšců")
        yield from pawn_moves

        # Generate pawn captures.
        print("Generuji braní pěšců")
        capturers = pawns
        capture_count = 0
        for from_square in scan_reversed(capturers):
            targets = (
                chess.BB_PAWN_ATTACKS[self.turn][from_square] &
                self.occupied_co[not self.turn] & to_mask)

            for to_square in scan_reversed(targets):
                capture_count += 1
                if chess.square_rank(to_square) in [0, 7]:
                    yield Move(from_square, to_square, chess.QUEEN)
                    yield Move(from_square, to_square, chess.ROOK)
                    yield Move(from_square, to_square, chess.BISHOP)
                    yield Move(from_square, to_square, chess.KNIGHT)
                else:
                    yield Move(from_square, to_square)
        print(f"Vygenerováno {capture_count} braní pěšců")

        print("generate_pseudo_legal_moves končí")

    def _is_attacked_by_new_piece(self, king: chess.Square, attacker_square: chess.Square, piece_type: chess.PieceType) -> bool:
        print(f"_is_attacked_by_new_piece: cíl (král) {chess.SQUARE_NAMES[king]}, útočník {chess.SQUARE_NAMES[attacker_square]}")
        if piece_type == AMAZON:
            return bool(self.attacks_mask(chess.KNIGHT, attacker_square, self.occupied) & chess.BB_SQUARES[king] or
                        self.attacks_mask(chess.BISHOP, attacker_square, self.occupied) & chess.BB_SQUARES[king] or
                        self.attacks_mask(chess.ROOK, attacker_square, self.occupied) & chess.BB_SQUARES[king])
        elif piece_type == CYRIL:
            return bool(self.attacks_mask(chess.KNIGHT, attacker_square, self.occupied) & chess.BB_SQUARES[king] or
                        self.attacks_mask(chess.ROOK, attacker_square, self.occupied) & chess.BB_SQUARES[king])
        elif piece_type == EVA:
            return bool(self.attacks_mask(chess.KNIGHT, attacker_square, self.occupied) & chess.BB_SQUARES[king] or
                        self.attacks_mask(chess.BISHOP, attacker_square, self.occupied) & chess.BB_SQUARES[king])
        return False

    def is_check(self) -> bool:
        print("is_check")
        king = self.king(self.turn)
        if king is None:
            return False

        print(f"Kontrola šachu pro krále na {chess.SQUARE_NAMES[king]}")
        opponent_color = not self.turn
        for attacker_square in chess.SQUARES:
            piece = self.piece_at(attacker_square)
            if piece and piece.color == opponent_color:
                if piece.piece_type in [AMAZON, CYRIL, EVA]:
                    if self._is_attacked_by_new_piece(king, attacker_square, piece.piece_type):
                        print(f"Král je v šachu od {self.piece_symbol(piece)} na poli {chess.SQUARE_NAMES[attacker_square]}")
                        return True
                elif self.attacks_mask(piece.piece_type, attacker_square, self.occupied) & chess.BB_SQUARES[king]:
                    print(f"Král je v šachu od {self.piece_symbol(piece)} na poli {chess.SQUARE_NAMES[attacker_square]}")
                    return True

        print("Král není v šachu")
        return False

    def is_legal(self, move: Move) -> bool:
        print(f"Kontrola legality tahu: {move}")
        if not super().is_pseudo_legal(move):
            print("Tah není pseudo-legální")
            return False

        self.push(move)
        legal = not self.is_check()
        self.pop()

        if legal:
            print("Tah je legální")
        else:
            print("Tah by vystavil vlastního krále šachu")
        return legal

@lru_cache(maxsize=None)
def simplify_fen_string(fen):
    print("simplify_fen_string")
    parts = fen.split(' ')
    simplified_fen = ' '.join(parts[:4] + ["0", "1"])
    return simplified_fen

def print_board(fen):
    print("print_board")
    board = CustomBoard(fen)
    print(board.unicode(invert_color=True))
    print()  # Add an empty line after the board

def format_time(seconds):
    print("format_time")
    hours, remainder = divmod(seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    return f"{int(hours):02d}h {int(minutes):02d}m {int(seconds):02d}s"

def analyze_positions(start_fen, min_depth=5):
    print("analyze_positions začíná")
    print("Initializing analysis...")
    POZ = {1: start_fen}
    AR = {simplify_fen_string(start_fen): {'depth': 0, 'value': None}}
    N = 1

    start_time = time.time()
    total_time = 0

    print("Generování pozic...")
    while True:
        new_positions = False
        for i in range(1, N + 1):
            current_fen = POZ[i]
            board = CustomBoard(current_fen)
            simplified_fen = simplify_fen_string(current_fen)

            if AR[simplified_fen]['depth'] == 0:
                AR[simplified_fen]['depth'] = 1
                legal_moves = list(board.legal_moves)
                print(f"Pozice {i}: {simplified_fen}, Počet legálních tahů: {len(legal_moves)}")
                for move in legal_moves:
                    board.push(move)
                    new_fen = board.fen()
                    simplified_new_fen = simplify_fen_string(new_fen)
                    if simplified_new_fen not in AR:
                        N += 1
                        POZ[N] = new_fen
                        AR[simplified_new_fen] = {'depth': 0, 'value': None}
                        new_positions = True
                    board.pop()

        if not new_positions:
            break

    print(f"Celkový počet pozic: {N}")

    # Inicializace koncových pozic
    terminal_positions = 0
    for i in range(1, N + 1):
        current_fen = POZ[i]
        board = CustomBoard(current_fen)
        simplified_fen = simplify_fen_string(current_fen)

        if board.is_checkmate():
            AR[simplified_fen]['value'] = -1000 if board.turn == chess.WHITE else 1000
            AR[simplified_fen]['depth'] = min_depth
            terminal_positions += 1
        elif board.is_stalemate() or board.is_insufficient_material() or board.is_seventyfive_moves() or board.is_fivefold_repetition():
            AR[simplified_fen]['value'] = 0
            AR[simplified_fen]['depth'] = min_depth
            terminal_positions += 1

    print(f"Počet koncových pozic: {terminal_positions}")

    max_depth = 0
    while max_depth < min_depth:
        max_depth += 1
        level_start_time = time.time()
        print(f"Výpočet v hloubce {max_depth}")

        changes = False
        positions_evaluated = 0
        for i in range(1, N + 1):
            current_fen = POZ[i]
            board = CustomBoard(current_fen)
            simplified_fen = simplify_fen_string(current_fen)
            if AR[simplified_fen]['depth'] < max_depth and AR[simplified_fen]['value'] is None:
                positions_evaluated += 1
                best_value = -2000 if board.turn == chess.WHITE else 2000
                all_moves_evaluated = True
                legal_moves = list(board.legal_moves)
              ##  print(f"Pozice {i}: {simplified_fen}, Počet legálních tahů: {len(legal_moves)}")
                for move in legal_moves:
                    board.push(move)
                    next_fen = board.fen()
                    simplified_next_fen = simplify_fen_string(next_fen)
                    if simplified_next_fen not in AR:
                        print(f"Warning: Position not found in AR: {simplified_next_fen}")
                        all_moves_evaluated = False
                        break
                    next_value = AR[simplified_next_fen]['value']
                    if next_value is None:
                        all_moves_evaluated = False
                        break
                    if board.turn == chess.WHITE:
                        best_value = max(best_value, -next_value)
                    else:
                        best_value = min(best_value, -next_value)
                    board.pop()

                if all_moves_evaluated:
                    AR[simplified_fen]['value'] = best_value
                    AR[simplified_fen]['depth'] = max_depth
                    changes = True

        level_end_time = time.time()
        level_elapsed_time = level_end_time - level_start_time
        total_time += level_elapsed_time

        formatted_total_time = format_time(total_time)
        formatted_level_time = format_time(level_elapsed_time)

        print(f"Hloubka {max_depth} dokončena")
        print(f"Pozice vyhodnoceny: {positions_evaluated}")
        print(f"Změny provedeny: {'Ano' if changes else 'Ne'}")
        print(f"Čas: {formatted_total_time} / {formatted_level_time}")

    print(f"Analýza dokončena do hloubky {max_depth}")
    print(f"Celkový čas: {format_time(total_time)}")

    current_fen = start_fen
    optimal_moves = []
    while True:
        board = CustomBoard(current_fen)
        simplified_fen = simplify_fen_string(current_fen)
        if simplified_fen not in AR:
            print(f"Warning: Position not found in AR: {simplified_fen}")
            break
        current_value = AR[simplified_fen]['value']

        if current_value in [-1000, 0, 1000] or current_value is None:
            break

        best_move = None
        for move in board.legal_moves:
            board.push(move)
            next_fen = board.fen()
            simplified_next_fen = simplify_fen_string(next_fen)
            if simplified_next_fen not in AR:
                print(f"Warning: Position not found in AR: {simplified_next_fen}")
                continue
            next_value = AR[simplified_next_fen]['value']
            if next_value is not None and next_value == -current_value:
                best_move = move
                break
            board.pop()

        if best_move is None:
            break

        optimal_moves.append((current_fen, best_move))
        board.push(best_move)
        current_fen = board.fen()

    print("\nOptimální tahy:")
    for fen, move in optimal_moves:
        board = CustomBoard(fen)
        print(f"{board.fullmove_number}. {'Bílý' if board.turn == chess.WHITE else 'Černý'}: {move}")
        print_board(fen)
        print()

    print("Konečná pozice:")
    print_board(current_fen)

    print("analyze_positions končí")

if __name__ == "__main__":
    print("Program začíná")
    start_fen = "1a6/3k4/8/8/8/8/8/6K1 b - - 0 1"
    print("Starting analysis...")
    print("Initial position:")
    print_board(start_fen)

    try:
        print("Creating initial board...")
        initial_board = CustomBoard(start_fen)
        print("Initial board created successfully.")

        print("Piece map:")
        for square, piece in initial_board._custom_piece_map.items():
            print(f"Square {chess.SQUARE_NAMES[square]}: {initial_board.piece_symbol(piece)}")

        print("Generating legal moves for initial position...")
        legal_moves = list(initial_board.legal_moves)
        print(f"Number of legal moves: {len(legal_moves)}")
        for move in legal_moves:
            print(f"Move: {move}")

        print("Starting position analysis...")
        analyze_positions(start_fen, min_depth=5)
    except Exception as e:
        print(f"An error occurred: {str(e)}")
        traceback.print_exc()

    print("Program končí")