From 4aa67bd83057ccf5253184b8a0f1492cabdc0875 Mon Sep 17 00:00:00 2001
From: Tibor Bizjak <tibor.bizjak@gmail.com>
Date: Tue, 10 Sep 2019 14:01:17 +0200
Subject: [PATCH] Added castling logic, fixed en passant and promotion

---
 chess.py | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 6 deletions(-)

diff --git a/chess.py b/chess.py
index 1c3b235..3c28193 100644
--- a/chess.py
+++ b/chess.py
@@ -119,6 +119,11 @@ class Game:
         self.moves = []
         self.stack = []
         self.turn  = "white"
+        self.can_castle = {"white" : {"queen" : True,
+                                      "king"  : True},
+                           "black" : {"queen" : True,
+                                      "king"  : True}
+                          }
 
     def __repr__(self):
         # Unicode board representation
@@ -148,6 +153,11 @@ class Game:
             self.board[sq].color = "white"
         for sq in cross(files, "78"):
             self.board[sq].color = "black"
+        # Castling markers
+        self.board["a1"].side = "queen"
+        self.board["a8"].side = "queen"
+        self.board["h1"].side = "king"
+        self.board["h8"].side = "king"
 
     def is_empty(self, sq):
         return self.board[sq] == None
@@ -155,23 +165,66 @@ class Game:
     def occupying(self, piece):
         return [sq for sq in squares if self.board[sq] == piece]
 
+    def castle(self, side):
+        color = self.turn
+        rank = home_ranks[self.turn]
+
+        if side not in ("queen", "king"):
+            return False
+        if not self.can_castle[color][side]:
+            return False
+
+        if side == "queen":
+            check_if_empty    = cross("bcd", rank)
+            check_if_attacked = cross("cde", rank)
+            rook_move = "a"+rank, "d"+rank
+            king_move = "e"+rank, "c"+rank
+        elif side == "king":
+            check_if_empty    = cross("fg", rank)
+            check_if_attacked = check_if_empty
+            rook_move = "h"+rank, "f"+rank
+            king_move = "e"+rank, "g"+rank
+        are_empty    = all(map(self.is_empty, check_if_empty))
+        not_attacked = not any(self.is_attacked(color, sq)
+                               for sq in check_if_attacked)
+
+        if not (are_empty and not_attacked):
+            return False
+
+        self.board[king_move[1]] = self.board[king_move[0]]
+        self.board[rook_move[1]] = self.board[rook_move[0]]
+        self.board[king_move[0]] = Piece()
+        self.board[rook_move[0]] = Piece()
+        self.can_castle[color]["queen"] = False
+        self.can_castle[color]["king"]  = False
+
     def move(self, source, target, promotion=None):
         if not self.is_legal(source, target, promotion):
             return False
         board = self.board
         moved = board[source]
         eaten = board[target]
+
         if promotion == None:
             promotion = moved
         else:
             promotion = Piece(self.turn, promotion)
 
-        if moved.piece == "pawn" and target in moves(source, bishop_dirs):
+        if moved.piece == "pawn" and target in moves(source, bishop_dirs) and eaten==None:
             # En passant
             a, b = move(target, up), move(target, down)
             board[a] = Piece()
             board[b] = Piece()
 
+        # Castling
+        if moved.piece == "rook":
+            self.can_castle[self.turn][moved.side] = False
+        if eaten.piece == "rook":
+            self.can_castle[eaten.color][eaten.side] = False
+        if moved.piece == "king":
+            self.can_castle[self.turn]["queen"] = False
+            self.can_castle[self.turn]["king"]  = False
+
         self.moves.append((source, target))
         self.stack.append(eaten)
 
@@ -189,7 +242,7 @@ class Game:
             return False
         if promotion == None and pawn_on_last_rank:
             return False
-        if promotion == "pawn":
+        if promotion in ("king", "pawn"):
             return False
         return color == self.turn and target in self.possible_moves(source)
 
@@ -289,14 +342,16 @@ def test():
     game = Game()
     assert len(squares) == 8**2
     assert sum(map(len, init_positions.values())) == 8*4
-    moves = [("a2", "a4"), ("b8", "c6"), ("a4", "a5"), ("b7", "b5"), ("a5", "b6"),
-             ("c6", "e5"), ("b6", "b7"), ("h7", "h5"), ("b7", "b8", "pawn")]
+    moves = [('d2', 'd4'), ('d7', 'd5'), ('d1', 'd3'), ('h7', 'h5'),
+             ('c1', 'h6'), ('g7', 'h6'), ('b1', 'c3'), ('d8', 'd6'),
+             ('h2', 'h4'), ('d6', 'f4'), ('e2', 'e3'), ('f4', 'g4'), ('e1', 'd1')]
     for m in moves:
         game.move(*m)
+        print (game.can_castle)
         print (game)
 
-    game.board["c3"] = Piece("black", "king")
-
+    game.castle("queen")
+    print(game)
     print (game.attacks("c3"))
     print(game.possible_moves("c3"))
     print(game.is_attacked("black", "d8"))