TDT4109/Exercise 10/chess/piece.py

from typing import Iterable, Callable
from itertools import product
from copy import deepcopy


class Piece:

  def __init__(self, type, color):
    self.type = type
    self.color = color

  def __str__(self):
    return self.type.upper() if self.color == 'white' else self.type

  @staticmethod
  def possibleMoves(x, y, board, legalMoves=None) -> Callable[[int, int], Iterable[tuple]]:
    piece = board.getPieceAt(x, y)
    moves = []

    # TODO: If player is in check, only validate the moves that are leading to places that can be blocked in order to save the king. Make a function that overwrites `moves`. Notice: The king is the only piece that can 'Escape' to a position that may or may not be in the list of positions to be blocked. Make an exception

    pieceIsEnemyColor = lambda pieceToCheck: pieceToCheck != None and pieceToCheck.color != piece.color
    pieceIsEmpty = lambda pieceToCheck: pieceToCheck == None
    pieceIsEmptyOrEnemyColor = lambda pieceToCheck: pieceToCheck == None or pieceToCheck.color != piece.color

    def addMoveIfTrue(xOffset, yOffset, condition: Callable[[Piece], bool]):
      """Tests a condition against a position away from self. Adds if condition returns true"""
      if condition(board.getPieceAt(x + xOffset, y + yOffset)):
        moves.append((x + xOffset, y + yOffset))
        return True
      return False

    def assertNotCheck(newX, newY) -> bool:
      testBoard = deepcopy(board)
      testBoard.movePiece((x, y), (newX, newY))   # J
      return not testBoard.checkCheck()

    def addWhileInsideBoard(direction: tuple):
      localX = x
      localY = y
      while localX not in [0, 7] and localY not in [0, 7]:
        localX += direction[0]
        localY += direction[1]
        if board.getPieceAt(localX, localY) == None:
          moves.append((localX, localY))
        else:
          if board.getPieceAt(localX, localY).color != piece.color:
            moves.append((localX, localY))
          return

    if piece.type == 'p':
      addMoveIfTrue(1, 1, pieceIsEnemyColor)
      addMoveIfTrue(-1, 1, pieceIsEnemyColor)
      if addMoveIfTrue(0, 1, pieceIsEmpty):
        addMoveIfTrue(0, 2, lambda pieceToCheck: pieceToCheck == None and piece.moves == 0)

    elif piece.type == 'n':
      positions = [(-2, -1), (-2, 1), (-1, -2), (-1, 2), (1, -2), (1, 2), (2, -1), (2, 1)]
      for position in positions:
        addMoveIfTrue(*position, pieceIsEmptyOrEnemyColor)

    elif piece.type == 'k':
      positions = list(product([-1, 0, 1], repeat=2))
      positions.remove((0, 0))
      for position in positions:
        addMoveIfTrue(*position, pieceIsEmptyOrEnemyColor)

      moves = [position for position in moves if assertNotCheck(*position)]

    elif piece.type == 'r':
      for direction in product([0, 1], repeat=2):
        addWhileInsideBoard(direction)

    elif piece.type == 'b':
      for direction in product([-1, 1], repeat=2):
        addWhileInsideBoard(direction)

    elif piece.type == 'q':
      directions = list(product([-1, 0, 1], repeat=2))
      directions.remove((0, 0))
      for direction in directions:
        addWhileInsideBoard(direction)

    #TODO: remove moves that will put the king in check

    if legalMoves != None and piece.type != 'k':
      moves = [move for move in moves if move in legalMoves]

    return moves