rychle a spravny pocet pozic

import chess
from typing import Iterator, Optional, Dict, Tuple
from chess import Move, BB_ALL, Bitboard, PieceType, Color
import time
from collections import deque
import threading

# Definice nových figur
AMAZON = 7
CYRIL = 8
EVE = 9

# Rozšíření seznamu PIECE_SYMBOLS
chess.PIECE_SYMBOLS.append('a')
chess.PIECE_SYMBOLS.append('c')
chess.PIECE_SYMBOLS.append('e')

class CustomBoard(chess.Board):
    def __init__(self, fen=None):
        self.amazons_white = chess.BB_EMPTY
        self.amazons_black = chess.BB_EMPTY
        self.cyrils_white = chess.BB_EMPTY
        self.cyrils_black = chess.BB_EMPTY
        self.eves_white = chess.BB_EMPTY
        self.eves_black = chess.BB_EMPTY
        super().__init__(None)
        if fen:
            self.set_custom_fen(fen)

    def clear_square(self, square):
        super()._remove_piece_at(square)
        self.amazons_white &= ~chess.BB_SQUARES[square]
        self.amazons_black &= ~chess.BB_SQUARES[square]
        self.cyrils_white &= ~chess.BB_SQUARES[square]
        self.cyrils_black &= ~chess.BB_SQUARES[square]
        self.eves_white &= ~chess.BB_SQUARES[square]
        self.eves_black &= ~chess.BB_SQUARES[square]

    def set_custom_fen(self, fen):
        parts = fen.split()
        board_part = parts[0]

        self.clear()
        self.amazons_white = chess.BB_EMPTY
        self.amazons_black = chess.BB_EMPTY
        self.cyrils_white = chess.BB_EMPTY
        self.cyrils_black = chess.BB_EMPTY
        self.eves_white = chess.BB_EMPTY
        self.eves_black = chess.BB_EMPTY

        square = 56
        for c in board_part:
            if c == '/':
                square -= 16
            elif c.isdigit():
                square += int(c)
            else:
                color = chess.WHITE if c.isupper() else chess.BLACK
                if c.upper() == 'A':
                    if color == chess.WHITE:
                        self.amazons_white |= chess.BB_SQUARES[square]
                    else:
                        self.amazons_black |= chess.BB_SQUARES[square]
                    piece_type = AMAZON
                elif c.upper() == 'C':
                    if color == chess.WHITE:
                        self.cyrils_white |= chess.BB_SQUARES[square]
                    else:
                        self.cyrils_black |= chess.BB_SQUARES[square]
                    piece_type = CYRIL
                elif c.upper() == 'E':
                    if color == chess.WHITE:
                        self.eves_white |= chess.BB_SQUARES[square]
                    else:
                        self.eves_black |= chess.BB_SQUARES[square]
                    piece_type = EVE
                else:
                    piece_type = chess.PIECE_SYMBOLS.index(c.lower())
                self._set_piece_at(square, piece_type, color)
                square += 1

        self.turn = chess.WHITE if parts[1] == 'w' else chess.BLACK
        self.castling_rights = chess.BB_EMPTY
        if '-' not in parts[2]:
            if 'K' in parts[2]: self.castling_rights |= chess.BB_H1
            if 'Q' in parts[2]: self.castling_rights |= chess.BB_A1
            if 'k' in parts[2]: self.castling_rights |= chess.BB_H8
            if 'q' in parts[2]: self.castling_rights |= chess.BB_A8
        self.ep_square = chess.parse_square(parts[3]) if parts[3] != '-' else None

    def _set_piece_at(self, square: chess.Square, piece_type: PieceType, color: Color) -> None:
        self.clear_square(square)
        super()._set_piece_at(square, piece_type, color)
        if piece_type == AMAZON:
            if color == chess.WHITE:
                self.amazons_white |= chess.BB_SQUARES[square]
            else:
                self.amazons_black |= chess.BB_SQUARES[square]
        elif piece_type == CYRIL:
            if color == chess.WHITE:
                self.cyrils_white |= chess.BB_SQUARES[square]
            else:
                self.cyrils_black |= chess.BB_SQUARES[square]
        elif piece_type == EVE:
            if color == chess.WHITE:
                self.eves_white |= chess.BB_SQUARES[square]
            else:
                self.eves_black |= chess.BB_SQUARES[square]

    def piece_at(self, square: chess.Square) -> Optional[chess.Piece]:
        if self.amazons_white & chess.BB_SQUARES[square]:
            return chess.Piece(AMAZON, chess.WHITE)
        elif self.amazons_black & chess.BB_SQUARES[square]:
            return chess.Piece(AMAZON, chess.BLACK)
        elif self.cyrils_white & chess.BB_SQUARES[square]:
            return chess.Piece(CYRIL, chess.WHITE)
        elif self.cyrils_black & chess.BB_SQUARES[square]:
            return chess.Piece(CYRIL, chess.BLACK)
        elif self.eves_white & chess.BB_SQUARES[square]:
            return chess.Piece(EVE, chess.WHITE)
        elif self.eves_black & chess.BB_SQUARES[square]:
            return chess.Piece(EVE, chess.BLACK)
        return super().piece_at(square)

    def generate_pseudo_legal_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboard = BB_ALL) -> Iterator[Move]:
        our_pieces = self.occupied_co[self.turn]
        if self.turn == chess.WHITE:
            our_amazons = self.amazons_white
            our_cyrils = self.cyrils_white
            our_eves = self.eves_white
        else:
            our_amazons = self.amazons_black
            our_cyrils = self.cyrils_black
            our_eves = self.eves_black

        for from_square in chess.scan_forward(our_amazons & from_mask):
            attacks = self.amazon_attacks(from_square)
            valid_moves = attacks & ~our_pieces & to_mask
            for to_square in chess.scan_forward(valid_moves):
                yield Move(from_square, to_square)

        for from_square in chess.scan_forward(our_cyrils & from_mask):
            attacks = self.cyril_attacks(from_square)
            valid_moves = attacks & ~our_pieces & to_mask
            for to_square in chess.scan_forward(valid_moves):
                yield Move(from_square, to_square)

        for from_square in chess.scan_forward(our_eves & from_mask):
            attacks = self.eve_attacks(from_square)
            valid_moves = attacks & ~our_pieces & to_mask
            for to_square in chess.scan_forward(valid_moves):
                yield Move(from_square, to_square)

        for move in super().generate_pseudo_legal_moves(from_mask, to_mask):
            piece = self.piece_at(move.from_square)
            if piece and piece.piece_type not in [AMAZON, CYRIL, EVE]:
                yield move

    def queen_attacks(self, square):
        return self.bishop_attacks(square) | self.rook_attacks(square)

    def bishop_attacks(self, square):
        return chess.BB_DIAG_ATTACKS[square][self.occupied & chess.BB_DIAG_MASKS[square]]

    def rook_attacks(self, square):
        return (chess.BB_RANK_ATTACKS[square][self.occupied & chess.BB_RANK_MASKS[square]] |
                chess.BB_FILE_ATTACKS[square][self.occupied & chess.BB_FILE_MASKS[square]])

    def amazon_attacks(self, square):
        return self.queen_attacks(square) | chess.BB_KNIGHT_ATTACKS[square]

    def cyril_attacks(self, square):
        return self.rook_attacks(square) | chess.BB_KNIGHT_ATTACKS[square]

    def eve_attacks(self, square):
        return self.bishop_attacks(square) | chess.BB_KNIGHT_ATTACKS[square]

    def is_pseudo_legal(self, move):
        from_square = move.from_square
        to_square = move.to_square
        piece = self.piece_at(from_square)

        if not piece or piece.color != self.turn:
            return False

        if self.occupied_co[self.turn] & chess.BB_SQUARES[to_square]:
            return False

        if self.is_castling(move):
            return True

        if piece.piece_type == AMAZON:
            return bool(self.amazon_attacks(from_square) & chess.BB_SQUARES[to_square])
        elif piece.piece_type == CYRIL:
            return bool(self.cyril_attacks(from_square) & chess.BB_SQUARES[to_square])
        elif piece.piece_type == EVE:
            return bool(self.eve_attacks(from_square) & chess.BB_SQUARES[to_square])
        else:
            return super().is_pseudo_legal(move)

    def is_legal(self, move):
        if not self.is_pseudo_legal(move):
            return False

        from_square = move.from_square
        to_square = move.to_square
        piece = self.piece_at(from_square)
        captured_piece = self.piece_at(to_square)

        self.clear_square(from_square)
        self.clear_square(to_square)
        self._set_piece_at(to_square, piece.piece_type, piece.color)

        king_square = to_square if piece.piece_type == chess.KING else self.king(self.turn)
        is_check, attacker_square = self._is_attacked_by(not self.turn, king_square)

        self.clear_square(to_square)
        self._set_piece_at(from_square, piece.piece_type, piece.color)
        if captured_piece:
            self._set_piece_at(to_square, captured_piece.piece_type, captured_piece.color)

        return not is_check

    def _is_attacked_by(self, color, square):
        attackers = self.attackers(color, square)
        if attackers:
            for attacker_square in chess.scan_forward(attackers):
                return True, attacker_square
        return False, None

    def attackers(self, color, square):
        attackers = chess.BB_EMPTY

        knights = self.knights & self.occupied_co[color]
        attackers |= knights & chess.BB_KNIGHT_ATTACKS[square]

        king = self.kings & self.occupied_co[color]
        attackers |= king & chess.BB_KING_ATTACKS[square]

        pawns = self.pawns & self.occupied_co[color]
        if color == chess.WHITE:
            attackers |= pawns & chess.BB_PAWN_ATTACKS[chess.BLACK][square]
        else:
            attackers |= pawns & chess.BB_PAWN_ATTACKS[chess.WHITE][square]

        queens = self.queens & self.occupied_co[color]
        bishops = (self.bishops | queens) & self.occupied_co[color]
        rooks = (self.rooks | queens) & self.occupied_co[color]

        attackers |= chess.BB_DIAG_ATTACKS[square][self.occupied & chess.BB_DIAG_MASKS[square]] & bishops
        attackers |= (chess.BB_RANK_ATTACKS[square][self.occupied & chess.BB_RANK_MASKS[square]] |
                      chess.BB_FILE_ATTACKS[square][self.occupied & chess.BB_FILE_MASKS[square]]) & rooks

        amazons = self.amazons_white if color == chess.WHITE else self.amazons_black
        for amazon_square in chess.scan_forward(amazons):
            if self.amazon_attacks(amazon_square) & chess.BB_SQUARES[square]:
                attackers |= chess.BB_SQUARES[amazon_square]

        cyrils = self.cyrils_white if color == chess.WHITE else self.cyrils_black
        for cyril_square in chess.scan_forward(cyrils):
            if self.cyril_attacks(cyril_square) & chess.BB_SQUARES[square]:
                attackers |= chess.BB_SQUARES[cyril_square]

        eves = self.eves_white if color == chess.WHITE else self.eves_black
        for eve_square in chess.scan_forward(eves):
            if self.eve_attacks(eve_square) & chess.BB_SQUARES[square]:
                attackers |= chess.BB_SQUARES[eve_square]

        return attackers

    def push(self, move):
        if not self.is_legal(move):
            raise ValueError(f"Move {move} is not legal in position {self.fen()}")

        piece = self.piece_at(move.from_square)
        captured_piece = self.piece_at(move.to_square)

        self.clear_square(move.from_square)
        self.clear_square(move.to_square)
        self._set_piece_at(move.to_square, piece.piece_type, piece.color)

        self.turn = not self.turn

        self.move_stack.append((move, captured_piece))

    def pop(self):
        if not self.move_stack:
            return None

        move, captured_piece = self.move_stack.pop()

        piece = self.piece_at(move.to_square)

        self.clear_square(move.from_square)
        self.clear_square(move.to_square)

        self._set_piece_at(move.from_square, piece.piece_type, piece.color)

        if captured_piece:
            self._set_piece_at(move.to_square, captured_piece.piece_type, captured_piece.color)

        self.turn = not self.turn

        return move

    def is_check(self):
        king_square = self.king(self.turn)
        if king_square is None:
            return False
        return self._is_attacked_by(not self.turn, king_square)[0]

    def is_checkmate(self):
        if not self.is_check():
            return False
        return not any(self.generate_legal_moves())

    def is_stalemate(self):
        if self.is_check():
            return False
        return not any(self.generate_legal_moves())

    def is_insufficient_material(self):
            return (self.pawns | self.rooks | self.queens | self.amazons_white | self.amazons_black |
                    self.cyrils_white | self.cyrils_black | self.eves_white | self.eves_black) == 0 and (
                chess.popcount(self.occupied) <= 3
            )

    def is_game_over(self):
        return self.is_checkmate() or self.is_stalemate() or self.is_insufficient_material()

    def piece_symbol(self, piece):
        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 == EVE:
            return 'E' if piece.color == chess.WHITE else 'e'
        return piece.symbol()

    def piece_type_at(self, square):
        if (self.amazons_white | self.amazons_black) & chess.BB_SQUARES[square]:
            return AMAZON
        if (self.cyrils_white | self.cyrils_black) & chess.BB_SQUARES[square]:
            return CYRIL
        if (self.eves_white | self.eves_black) & chess.BB_SQUARES[square]:
            return EVE
        return super().piece_type_at(square)

    def color_at(self, square):
        if self.amazons_white & chess.BB_SQUARES[square]:
            return chess.WHITE
        if self.amazons_black & chess.BB_SQUARES[square]:
            return chess.BLACK
        if self.cyrils_white & chess.BB_SQUARES[square]:
            return chess.WHITE
        if self.cyrils_black & chess.BB_SQUARES[square]:
            return chess.BLACK
        if self.eves_white & chess.BB_SQUARES[square]:
            return chess.WHITE
        if self.eves_black & chess.BB_SQUARES[square]:
            return chess.BLACK
        return super().color_at(square)

    @property
    def legal_moves(self):
        return [move for move in self.generate_pseudo_legal_moves() if self.is_legal(move)]

    def __str__(self):
        builder = []
        for square in chess.SQUARES_180:
            piece = self.piece_at(square)
            symbol = self.piece_symbol(piece) if piece else '.'
            builder.append(symbol)
            if chess.square_file(square) == 7:
                if square != chess.H1:
                    builder.append('\n')
        return ''.join(builder)

