zeropod/plugin/chess/draw.py

from typing import Optional
from icecream import ic
from graphics import Graphics
from chess import Board, Outcome, Color, WHITE

class Draw:
    BOARD_SIZE = 64
    SQUARE_SIZE = 8

    PIECES = {
            "p": [
                " . . . . . . ",
                " . . ■ ■ . . ",
                " . ■ - - ■ . ",
                " . ■ - - ■ . ",
                " . . ■ ■ . . ",
                " . . . . . . "],
            "r": [
                " ■ ■ ■ ■ ■ ■ ",
                " ■ - - - - ■ ",
                " . ■ - - ■ . ",
                " . ■ - - ■ . ",
                " ■ - - - - ■ ",
                " ■ ■ ■ ■ ■ ■ "],
            "n": [
                " . . ■ ■ ■ . ",
                " . ■ - - - ■ ",
                " ■ - - - - ■ ",
                " . ■ - - - ■ ",
                " . . ■ - - ■ ",
                " . . ■ ■ ■ ■ "],
            "b": [
                " . . ■ ■ . . ",
                " . ■ - - ■ . ",
                " . ■ - - ■ . ",
                " . ■ - - ■ . ",
                " ■ - - - - ■ ",
                " ■ ■ ■ ■ ■ ■ "],
            "q": [
                " ■ . . . . ■ ",
                " ■ ■ . . ■ ■ ",
                " ■ - ■ ■ - ■ ",
                " ■ - - - - ■ ",
                " ■ - - - - ■ ",
                " ■ ■ ■ ■ ■ ■ "],

            "k": [
                " . . ■ ■ . . ",
                " ■ ■ - - ■ ■ ",
                " ■ - - - - ■ ",
                " ■ ■ - - ■ ■ ",
                " . ■ - - ■ . ",
                " . ■ ■ ■ ■ . "],
            }

    MENU = [
            "Play",
            "Quit",
            "New (W)",
            "New (B)"]

    def __init__(self, graphics: Graphics):
        self._graphics = graphics

    def _draw_piece(self, piece: list[str], white: bool, sqx, sqy, sqc):
        c = 0 if sqc == 1 else 1
        filled = (white and sqc == 0) or (not white and sqc == 1)
        for row in range(len(piece)):
            pixels = piece[row].replace(" ", "")
            for col in range(len(pixels)):
                x = (sqx * self.SQUARE_SIZE) + col + 1
                y = (sqy * self.SQUARE_SIZE) + row + 1
                pixel = pixels[col]
                if pixel == '■':
                    self._graphics.pixel(x, y, c)
                elif pixel =='-':
                    self._graphics.pixel(x, y, c if filled else sqc)

    def _clear_info(self):
        self._graphics.fill_rect(
                self.BOARD_SIZE, 0, self.BOARD_SIZE, self.BOARD_SIZE, 0)

    def _format_move(self, move: str):
        if len(move) == 4:
            return move.upper()[:2] + "-" + move.upper()[2:]
        return move

    def _invert_pixel(self, x: int, y: int):
        c = self._graphics.pixel(x, y)
        c = 0 if c != 0 else 1
        self._graphics.pixel(x, y, c)

    def _cursor_clear(self, except_for: list[str], player_color: Color):
        c = 1
        for row in range(8):
            for col in range(8): # 0,0 w = a8 b = h1 ... 7,7 w = h1 b = a8
                sqx = col if player_color == WHITE else 7 - col
                sqy = 7 - row if player_color == WHITE else row
                square = chr(sqx + 97) + str(sqy + 1)
                if not (square in except_for):
                    # see if square is right color
                    sqc = self._graphics.pixel(col * 8, row * 8)
                    if sqc != c:
                        self._cursor(square, player_color);
                c = 0 if c == 1 else 1
            c = 0 if c == 1 else 1

    def _cursor(self, square: str, player_color: Color):
        sqx = ord(square[0]) - 97
        sqy = int(square[1]) - 1
        if player_color == WHITE:
            sqy = 7 - sqy
        else:
            sqx = 7 - sqx
        for i in range(8):
            # top and bottom row
            self._invert_pixel(sqx * 8 + i, sqy * 8)
            self._invert_pixel(sqx * 8 + i, sqy * 8 + 7)
            if i > 0 and i < 7:
                # sides (skipping top and bottom row)
                self._invert_pixel(sqx * 8, sqy * 8 + i)
                self._invert_pixel(sqx * 8 + 7, sqy * 8 + i)

    def draw_board(self, board: Board, player_color: Color):
        self._graphics.fill_rect(0, 0, self.BOARD_SIZE, self.BOARD_SIZE, 0)
        nb = "".join(str(board).split()).replace(" ", "")
        nb = nb if player_color == WHITE else nb[::-1]
        c = True
        for row in range(8):
            for col in range(8):
                if c:
                    x = col * self.SQUARE_SIZE
                    y = row * self.SQUARE_SIZE
                    self._graphics.fill_rect(x, y, self.SQUARE_SIZE, self.SQUARE_SIZE, 1)
                p = nb[(row * 8) + col]
                if p.lower() in self.PIECES:
                    self._draw_piece(
                            self.PIECES[p.lower()], p.isupper(), col, row, c)
                c = not c
            c = not c
        self._graphics.show()

    def draw_menu(self, menu_index: int):
        self._clear_info()
        self._graphics.text("Menu:", 12, 0, 1)
        for i in range(len(self.MENU)):
            marker = "> " if i == menu_index else "  "
            self._graphics.text(marker + self.MENU[i], 12, i + 2, 1)
        self._graphics.show()

    def draw_thinking(self, last_move: str):
        self._clear_info()
        self._graphics.text(self._format_move(last_move), 12, 0, 1)
        self._graphics.text("Shh, I'm", 12, 2, 1)
        self._graphics.text("thinking!", 12, 3, 1)
        self._graphics.show()

    def draw_checkmate(self, move: str, player_won: bool):
        self._clear_info()
        self._graphics.text(move, 12, 0, 1)
        self._graphics.text("Checkmate", 12, 1, 1)
        self._graphics.text(
                "you win!" if player_won else "you lose!", 12, 2, 1)
        self._graphics.show()

    def draw_select(self, last_move: str, player_color: Color, src: Optional[str], dst: Optional[str] = None):
        self._clear_info()
        self._graphics.text(self._format_move(last_move), 12, 0, 1)
        self._graphics.text("You move:", 12, 2, 1)
        if src != None and dst == None:
            self._graphics.text("<" + src.upper() + "> - __", 12, 3, 1)
            self._cursor_clear([src], player_color)
            self._cursor(src, player_color)
        elif src != None and dst != None:
            self._graphics.text(src.upper() + " - <" + dst.upper() + ">", 12, 3, 1)
            self._cursor_clear([src, dst], player_color)
            self._cursor(src, player_color)
            self._cursor(dst, player_color)
        else:
            self._cursor_clear(list(), player_color)

        self._graphics.show()

    def draw_game_over(self, outcome: Optional[Outcome]):
        self._clear_info()
        self._graphics.text("Game over", 12, 0, 1)
        if outcome != None:
            if outcome.winner == None:
                self._graphics.text("Draw!", 12, 2, 1)
            else:
                self._graphics.text("Winner:", 12, 2, 1)
                self._graphics.text("White!" if outcome.winner == WHITE else "Black!",
                        12, 3, 1)
        self._graphics.show()