push a pop implementovany dost dobre ale s malo vystupy KB bez A

import chess
from typing import Iterator
from chess import Move, BB_ALL, Bitboard

# Definice nových figur
AMAZON = 7

class CustomBoard(chess.Board):
    def __init__(self, fen=None):
        self.amazons_white = chess.BB_EMPTY
        self.amazons_black = chess.BB_EMPTY
        self.custom_bishops_white = chess.BB_EMPTY
        self.custom_bishops_black = chess.BB_EMPTY
        super().__init__(None)
        if fen:
            self.set_custom_fen(fen)
        print("Šachovnice inicializována")
        self.debug_amazons()
        self.debug_custom_bishops()

    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.custom_bishops_white = chess.BB_EMPTY
        self.custom_bishops_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() == 'B':
                    if color == chess.WHITE:
                        self.custom_bishops_white |= chess.BB_SQUARES[square]
                    else:
                        self.custom_bishops_black |= chess.BB_SQUARES[square]
                    piece_type = chess.BISHOP
                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
        self.halfmove_clock = int(parts[4])
        self.fullmove_number = int(parts[5])

    def _set_piece_at(self, square, piece_type, color, promoted=None): # Add promoted argument
            super()._set_piece_at(square, piece_type, color, promoted)
            if piece_type is not None:
                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 == chess.BISHOP:
                    if color == chess.WHITE:
                        self.custom_bishops_white |= chess.BB_SQUARES[square]
                    else:
                        self.custom_bishops_black |= chess.BB_SQUARES[square]
            else:
                self.amazons_white &= ~chess.BB_SQUARES[square]
                self.amazons_black &= ~chess.BB_SQUARES[square]
                self.custom_bishops_white &= ~chess.BB_SQUARES[square]
                self.custom_bishops_black &= ~chess.BB_SQUARES[square]

    def generate_pseudo_legal_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboard = BB_ALL) -> Iterator[Move]:
        for move in super().generate_pseudo_legal_moves(from_mask, to_mask):
            if self.piece_type_at(move.from_square) not in [AMAZON, chess.BISHOP]:
                yield move

        our_bishops = self.custom_bishops_white if self.turn == chess.WHITE else self.custom_bishops_black
        our_bishops &= from_mask
        for from_square in chess.scan_forward(our_bishops):
            attacks = self.bishop_attacks(from_square)
            valid_moves = attacks & ~self.occupied & to_mask
            for to_square in chess.scan_forward(valid_moves):
                yield Move(from_square, to_square)

        our_amazons = self.amazons_white if self.turn == chess.WHITE else self.amazons_black
        our_amazons &= from_mask
        for from_square in chess.scan_forward(our_amazons):
            attacks = self.amazon_attacks(from_square)
            valid_moves = attacks & ~self.occupied & to_mask
            for to_square in chess.scan_forward(valid_moves):
                yield Move(from_square, to_square)

    def amazon_attacks(self, square):
        return self.queen_attacks(square) | self.knight_attacks(square)

    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 knight_attacks(self, square):
        return chess.BB_KNIGHT_ATTACKS[square]

    def is_pseudo_legal(self, move):
        piece = self.piece_at(move.from_square)
        if not piece or piece.color != self.turn:
            return False

        if self.is_castling(move):
            return True

        if piece.piece_type == AMAZON:
            return bool(self.amazon_attacks(move.from_square) & chess.BB_SQUARES[move.to_square])
        elif piece.piece_type == chess.BISHOP and ((self.custom_bishops_white & chess.BB_SQUARES[move.from_square]) or (self.custom_bishops_black & chess.BB_SQUARES[move.from_square])):
            return bool(self.bishop_attacks(move.from_square) & chess.BB_SQUARES[move.to_square])
        else:
            return super().is_pseudo_legal(move)

    def attackers_mask(self, color: chess.Color, square: chess.Square) -> Bitboard:
        attackers = super().attackers_mask(color, square)
        our_amazons = self.amazons_white if color == chess.WHITE else self.amazons_black
        for amazon_square in chess.scan_forward(our_amazons):
            if self.amazon_attacks(amazon_square) & chess.BB_SQUARES[square]:
                attackers |= chess.BB_SQUARES[amazon_square]
        return attackers

    def is_check(self):
        return bool(self.attackers_mask(not self.turn, self.king(self.turn)))

    def is_into_check(self, move):
        self.push(move)
        is_check = self.is_check()
        self.pop()
        return is_check

    def is_legal(self, move):
        return self.is_pseudo_legal(move) and not self.is_into_check(move)

    def piece_symbol(self, piece):
        if piece is None:
            return '.'
        if piece.piece_type == AMAZON:
            return 'A' if piece.color == chess.WHITE else 'a'
        return piece.symbol()

    def piece_type_at(self, square):
        if (self.amazons_white | self.amazons_black) & chess.BB_SQUARES[square]:
            return AMAZON
        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
        return super().color_at(square)

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

        super().push(move)

        if piece.piece_type == AMAZON:
            if piece.color == chess.WHITE:
                self.amazons_white &= ~chess.BB_SQUARES[move.from_square]
                self.amazons_white |= chess.BB_SQUARES[move.to_square]
            else:
                self.amazons_black &= ~chess.BB_SQUARES[move.from_square]
                self.amazons_black |= chess.BB_SQUARES[move.to_square]
        elif piece.piece_type == chess.BISHOP:
            if piece.color == chess.WHITE:
                self.custom_bishops_white &= ~chess.BB_SQUARES[move.from_square]
                self.custom_bishops_white |= chess.BB_SQUARES[move.to_square]
            else:
                self.custom_bishops_black &= ~chess.BB_SQUARES[move.from_square]
                self.custom_bishops_black |= chess.BB_SQUARES[move.to_square]

        if captured_piece:
            if captured_piece.piece_type == AMAZON:
                if captured_piece.color == chess.WHITE:
                    self.amazons_white &= ~chess.BB_SQUARES[move.to_square]
                else:
                    self.amazons_black &= ~chess.BB_SQUARES[move.to_square]
            elif captured_piece.piece_type == chess.BISHOP:
                if captured_piece.color == chess.WHITE:
                    self.custom_bishops_white &= ~chess.BB_SQUARES[move.to_square]
                else:
                    self.custom_bishops_black &= ~chess.BB_SQUARES[move.to_square]


    def pop(self):
            move = self.move_stack[-1]
            piece = self.piece_at(move.to_square)

            last_move = super().pop()

            if piece.piece_type == AMAZON:
                if piece.color == chess.WHITE:
                    self.amazons_white &= ~chess.BB_SQUARES[move.to_square]
                    self.amazons_white |= chess.BB_SQUARES[move.from_square]
                else:
                    self.amazons_black &= ~chess.BB_SQUARES[move.to_square]
                    self.amazons_black |= chess.BB_SQUARES[move.from_square]
            elif piece.piece_type == chess.BISHOP:
                if piece.color == chess.WHITE:
                    self.custom_bishops_white &= ~chess.BB_SQUARES[move.to_square]
                    self.custom_bishops_white |= chess.BB_SQUARES[move.from_square]
                else:
                    self.custom_bishops_black &= ~chess.BB_SQUARES[move.to_square]
                    self.custom_bishops_black |= chess.BB_SQUARES[move.from_square]

            # Check if a piece was captured using the piece_at method
            if self.piece_at(move.to_square) is None and move.from_square != move.to_square:
                captured_square = move.to_square
                captured_color = not piece.color
                if self.piece_type_at(captured_square) == AMAZON: # Determine the captured piece type
                    if captured_color == chess.WHITE:
                        self.amazons_white |= chess.BB_SQUARES[captured_square]
                    else:
                        self.amazons_black |= chess.BB_SQUARES[captured_square]
                elif self.piece_type_at(captured_square) == chess.BISHOP:
                    if captured_color == chess.WHITE:
                        self.custom_bishops_white |= chess.BB_SQUARES[captured_square]
                    else:
                        self.custom_bishops_black |= chess.BB_SQUARES[captured_square]

            return last_move

