ultimate very good amazonka

from typing import List, Iterable, Tuple

# Definice bitboardu a dalších typů
Bitboard = int
Square = int

# Konstanty
BB_EMPTY = 0
BB_ALL = (1 << 64) - 1
SQUARES = list(range(64))
BB_SQUARES = [1 << sq for sq in SQUARES]

# Pomocné funkce
def square_distance(sq1: Square, sq2: Square) -> int:
    file1, rank1 = sq1 % 8, sq1 // 8
    file2, rank2 = sq2 % 8, sq2 // 8
    return max(abs(file1 - file2), abs(rank1 - rank2))

# Funkce pro generování útoků
def _sliding_attacks(square: Square, occupied: Bitboard, deltas: Iterable[int]) -> Bitboard:
    attacks = BB_EMPTY

    for delta in deltas:
        sq = square

        while True:
            sq += delta
            if not (0 <= sq < 64) or square_distance(sq, sq - delta) > 2:
                break

            attacks |= BB_SQUARES[sq]

            if occupied & BB_SQUARES[sq]:
                break

    return attacks

def _step_attacks(square: Square, deltas: Iterable[int]) -> Bitboard:
    attacks = BB_EMPTY
    for delta in deltas:
        sq = square + delta
        if 0 <= sq < 64 and square_distance(sq, square) <= 2:
            attacks |= BB_SQUARES[sq]
    return attacks

# Funkce pro generování tabulek útoků
def _attack_table(deltas: List[int]) -> Tuple[List[Bitboard], List[Bitboard]]:
    masks = []
    attacks = []

    for sq in SQUARES:
        mask = 0
        attack = _sliding_attacks(sq, BB_EMPTY, deltas)
        for i in range(64):
            if attack & BB_SQUARES[i]:
                mask |= BB_SQUARES[i]
        masks.append(mask)
        attacks.append(attack)

    return masks, attacks

# Útoky pro KNIGHTa a BISHOPa
KNIGHT_DELTAS = [-17, -15, -10, -6, 6, 10, 15, 17]
BISHOP_DELTAS = [-9, -7, 7, 9]

# Získání masky a útoků pro KNIGHTa a BISHOPa
BB_KNIGHT_MASKS, BB_KNIGHT_ATTACKS = _attack_table(KNIGHT_DELTAS)
BB_BISHOP_MASKS, BB_BISHOP_ATTACKS = _attack_table(BISHOP_DELTAS)

# Kombinované útoky pro figurku AMAZON
BB_AMAZON_ATTACKS = [knight | bishop for knight, bishop in zip(BB_KNIGHT_ATTACKS, BB_BISHOP_ATTACKS)]

# Funkce pro tisk šachovnice v ASCII formátu s FEN stringem
def print_board_with_fen(fen: str) -> None:
    parts = fen.split()
    board_part = parts[0]
    rows = board_part.split('/')
    for row in rows:
        row_str = ''
        for char in row:
            if char.isdigit():
                row_str += '.' * int(char)
            else:
                row_str += char
        print(row_str)

# Funkce pro parsování FEN stringu a získání pozice Amazonky
def parse_fen(fen: str) -> Tuple[List[Square], List[Square], str, Bitboard, Bitboard]:
    parts = fen.split()
    board_part = parts[0]
    turn = parts[1]
    rows = board_part.split('/')
    white_amazons = []
    black_amazons = []
    white_pieces = BB_EMPTY
    black_pieces = BB_EMPTY
    for rank in range(8):
        row = rows[rank]
        file = 0
        for char in row:
            if char.isdigit():
                file += int(char)
            elif char == 'A':
                white_amazons.append((7 - rank) * 8 + file)
                white_pieces |= BB_SQUARES[(7 - rank) * 8 + file]
                file += 1
            elif char == 'a':
                black_amazons.append((7 - rank) * 8 + file)
                black_pieces |= BB_SQUARES[(7 - rank) * 8 + file]
                file += 1
            elif char.isupper():
                white_pieces |= BB_SQUARES[(7 - rank) * 8 + file]
                file += 1
            elif char.islower():
                black_pieces |= BB_SQUARES[(7 - rank) * 8 + file]
                file += 1
    return white_amazons, black_amazons, turn, white_pieces, black_pieces

