TDT4109/Exercise 10/chess/board.py

from typing import Callable, Iterable
from os import system
from shutil import get_terminal_size as getTerminalSize

from piece import Piece


class Board:

  def __init__(self):
    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']],
    ]

  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, column 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):
      """highlight contents of a piece with xterm-256colors modifiers"""
      stringArray[y * 2 +
                  1][x * 4 +
                     1] = config['highlightEscapeCodes'][0] + stringArray[y * 2 + 1][x * 4 + 1]
      stringArray[y * 2 + 1][x * 4 + 3] += config['highlightEscapeCodes'][1]

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

      characterMap = {
        '─': '═',
        '│': '║',
        '┼': '╬',
        '╰': '╚',
        '╯': '╝',
        '╭': '╔',
        '╮': '╗',
        '├': '╠',
        '┴': '╩',
        '┤': '╣',
        '┬': '╦',
      }

      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] = characterMap[
          stringArray[y][x]] if stringArray[y][x] in characterMap else stringArray[y][x]

    for x, y in config['highlightedBoxes']:
      highlightBox(x, y)
    for x, y in config['highlightedContent']:
      highlightContent(x, y)

    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')
      menuString = '\n' + player.name + '\n\n'
      menuString += self.draw({'highlightedBoxes': [(x, y)]}) + '\n'
      inputString = f"  W E\nA S D  <- Enter : "

      def centerText(text):
        terminalWidth = getTerminalSize((60, 0))[0]   # Column size 60 as fallback
        return "\n".join(line.center(terminalWidth) for line in text.split('\n'))

      def centerBlockText(text):
        terminalWidth = getTerminalSize((60, 0))[0]   # Column size 60 as fallback
        textArray = text.split('\n')
        offset = int((terminalWidth - len(textArray[0])) / 2)
        return "\n".join(offset * ' ' + line for line in textArray)

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

      print(menuString)

      key = input(inputString)[0]
      if key in ['s', 'j'] and y != 7: y += 1
      elif key in ['w', 'k'] and y != 0: y -= 1
      elif key in ['d', 'l'] and x != 7: x += 1
      elif key in ['a', 'h'] and x != 0: x -= 1
      elif key == 'e': return (x, y)

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

  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):
        if condition(piece):
          result.append((x, y))
    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)
    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)
    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):
          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 checkStaleMate(self, color) -> bool:
    """Check whether a team is caught in stalemate. The color is the color of the team to check"""
    enemyPieces = self.getPositionsWhere(lambda piece: piece.color == color)
    getLegalMoves = lambda piece: Piece.possibleMoves(*piece, self, legalMoves = self.getPositionsToProtectKing(color))
    piecesHasNoLegalMoves = any( getLegalMoves(piece) == None for piece in enemyPieces)
    return (not self.checkCheck(color)) and piecesHasNoLegalMoves

  def checkCheckMate(self, color) -> bool:
    """Check whether a team is caught in checkmate. The color is the color of the team to check"""
    enemyPieces = self.getPositionsWhere(lambda piece: piece.color == color)
    getLegalMoves = lambda piece: Piece.possibleMoves(*piece, self, legalMoves = self.getPositionsToProtectKing(color))
    piecesHasNoLegalMoves = any( getLegalMoves(piece) == None for piece in enemyPieces)
    return self.checkCheck(color) and piecesHasNoLegalMoves

  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