Chess Game

from collections import namedtuple

def chess_to_list(pos: namedtuple) -> tuple:  # position(4, 2) -> (6,3)
  return (8 - pos.y, pos.x - 1)

def list_to_chess(pos: tuple) -> namedtuple:  # (6,3) -> position(4, 2)
  return position(pos[1] + 1, 8 - pos[0])

class Piece:

  def __init__(self, colour, pos):
    self.colour = colour
    self.pos = pos

  def __str__(self):
    return f"{self.type}"

  def __repr__(self):
    return f"{self.repr_type}"


class Pawn(Piece):
  __slots__ = 'first_move', 'pinned', 'type', 'repr_type', 'direction'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.first_move = 6 if self.colour == 'W' else 1
    self.pinned = False
    self.type = '♙' if self.colour == 'W' else '♟︎'
    self.repr_type = 'wP' if self.colour == 'W' else 'bP'
    self.direction = 1 if self.colour == 'W' else -1

  def legal_moves(self, board):
    legal = []
    current = chess_to_list(self.pos)
    current_x = current[1]
    current_y = current[0]
    if current_y <= 7:
      # d+2, check d3 and d4 == (6,3) check (5,3) and (4,3)
      if board[current_y - self.direction][current_x] is None:
        legal.append(list_to_chess((current_y - self.direction, current_x)))

        # Test Forward moves, but only if there is nothing one ahead of the pawn
        if current_y == self.first_move:
          if board[current_y - (2 * self.direction)][current_x] is None:
            legal.append(
              list_to_chess((current_y - (2 * self.direction), current_x)))

      # Test diagonal captures - Short Circuit Evaluation

      # Check to the Right
      if current_x <= 6:
        if (board[current_y - self.direction][current_x + 1]) is not None and (
            board[current_y - self.direction][current_x +
                                              1]).colour != self.colour:
          legal.append(
            list_to_chess((current_y - self.direction, current_x + 1)))

      # Check to the Left
      if current_x <= 2:
        if (board[current_y - self.direction][current_x - 1]) is not None and (
            board[current_y - self.direction][current_x -
                                              1]).colour != self.colour:
          legal.append(
            list_to_chess((current_y - self.direction, current_x + 1)))

    # test if move puts in check/keeps king in check, but causes error <---

    new_legal_moves = []
    for move in legal:
      g.board.move_piece(self.pos, move)
      if not g.in_check(self.colour):
        new_legal_moves.append(move)

      g.board.move_piece(move, self.pos)

    return new_legal_moves

    return legal

class Bishop(Piece):

  __slots__ = 'pinned', 'type', 'repr_type'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.pinned = False
    self.type = '♗' if self.colour == 'W' else '♝'
    self.repr_type = 'wB' if self.colour == 'W' else 'bB'

  def legal_moves(self, board):
    legal = []
    current_x = self.pos.x
    current_y = self.pos.y

    # Check up-right (Positive Gradient)
    while current_y != 8 and current_x != 8:
      current_y += 1
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check down-left (Positive Gradient)
    while current_y != 1 and current_x != 1:
      current_y -= 1
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check up-left (Negative Gradient)
    while current_y != 8 and current_x != 1:
      current_y += 1
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check down-right (Negative Gradient)
    while current_y != 1 and current_x != 8:
      current_y -= 1
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    return legal


class Knight(Piece):

  __slots__ = 'pinned', 'type', 'repr_type'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.pinned = False
    self.type = '♘' if self.colour == 'W' else '♞'
    self.repr_type = 'wN' if self.colour == 'W' else 'bN'

  def legal_moves(self, board):
    legal = []
    knight_moves = [(2, 1), (2, -1), (-2, 1), (-2, -1), (1, 2), (1, -2),
                    (-1, 2), (-1, -2)]

    for move in knight_moves:
      new_x = self.pos.x + move[0]
      new_y = self.pos.y + move[1]

      if 1 <= new_x <= 7 and 1 <= new_y <= 7:
        current = chess_to_list(position(new_x, new_y))

        target = board[current[0]][current[1]]

        if target is None or target.colour != self.colour:
          legal.append(position(new_x, new_y))

    return legal


class Rook(Piece):
  __slots__ = 'has_moved', 'pinned', 'type', 'repr_type'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.has_moved = False
    self.pinned = False
    self.type = '♖' if self.colour == 'W' else '♜'
    self.repr_type = 'wR' if self.colour == 'W' else 'bR'

  def legal_moves(self, board):
    legal = []
    current = chess_to_list(self.pos)
    current_x = current[1]
    current_y = current[0]

    # Check Up
    while current_y != 8:
      current_y += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Right
    while current_x != 8:
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Down
    while current_y != 1:
      current_y -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Left
    while current_x != 1:
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    return legal


