TDT4109/Exercise 10/chess/board.py

from typing import Callable, Iterable, Union
from os import system

from piece import Piece
from util import centerText, centerBlockText, boxCharsToBoldBoxCharsMap, determineMove


class Board:

  def __init__(self, boardState=None):
    self.boardArray = [
      [Piece(type, 'black') for type in ['r', 'n', 'b', 'q', 'k', 'b', 'n', 'r']],
      [Piece('p', 'black') for _ in range(8)],
      *[[None for _ in range(8)] for _ in range(4)],
      [Piece('p', 'white') for _ in range(8)],
      [Piece(type, 'white') for type in ['r', 'n', 'b', 'q', 'k', 'b', 'n', 'r']],
    ] if boardState == None else boardState

  def draw(
    self,
    config={
      'highlightedContent': [],
      'highlightEscapeCodes': ('\033[32;5;7m', '\033[0m'),
      'highlightedBoxes': [],
    }
  ) -> str:
    """Returns a string representing the board

    config options:
      highlightedContent: [(x,y)] - Pieces to color
      highlightEscapeCodes: (str, str) - Terminal escape codes to color highlightedContent with
      highlightedBoxes: [(x,y)] - Boxes to make bold
    """

    def fillConfigDefaultValue(key, defaultValue):
      if key not in config:
        config[key] = defaultValue

    fillConfigDefaultValue('highlightedContent', [])
    fillConfigDefaultValue('highlightedBoxes', [])
    fillConfigDefaultValue('highlightEscapeCodes', ('\033[32;5;7m', '\033[0m'))

    # Draw general outline
    stringArray = [list('┼' + '───┼' * 8)] + [[None] for _ in range(8 * 2)]
    for y, row in enumerate(self.boardArray):
      for x, _ in enumerate(row):
        stringArray[2 * y + 1][4 * x] = '│'
        stringArray[2 * y + 2][4 * x] = '┼'

        stringArray[2 * y + 1] += list(
          ' {} │'.format(str(self.boardArray[y][x]) if self.boardArray[y][x] != None else ' '))
        stringArray[2 * y + 2] += list('───┼')

    # Overwrite corners
    stringArray[0][0] = '╭'
    stringArray[0][-1] = '╮'
    stringArray[-1][0] = '╰'
    stringArray[-1][-1] = '╯'

    # Overwrite T-junctions
    for i in range(int(len(stringArray[0]) / 4) - 1):
      stringArray[0][i * 4 + 4] = '┬'
      stringArray[-1][i * 4 + 4] = '┴'
    for i in range(int(len(stringArray) / 2) - 1):
      stringArray[i * 2 + 2][0] = '├'
      stringArray[i * 2 + 2][-1] = '┤'

    def highlightContent(x, y, escapeCodes=config['highlightEscapeCodes']):
      """highlight contents of a piece with xterm-256colors modifiers"""
      stringArray[y * 2 + 1][x * 4 + 1] = \
        escapeCodes[0] + stringArray[y * 2 + 1][x * 4 + 1]
      stringArray[y * 2 + 1][x * 4 + 3] += escapeCodes[1]

    def highlightBox(x, y):
      """Make box around a piece bold"""

      pointsToChange = \
      [(x * 4 + 0, y * 2 + i) for i in range(3)] + \
      [(x * 4 + 4, y * 2 + i) for i in range(3)] + \
      [(x * 4 + i, y * 2 + 0) for i in range(1,4)] + \
      [(x * 4 + i, y * 2 + 2) for i in range(1,4)]

      for x, y in pointsToChange:
        stringArray[y][x] = boxCharsToBoldBoxCharsMap[
          stringArray[y]
          [x]] if stringArray[y][x] in boxCharsToBoldBoxCharsMap else stringArray[y][x]

      # Color white pieces
      for piece in self.getPositionsWhere(lambda piece: piece.color == 'white'):
        highlightContent(*piece, ('\033[7m', '\033[0m'))

    for box in config['highlightedBoxes']:
      highlightBox(*box)
    for piece in config['highlightedContent']:
      highlightContent(*piece)

    return '\n'.join([''.join(line) for line in stringArray])

  def selectPiece(self, player, x=0, y=0, centering=True) -> tuple:
    """Lets the user select a piece from a graphic board"""

    while True:
      system('clear')
      playerString = '\n' + player.name + '\n\n'
      menuString = self.draw({'highlightedBoxes': [(x, y)]}) + '\n'
      inputString = f"  W E\nA S D  <- Enter : "

      if centering:
        playerString = centerText(playerString)
        menuString = centerBlockText(menuString)
        inputString = centerBlockText(inputString)

      print(playerString)
      print(menuString)

      try:
        key = input(inputString)[0]
      except IndexError:
        key = ''
      try:
        if move := determineMove(key, x, y, (0, 7)):
          x += move[0]
          y += move[1]
        elif key == 'e' and self.getPieceAt(x, y).color==player.color:
          return (x, y)
      except AttributeError:
        pass

  def selectMove(self, player, x, y, legalMoves, centering=True) -> Union[tuple, bool]:
    """Lets the user select a move to make from a graphic board"""

    while True:
      system('clear')
      playerString = '\n' + player.name + '\n\n'
      menuString = self.draw({
        'highlightedBoxes': [(x, y)],
        'highlightedContent': legalMoves
      }) + '\n'
      inputString = f"Q W E\nA S D  <- Enter : "

      if centering:
        playerString = centerText(playerString)
        menuString = centerBlockText(menuString)
        inputString = centerBlockText(inputString)

      print(playerString)
      print(menuString)

      try:
        key = input(inputString)[0]
      except IndexError:
        key = ''
      if move := determineMove(key, x, y, (0, 7)):
        x += move[0]
        y += move[1]
      elif key == 'q':
        return False
      elif key == 'e' and (x, y) in legalMoves:
        return (x, y)

  def getPieceAt(self, x, y) -> Piece:
    try:
      return self.boardArray[y][x]
    except IndexError:
      return None

  def getPositionsWhere(self, condition: Callable[[Piece], bool]) -> Iterable[tuple]:
    """ Returns a list of xy pairs of the pieces where a condition is met """

    result = []
    for y, row in enumerate(self.boardArray):
      for x, piece in enumerate(row):
        try:
          if condition(piece):
            result.append((x, y))
        except AttributeError:
          pass
    return result

  def checkCheck(self, color) -> bool:
    """Check whether a team is caught in check. The color is the color of the team to check"""
    king = self.getPositionsWhere(lambda piece: piece.type == 'k' and piece.color == color)[0]
    piecesToCheck = self.getPositionsWhere(lambda piece: piece.color != color)
    return any([king in Piece.possibleMoves(*piece, self) for piece in piecesToCheck])

  def getPositionsToProtectKing(self, color) -> Iterable[tuple]:
    """Get a list of the positions to protect in order to protect the king when in check. The color is the color of the team who's in check"""
    king = self.getPositionsWhere(lambda piece: piece.type == 'k' and piece.color == color)[0]
    piecesToCheck = self.getPositionsWhere(lambda piece: piece.color != color)
    for piece in piecesToCheck:
      if king not in Piece.possibleMoves(*piece, self):
        piecesToCheck.remove(piece)
    result = []
    for piece in piecesToCheck:
      result.append(piece)
      if self.getPieceAt(*piece).type not in ['p', 'n', 'k']:

        def getDirection(fromPosition, toPosition) -> tuple:
          x = -1 if toPosition[0] > fromPosition[0] else \
            0 if toPosition[0] == fromPosition[0] else 1
          y = -1 if toPosition[1] > fromPosition[1] else \
            0 if toPosition[1] == fromPosition[1] else 1
          return (x, y)

        direction = getDirection(piece, king)

        def getPositionsUntilKing(x, y, direction) -> Iterable[tuple]:
          result = []
          while self.getPieceAt(x, y) == None:
            result.append((x, y))
            x += direction[0]
            y += direction[1]
          return result

        result += getPositionsUntilKing(*piece, direction)

    return result

  def playerHasLegalMoves(self, color) -> bool:
    enemyPieces = self.getPositionsWhere(lambda piece: piece.color == color)
    getLegalMoves = lambda piece: Piece.possibleMoves(
      *piece, self, legalMoves=self.getPositionsToProtectKing(color))
    return not any(getLegalMoves(piece) == None for piece in enemyPieces)

  def checkStaleMate(self, color) -> bool:
    """Check whether a team is caught in stalemate. The color is the color of the team to check"""
    return (not self.checkCheck(color)) and not self.playerHasLegalMoves(color)

  def checkCheckMate(self, color) -> bool:
    """Check whether a team is caught in checkmate. The color is the color of the team to check"""
    return self.checkCheck(color) and not self.playerHasLegalMoves(color)

  def movePiece(self, position, toPosition, piecesToRemove=None):
    x, y = position
    toX, toY = toPosition
    self.boardArray[toY][toX] = self.boardArray[y][x]
    self.boardArray[y][x] = None

    if piecesToRemove != None:
      for x, y in piecesToRemove:
        self.boardArray[y][x] = None