# ... (rest of your code remains the same)

    def debug_amazons(self):
        print(f"Bitboard bílých amazonek: {format(self.amazons_white, '064b')}")
        print(f"Bitboard černých amazonek: {format(self.amazons_black, '064b')}")
        for square in chess.SQUARES:
            if self.amazons_white & chess.BB_SQUARES[square]:
                print(f"Bílá amazonka na {chess.SQUARE_NAMES[square]}")
            if self.amazons_black & chess.BB_SQUARES[square]:
                print(f"Černá amazonka na {chess.SQUARE_NAMES[square]}")

    def debug_custom_bishops(self):
        print(f"Bitboard bílých vlastních střelců: {format(self.custom_bishops_white, '064b')}")
        print(f"Bitboard černých vlastních střelců: {format(self.custom_bishops_black, '064b')}")
        for square in chess.SQUARES:
            if self.custom_bishops_white & chess.BB_SQUARES[square]:
                print(f"Bílý vlastní střelec na {chess.SQUARE_NAMES[square]}")
            if self.custom_bishops_black & chess.BB_SQUARES[square]:
                print(f"Černý vlastní střelec na {chess.SQUARE_NAMES[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)
            if (self.amazons_white & chess.BB_SQUARES[square]):
                symbol = 'A'
            elif (self.amazons_black & chess.BB_SQUARES[square]):
                symbol = 'a'
            else:
                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)

if __name__ == "__main__":
    start_fen = "1a6/3k4/8/8/1B6/8/A7/6K1 w - - 0 1"

    print(f"Vytváření šachovnice s FEN: {start_fen}")
    board = CustomBoard(start_fen)

    print("Současná pozice:")
    print(board)

    print("Generování legálních tahů pro počáteční pozici...")
    legal_moves = list(board.legal_moves)
    print(f"Počet legálních tahů: {len(legal_moves)}")

    print("Legální tahy:")
    for move in legal_moves:
        from_square = chess.SQUARE_NAMES[move.from_square]
        to_square = chess.SQUARE_NAMES[move.to_square]
        piece = board.piece_at(move.from_square)
        print(f"{board.piece_symbol(piece)}: {from_square}-{to_square}")