bishops amazons BB good OK

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

class CustomBoard(chess.Board):
    def __init__(self, fen=None):
        self.amazons = chess.BB_EMPTY
        self.custom_bishops = chess.BB_EMPTY
        super().__init__(None)  # Initialize with empty board
        if fen:
            self.set_custom_fen(fen)

    def set_custom_fen(self, fen):
        parts = fen.split()
        board_part = parts[0]
        self.clear()
        self.amazons = chess.BB_EMPTY
        self.custom_bishops = chess.BB_EMPTY
        square = 0
        for c in board_part:
            if c == '/':
                continue
            elif c.isdigit():
                square += int(c)
            else:
                color = chess.WHITE if c.isupper() else chess.BLACK
                piece_type = chess.PIECE_SYMBOLS.index(c.lower()) if c.lower() not in ['a', 'b'] else (AMAZON if c.lower() == 'a' else chess.BISHOP)
                self._set_piece_at(square, chess.Piece(piece_type, color))
                square += 1

        self.turn = chess.WHITE if parts[1] == 'w' else chess.BLACK
        self.set_castling_fen(parts[2])
        self.ep_square = chess.parse_square(parts[3]) if parts[3] != '-' else None
        self.halfmove_clock = int(parts[4]) if len(parts) > 4 else 0
        self.fullmove_number = int(parts[5]) if len(parts) > 5 else 1
        print(f"After setting FEN, amazons bitboard: {self.amazons:064b}")
        print(f"After setting FEN, custom bishops bitboard: {self.custom_bishops:064b}")

    def _set_piece_at(self, square, piece, promoted=False):
        super()._set_piece_at(square, piece, promoted)
        if piece:
            if piece.piece_type == AMAZON:
                self.amazons |= chess.BB_SQUARES[square]
            elif piece.piece_type == chess.BISHOP:
                self.custom_bishops |= chess.BB_SQUARES[square]
        else:
            self.amazons &= ~chess.BB_SQUARES[square]
            self.custom_bishops &= ~chess.BB_SQUARES[square]
        print(f"After _set_piece_at, amazons bitboard: {self.amazons:064b}")
        print(f"After _set_piece_at, custom bishops bitboard: {self.custom_bishops:064b}")

    def piece_type_at(self, square):
        if self.amazons & chess.BB_SQUARES[square]:
            return AMAZON
        if self.custom_bishops & chess.BB_SQUARES[square]:
            return chess.BISHOP
        return super().piece_type_at(square)

    def piece_at(self, square):
        if self.amazons & chess.BB_SQUARES[square]:
            color = bool(self.occupied_co[chess.WHITE] & chess.BB_SQUARES[square])
            return chess.Piece(AMAZON, color)
        if self.custom_bishops & chess.BB_SQUARES[square]:
            color = bool(self.occupied_co[chess.WHITE] & chess.BB_SQUARES[square])
            return chess.Piece(chess.BISHOP, color)
        return super().piece_at(square)

    def generate_pseudo_legal_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboard = BB_ALL) -> Iterator[Move]:
        print("Generating pseudo-legal moves...")
        print(f"Current amazons bitboard: {self.amazons:064b}")
        print(f"Current custom bishops bitboard: {self.custom_bishops:064b}")
        print(f"Current turn: {'White' if self.turn == chess.WHITE else 'Black'}")

        # Generate standard moves
        for move in super().generate_pseudo_legal_moves(from_mask, to_mask):
            print(f"Standard pseudo-legal move: {move}")
            yield move

        # Generate amazon moves
        our_amazons = self.amazons & self.occupied_co[self.turn] & from_mask
        print(f"Our amazons: {our_amazons:064b}")
        for from_square in chess.scan_forward(our_amazons):
            attacks = self.amazon_attacks(from_square)
            print(f"Amazon at {chess.SQUARE_NAMES[from_square]}, attacks: {attacks:064b}")
            for to_square in chess.scan_forward(attacks & ~self.occupied & to_mask):
                move = Move(from_square, to_square)
                print(f"Amazon pseudo-legal move: {move}")
                yield move

        # Generate custom bishop moves
        our_bishops = self.custom_bishops & self.occupied_co[self.turn] & from_mask
        print(f"Our custom bishops: {our_bishops:064b}")
        for from_square in chess.scan_forward(our_bishops):
            attacks = self.bishop_attacks(from_square)
            print(f"Custom bishop at {chess.SQUARE_NAMES[from_square]}, attacks: {attacks:064b}")
            for to_square in chess.scan_forward(attacks & ~self.occupied & to_mask):
                move = Move(from_square, to_square)
                print(f"Custom bishop pseudo-legal move: {move}")
                yield move

    def amazon_attacks(self, square):
        queen_attacks = self.attacks_mask(chess.QUEEN, square, self.occupied)
        knight_attacks = self.attacks_mask(chess.KNIGHT, square, self.occupied)
        print(f"Amazon at {chess.SQUARE_NAMES[square]}:")
        print(f"  Queen attacks: {queen_attacks:064b}")
        print(f"  Knight attacks: {knight_attacks:064b}")
        return queen_attacks | knight_attacks

    def bishop_attacks(self, square):
        return self.attacks_mask(chess.BISHOP, square, self.occupied)

    def is_check(self):
        king = self.king(self.turn)
        if king is None:
            return False

        if super().is_check():
            return True

        opponent_amazons = self.amazons & self.occupied_co[not self.turn]
        for amazon_square in chess.scan_forward(opponent_amazons):
            if self.amazon_attacks(amazon_square) & chess.BB_SQUARES[king]:
                return True

        return False

    def is_legal(self, move: Move) -> bool:
        print(f"Checking legality of move: {move}")
        if not self.is_pseudo_legal(move):
            print("Move is not pseudo-legal")
            return False

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

        if legal:
            print("Move is legal")
        else:
            print("Move is illegal (leaves king in check)")
        return legal

    def push(self, move):
        piece = self.piece_at(move.from_square)
        super().push(move)
        if piece:
            if piece.piece_type == AMAZON:
                self.amazons &= ~chess.BB_SQUARES[move.from_square]
                self.amazons |= chess.BB_SQUARES[move.to_square]
            elif piece.piece_type == chess.BISHOP:
                self.custom_bishops &= ~chess.BB_SQUARES[move.from_square]
                self.custom_bishops |= chess.BB_SQUARES[move.to_square]

    def pop(self):
        move = self.move_stack[-1]
        piece = self.piece_at(move.to_square)
        result = super().pop()
        if piece:
            if piece.piece_type == AMAZON:
                self.amazons &= ~chess.BB_SQUARES[move.to_square]
                self.amazons |= chess.BB_SQUARES[move.from_square]
            elif piece.piece_type == chess.BISHOP:
                self.custom_bishops &= ~chess.BB_SQUARES[move.to_square]
                self.custom_bishops |= chess.BB_SQUARES[move.from_square]
        return result

    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 == chess.BISHOP:
            return 'B' if piece.color == chess.WHITE else 'b'
        return piece.symbol()

    def fen(self):
        fen = []
        empty = 0

        for square in chess.SQUARES_180:
            piece = self.piece_at(square)

            if not piece:
                empty += 1
            else:
                if empty:
                    fen.append(str(empty))
                    empty = 0
                fen.append(self.piece_symbol(piece))

            if chess.BB_SQUARES[square] & chess.BB_FILE_H:
                if empty:
                    fen.append(str(empty))
                    empty = 0

                if square != chess.H1:
                    fen.append("/")

        return " ".join([
            "".join(fen),
            "w" if self.turn == chess.WHITE else "b",
            self.castling_xfen(),
            chess.SQUARE_NAMES[self.ep_square] if self.ep_square is not None else "-",
            str(self.halfmove_clock),
            str(self.fullmove_number)
        ])