# Funkce pro generování možných tahů Amazonky
def generate_amazon_moves(amazon_squares: List[Square], white_pieces: Bitboard, black_pieces: Bitboard, turn: str) -> List[Tuple[Square, Square]]:
    moves = []
    occupied = white_pieces | black_pieces
    for amazon_square in amazon_squares:
        attacks = BB_AMAZON_ATTACKS[amazon_square]
        for move in SQUARES:
            if attacks & BB_SQUARES[move]:
                if turn == 'w':
                    if not (occupied & BB_SQUARES[move] and (white_pieces & BB_SQUARES[move])):
                        moves.append((amazon_square, move))
                else:
                    if not (occupied & BB_SQUARES[move] and (black_pieces & BB_SQUARES[move])):
                        moves.append((amazon_square, move))
    return moves

# Funkce pro generování SAN notace
def generate_san_notation(amazon_square: Square, move_square: Square, is_white: bool) -> str:
    file_from = chr((amazon_square % 8) + ord('a'))
    rank_from = str(8 - (amazon_square // 8))
    file_to = chr((move_square % 8) + ord('a'))
    rank_to = str(8 - (move_square // 8))
    piece = 'A' if is_white else 'a'
    return f"{piece}{file_from}{rank_from}-{file_to}{rank_to}"

# Funkce pro aktualizaci FEN stringu s novou pozicí Amazonky
def update_fen_with_amazon(fen: str, amazon_square: Square, move_square: Square, is_white: bool) -> str:
    parts = fen.split()
    board_part = parts[0]
    rows = board_part.split('/')
    new_rows = []

    for rank in range(8):
        row = rows[rank]
        new_row = ''
        file = 0
        for char in row:
            if char.isdigit():
                empty_squares = int(char)
                for _ in range(empty_squares):
                    new_row += '1'
                    file += 1
            else:
                if (7 - rank) * 8 + file == amazon_square:
                    new_row += '1'
                else:
                    new_row += char
                file += 1
        new_rows.append(new_row)

    new_board_part = '/'.join(new_rows)
    new_fen = f"{new_board_part} {'b' if is_white else 'w'} - - 0 1"

    # Použití regulárních výrazů pro zkrácení sekvencí '1'
    new_fen = re.sub(r'1{8}', '8', new_fen)
    new_fen = re.sub(r'1{7}', '7', new_fen)
    new_fen = re.sub(r'1{6}', '6', new_fen)
    new_fen = re.sub(r'1{5}', '5', new_fen)
    new_fen = re.sub(r'1{4}', '4', new_fen)
    new_fen = re.sub(r'1{3}', '3', new_fen)
    new_fen = re.sub(r'1{2}', '2', new_fen)
    new_fen = re.sub(r'1', '', new_fen)

    # Update the Amazon position in the new FEN
    new_rows = new_fen.split('/')

    for rank in range(8):
        new_row = ''
        file = 0
        for char in new_rows[rank]:
            if char.isdigit():
                empty_squares = int(char)
                for _ in range(empty_squares):
                    if (7 - rank) * 8 + file == move_square:
                        new_row += 'A' if is_white else 'a'
                    else:
                        new_row += '1'
                    file += 1
            else:
                if (7 - rank) * 8 + file == move_square:
                    new_row += 'A' if is_white else 'a'
                else:
                    new_row += char
                file += 1
        new_rows[rank] = new_row

    new_fen = '/'.join(new_rows)
    new_fen = re.sub(r'1{8}', '8', new_fen)
    new_fen = re.sub(r'1{7}', '7', new_fen)
    new_fen = re.sub(r'1{6}', '6', new_fen)
    new_fen = re.sub(r'1{5}', '5', new_fen)
    new_fen = re.sub(r'1{4}', '4', new_fen)
    new_fen = re.sub(r'1{3}', '3', new_fen)
    new_fen = re.sub(r'1{2}', '2', new_fen)
    new_fen = re.sub(r'1', '', new_fen)

    return new_fen

# Parsování FEN stringu
fen_string = "1a6/2k5/8/8/3A4/8/8/6K1 b - - 0 1"
white_amazons, black_amazons, turn, white_pieces, black_pieces = parse_fen(fen_string)
amazons = white_amazons if turn == 'w' else black_amazons

# Výpis šachovnic pro všechny možné tahy Amazonky z výchozí pozice
print("Inicialní šachovnice:")
print_board_with_fen(fen_string)
print()
moves = generate_amazon_moves(amazons, white_pieces, black_pieces, turn)
for amazon_square, move in moves:
    new_fen = update_fen_with_amazon(fen_string, amazon_square, move, turn == 'w')
    san_notation = generate_san_notation(amazon_square, move, turn == 'w')
    print(f"SAN: {san_notation}")
    print(f"FEN: {new_fen}")
    print_board_with_fen(new_fen)
    print()