class Queen(Piece):

  __slots__ = 'pinned', 'type', 'repr_type'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.pinned = False
    self.type = '♕' if self.colour == 'W' else '♛'
    self.repr_type = 'wQ' if self.colour == 'W' else 'bQ'

  def legal_moves(self, board):
    legal = []
    current_x = self.pos.x
    current_y = self.pos.y

    # Check Rook Moves

    # Check Up
    while current_y != 8:
      current_y += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Right
    while current_x != 8:
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Down
    while current_y != 1:
      current_y -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check Left
    while current_x != 1:
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    # Check Bishop Moves

    # Check up-right (Positive Gradient)
    while current_y != 8 and current_x != 8:
      current_y += 1
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check down-left (Positive Gradient)
    while current_y != 1 and current_x != 1:
      current_y -= 1
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check up-left (Negative Gradient)
    while current_y != 8 and current_x != 1:
      current_y += 1
      current_x -= 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    current_x = self.pos.x
    current_y = self.pos.y

    # Check down-right (Negative Gradient)
    while current_y != 1 and current_x != 8:
      current_y -= 1
      current_x += 1
      test_pos = chess_to_list(position(current_x, current_y))

      if board[test_pos[0]][test_pos[1]] is None:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
      elif board[test_pos[0]][test_pos[1]].colour != self.colour:
        legal.append(list_to_chess((test_pos[0], test_pos[1])))
        break
      else:
        break

    return legal


class King(Piece):

  __slots__ = 'can_castle', 'pinned', 'type', 'repr_type'

  def __init__(self, colour, pos):
    super().__init__(colour, pos)
    self.can_castle = True
    self.pinned = False
    self.type = '♔' if self.colour == 'W' else '♚'
    self.repr_type = 'wK' if self.colour == 'W' else 'bK'

  def legal_moves(self, board):
    legal = []
    king_moves = [(0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1), (-1, 0),
                  (-1, 1)]

    for move in king_moves:
      new_x = self.pos.x + move[0]
      new_y = self.pos.y + move[1]

      if 1 <= new_x <= 7 and 1 <= new_y <= 7:
        current = chess_to_list(position(new_x, new_y))

        target = board[current[0]][current[1]]

        if target is None or target.colour != self.colour:
          legal.append(position(new_x, new_y))

    return legal


class Board:

  def __init__(self):
    self.board = [[None for i in range(8)] for j in range(8)]

  def __str__(self):
    board_str = ''
    c = 8
    for row in self.board:
      board_str += f'{c} '
      c -= 1
      for piece in row:
        if piece is None:
          board_str += '- '
        else:
          board_str += f"{piece} "
      board_str += '\n'

    board_str += '  a b c d e f g h'
    return board_str

  def add_piece(self, piece, chess_coord):
    list_coord = chess_to_list(chess_coord)
    self.board[list_coord[0]][list_coord[1]] = piece

  def move_piece(self,
                 start_coord: namedtuple,
                 end_coord: namedtuple,
                 captured=None) -> Piece | None:
    start_list_coord = chess_to_list(start_coord)

    end_list_coord = chess_to_list(end_coord)

    piece = self.board[start_list_coord[0]][start_list_coord[1]]
    print(f"{piece = !r}")
    captured = self.board[end_list_coord[0]][end_list_coord[1]]

    self.board[start_list_coord[0]][start_list_coord[1]] = None
    self.board[end_list_coord[0]][end_list_coord[1]] = piece
    piece.pos = end_coord

    return captured



types = [Pawn, Bishop, Knight, Rook, Queen, King]
colours = ['W', 'B']
position = namedtuple("Position", ['x', 'y'])

pieces = [Pawn('W', position(x, 2)) for x in range(1, 9)
          ] + [Pawn('B', position(x, 7)) for x in range(1, 9)] + [
            Rook('W', position(1, 1)),
            Rook('W', position(8, 1)),
            Rook('B', position(1, 8)),
            Rook('B', position(8, 8))
          ] + [
            Knight('W', position(2, 1)),
            Knight('W', position(7, 1)),
            Knight('B', position(2, 8)),
            Knight('B', position(7, 8))
          ] + [
            Bishop('W', position(3, 1)),
            Bishop('W', position(6, 1)),
            Bishop('B', position(3, 8)),
            Bishop('B', position(6, 8))
          ] + [
            Queen('W', position(4, 1)),
            King('W', position(5, 1)),
            Queen('B', position(4, 8)),
            King('B', position(5, 8))
          ]


class Game:

  def __init__(self):
    self.board = Board()

  def setup(self):
    for piece in pieces:
      chess_coord = piece.pos
      self.board.add_piece(piece, chess_coord)

  def restart(self):
    # Move pieces here at their default positions, so the game can be restarted
    self.board.board = [[None for i in range(8)] for j in range(8)]
    self.setup()

  def in_check(self, colour, checked=set()): # Have it so that you can't make a move if you are in check!!!!
    king_pos = None
    for piece in pieces:
      if isinstance(piece, King) and piece.colour == colour:
        king_pos = piece.pos
        break

    return bool([piece for piece in pieces if king_pos in piece.legal_moves(self.board.board)])

# setup
g = Game()
g.setup()

# Testing g.in_check()
g.in_check('W')