@lru_cache(maxsize=None)
def simplify_fen_string(fen):
    parts = fen.split(' ')
    return ' '.join(parts[:4])

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 analyze_positions(start_fen, min_depth=5):
    POZ = {1: start_fen}
    AR = {simplify_fen_string(start_fen): {'depth': 0, 'value': None}}
    N = 1

    start_time = time.time()
    total_time = 0

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

        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)
                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:
                        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:
            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:
                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("\nKonečná pozice:")
    print(f"FEN: {current_fen}")

if __name__ == "__main__":
    start_fen = "1b6/3k4/8/8/8/8/8/6K1 b - - 0 1"

    try:
        print(f"Creating board with FEN: {start_fen}")
        initial_board = CustomBoard(start_fen)

        print("Debugging amazons...")
        print(f"Amazons bitboard: {initial_board.amazons:064b}")
        print(f"Bishops bitboard: {initial_board.bishops:064b}")


        print("Piece positions:")
        for square in chess.SQUARES:
            piece = initial_board.piece_at(square)
            if piece:
                print(f"{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)}")

        print("Legal moves:")
        for move in legal_moves:
            from_square = chess.SQUARE_NAMES[move.from_square]
            to_square = chess.SQUARE_NAMES[move.to_square]
            piece = initial_board.piece_at(move.from_square)
            print(f"{initial_board.piece_symbol(piece)}: {from_square}-{to_square}")

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