def simplify_fen_string(fen):
    parts = fen.split(' ')
    return ' '.join(parts[:4])  # Zachováváme pouze informace o pozici, barvě na tahu, rošádách a en passant


import chess
from typing import Dict
from collections import deque
import time

def calculate_optimal_moves(start_fen: str) -> Dict[str, int]:
    board = CustomBoard(start_fen)
    AR = {}
    queue = deque([(start_fen, 0)])
    visited = set()

    start_time = time.time()
    current_level = 0
    positions_at_level = 0
    level_start_time = start_time

    while queue:
        fen, depth = queue.popleft()

        if depth > current_level:
            level_time = time.time() - level_start_time
            print(f"Hloubka {current_level}: {positions_at_level} pozic, Čas: {format_time(level_time)}")
            current_level = depth
            positions_at_level = 0
            level_start_time = time.time()

        if fen in visited:
            continue

        visited.add(fen)
        positions_at_level += 1
        board.set_custom_fen(fen)

        if board.is_checkmate():
            AR[fen] = -1000 if board.turn == chess.WHITE else 1000
        elif board.is_stalemate() or board.is_insufficient_material():
            AR[fen] = 1 if board.turn == chess.WHITE else -1
        else:
            legal_moves = list(board.legal_moves)
            for move in legal_moves:
                board.push(move)
                new_fen = simplify_fen_string(board.fen())
                if new_fen not in AR:
                    queue.append((new_fen, depth + 1))
                board.pop()

    # Procházení a aktualizace hodnot
    changed = True
    while changed:
        changed = False
        for fen in list(AR.keys()):
            board.set_custom_fen(fen)
            if board.is_game_over():
                continue

            legal_moves = list(board.legal_moves)
            if not legal_moves:
                continue

            values = []
            for move in legal_moves:
                board.push(move)
                new_fen = simplify_fen_string(board.fen())
                if new_fen in AR:
                    values.append(AR[new_fen])
                board.pop()

            if values:
                if board.turn == chess.WHITE:
                    best_value = max(values)
                    if best_value < 0:
                        new_value = best_value + 1
                    elif best_value > 0:
                        new_value = 1000 - (1000 - best_value + 1)
                    else:
                        new_value = 1
                else:
                    best_value = min(values)
                    if best_value > 0:
                        new_value = best_value - 1
                    elif best_value < 0:
                        new_value = -1000 + (-1000 - best_value - 1)
                    else:
                        new_value = -1

                if new_value != AR[fen]:
                    AR[fen] = new_value
                    changed = True

    total_time = time.time() - start_time
    print(f"\nCelkový čas výpočtu: {format_time(total_time)}")

    return AR

def format_time(seconds):
    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 simplify_fen_string(fen):
    parts = fen.split(' ')
    return ' '.join(parts[:4])

if __name__ == "__main__":
    start_fen = "8/3k4/8/8/8/8/A7/6K1 w - - 0 1"
    AR = calculate_optimal_moves(start_fen)

    # Výpis výsledků
    for fen, hodnota in AR.items():
        F = F + 1
        if F < 50 or hodnota != 0:
            print(f"FEN: {fen}")
            print(f"Hodnota: {hodnota}")
            print()