uz se zobrazuji kralove Kk! spolu s Aa

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 figur
AMAZON = 7

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

    def set_chess_fen(self, fen):
        super().set_chess_fen(fen)
        self.custom_bishops_white = self.bishops & self.occupied_co[chess.WHITE]
        self.custom_bishops_black = self.bishops & self.occupied_co[chess.BLACK]
        print(f"Po nastavení FEN, bitboard amazonek: {self.amazons:064b}")
        print(f"Po nastavení FEN, bitboard bílých vlastních střelců: {self.custom_bishops_white:064b}")
        print(f"Po nastavení FEN, bitboard černých vlastních střelců: {self.custom_bishops_black:064b}")

    def _set_piece_at(self, square, piece, promoted=False):
        super()._set_piece_at(square, piece, promoted)
        if piece is not None:
            piece_type = piece if isinstance(piece, int) else piece.piece_type
            if piece_type == chess.BISHOP:
                color = self.color_at(square)
                if color == chess.WHITE:
                    self.custom_bishops_white |= chess.BB_SQUARES[square]
                else:
                    self.custom_bishops_black |= chess.BB_SQUARES[square]
            elif piece_type == AMAZON:
                self.amazons |= chess.BB_SQUARES[square]
        else:
            self.custom_bishops_white &= ~chess.BB_SQUARES[square]
            self.custom_bishops_black &= ~chess.BB_SQUARES[square]
            self.amazons &= ~chess.BB_SQUARES[square]

    def set_piece_at(self, square, piece, promoted=False):
        self._set_piece_at(square, piece, promoted)

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

    def generate_pseudo_legal_moves(self, from_mask: Bitboard = BB_ALL, to_mask: Bitboard = BB_ALL) -> Iterator[Move]:
        print("Generování pseudo-legálních tahů...")

        # Generování standardních tahů
        print("Generování standardních tahů...")
        for move in super().generate_pseudo_legal_moves(from_mask, to_mask):
            if self.piece_type_at(move.from_square) != AMAZON:
                print(f"Standardní pseudo-legální tah: {move}")
                yield move

        # Generování tahů vlastních střelců
        print("Generování tahů vlastních střelců...")
        our_bishops = self.custom_bishops_white if self.turn == chess.WHITE else self.custom_bishops_black
        our_bishops &= from_mask
        print(f"Naši střelci: {our_bishops:064b}")
        for from_square in chess.scan_forward(our_bishops):
            print(f"Generování tahů pro střelce na {chess.SQUARE_NAMES[from_square]}")
            attacks = self.bishop_attacks(from_square)
            print(f"Útoky střelce: {attacks:064b}")
            valid_moves = attacks & ~self.occupied & to_mask
            print(f"Platné cílové pole: {valid_moves:064b}")
            for to_square in chess.scan_forward(valid_moves):
                move = Move(from_square, to_square)
                print(f"Pseudo-legální tah vlastního střelce: {move}")
                yield move

        # Generování tahů amazonek
        print("Generování tahů amazonek...")
        our_amazons = self.amazons & self.occupied_co[self.turn] & from_mask
        print(f"Naše amazonky: {our_amazons:064b}")
        for from_square in chess.scan_forward(our_amazons):
            print(f"Generování tahů pro amazonku na {chess.SQUARE_NAMES[from_square]}")
            attacks = self.amazon_attacks(from_square)
            print(f"Útoky amazonky: {attacks:064b}")
            valid_moves = attacks & ~self.occupied & to_mask
            print(f"Platné cílové pole: {valid_moves:064b}")
            for to_square in chess.scan_forward(valid_moves):
                move = Move(from_square, to_square)
                print(f"Pseudo-legální tah amazonky: {move}")
                yield move

        print("Dokončeno generování pseudo-legálních tahů")

    def amazon_attacks(self, square):
        return (self.attacks_mask(chess.QUEEN, square, self.occupied) |
                self.attacks_mask(chess.KNIGHT, square, self.occupied))

    def attacks_mask(self, piece_type, square=None, occupied=None, mask=BB_ALL):
        if square is None or occupied is None:
            return 0

        if piece_type == chess.QUEEN:
            return (chess.BB_DIAG_ATTACKS[square][occupied & chess.BB_DIAG_MASKS[square]] |
                    chess.BB_RANK_ATTACKS[square][occupied] |
                    chess.BB_FILE_ATTACKS[square][occupied]) & mask
        elif piece_type == chess.KNIGHT:
            return chess.BB_KNIGHT_ATTACKS[square] & mask
        elif piece_type == chess.BISHOP:
            return chess.BB_DIAG_ATTACKS[square][occupied & chess.BB_DIAG_MASKS[square]] & mask
        return 0

    def bishop_attacks(self, square):
        attacks = self.attacks_mask(chess.BISHOP, square, self.occupied)
        print(f"Bishop attacks from {chess.SQUARE_NAMES[square]}: {attacks:064b}")
        return attacks

    def is_pseudo_legal(self, move):
        # Kontrola, zda je tah v rámci šachovnice
        if not (chess.A1 <= move.from_square < chess.H8 and chess.A1 <= move.to_square < chess.H8):
            return False

        # Kontrola, zda je na výchozím poli figura
        piece = self.piece_at(move.from_square)
        if not piece:
            return False

        # Kontrola, zda figura patří hráči na tahu
        if piece.color != self.turn:
            return False

        # Kontrola, zda cílové pole není obsazeno vlastní figurou
        if self.color_at(move.to_square) == self.turn:
            return False

        # Speciální kontrola pro vlastní střelce
        if 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])

        # Speciální kontrola pro amazonky
        if piece.piece_type == AMAZON:
            return bool(self.amazon_attacks(move.from_square) & chess.BB_SQUARES[move.to_square])

        # Kontrola pro ostatní figury
        return move in super().generate_pseudo_legal_moves()

    def is_legal(self, move):
        print(f"Kontrola legality tahu: {move}")
        if not self.is_pseudo_legal(move):
            print(f"Tah {move} není pseudo-legální")
            return False
        if self.is_into_check(move):
            print(f"Tah {move} staví vlastního krále do šachu")
            return False
        print(f"Tah {move} je legální")
        return True

    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 debug_amazons(self):
        print(f"Bitboard amazonek: {format(self.amazons, '064b')}")
        for square in chess.SQUARES:
            if self.amazons & chess.BB_SQUARES[square]:
                print(f"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):
        legal_moves = [move for move in self.generate_pseudo_legal_moves() if self.is_legal(move)]
        for move in legal_moves:
            print(f"Legální tah: {move}")
        return legal_moves

    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)

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

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

        print("Ladění amazonek...")
        print(f"Bitboard amazonek: {initial_board.amazons:064b}")
        print(f"Bitboard bílých vlastních střelců: {initial_board.custom_bishops_white:064b}")
        print(f"Bitboard černých vlastních střelců: {initial_board.custom_bishops_black:064b}")

        initial_board.debug_amazons()
        initial_board.debug_custom_bishops()

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

        print("Generování legálních tahů pro počáteční pozici...")
        legal_moves = list(initial_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 = initial_board.piece_at(move.from_square)
            print(f"{initial_board.piece_symbol(piece)}: {from_square}-{to_square}")

    except Exception as e:
        print(f"Došlo k chybě: {str(e)}")
        traceback.print_exc()