From 36f3873816c7ede65008946910763bb9fd987908 Mon Sep 17 00:00:00 2001
From: ElnuDev <elnu@elnu.com>
Date: Wed, 15 Mar 2023 17:00:04 -0700
Subject: [PATCH] Fix checkmate logic to account for other piece taking
 attacker

---
 src/Board.java   | 17 ++++++++++++++---
 src/ChessAI.java |  6 +++---
 src/Piece.java   | 10 +++++-----
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/Board.java b/src/Board.java
index eaaa012..74ff175 100644
--- a/src/Board.java
+++ b/src/Board.java
@@ -197,7 +197,18 @@ public class Board {
                 break;
             }
         }
-        isGameOver = inCheck && oppositeKing.getLegalMoves(oppositeKingPosition, this).size() == 0;
+        if (inCheck) {
+            isGameOver = true;
+            for (Move move : getAllLegalMoves(!movedPiece.black)) {
+                move(move);
+                if (!oppositeKing.isInCheck(move, this)) {
+                    isGameOver = false;
+                    undoMove();
+                    break;
+                }
+                undoMove();
+            }
+        }
         victor = movedPiece.black;
     }
 
@@ -287,10 +298,10 @@ public class Board {
         forEachPiece((x, y, piece) -> tileAction.forEachTile(new BoardCoordinate(x, y), get(x, y)));
     }
 
-    public ArrayList<Move> getAllLegalMoves() {
+    public ArrayList<Move> getAllLegalMoves(boolean black) {
         ArrayList<Move> allLegalMoves = new ArrayList<>();
         forEachPiece((from, piece) -> {
-            if (!piece.black) return;
+            if (piece.black != black) return;
             ArrayList<BoardCoordinate> legalTiles = piece.getLegalMoves(from, this);
             for (BoardCoordinate to : legalTiles) {
                 allLegalMoves.add(new Move(from, to));
diff --git a/src/ChessAI.java b/src/ChessAI.java
index 24b642e..8fe8ce5 100644
--- a/src/ChessAI.java
+++ b/src/ChessAI.java
@@ -5,7 +5,7 @@ public class ChessAI {
         int bestScore = Integer.MIN_VALUE;
         Move bestMove = null;
 
-        for (Move move : board.getAllLegalMoves()) {
+        for (Move move : board.getAllLegalMoves(true)) {
             board.move(move);
             int score = minimax(board, MAX_DEPTH, Integer.MIN_VALUE, Integer.MAX_VALUE, false);
             board.undoMove();
@@ -27,7 +27,7 @@ public class ChessAI {
         int score;
         if (maximizingPlayer) {
             score = Integer.MIN_VALUE;
-            for (Move move : board.getAllLegalMoves()) {
+            for (Move move : board.getAllLegalMoves(true)) {
                 board.move(move);
                 score = Math.max(score, minimax(board, depth - 1, alpha, beta, false));
                 board.undoMove();
@@ -38,7 +38,7 @@ public class ChessAI {
             }
         } else {
             score = Integer.MAX_VALUE;
-            for (Move move : board.getAllLegalMoves()) {
+            for (Move move : board.getAllLegalMoves(true)) {
                 board.move(move);
                 score = Math.min(score, minimax(board, depth - 1, alpha, beta, true));
                 board.undoMove();
diff --git a/src/Piece.java b/src/Piece.java
index a98ad2e..c0579f0 100644
--- a/src/Piece.java
+++ b/src/Piece.java
@@ -38,7 +38,7 @@ public abstract class Piece {
                     possibleMove.y >= Board.BOARD_SIZE ||
 
                     // is in check
-                    (doCheckChecks && isInCheck(position, possibleMove, board))
+                    (doCheckChecks && isInCheck(new Move(position, possibleMove), board))
             ) {
                 legalMoves.remove(i);
                 i--;
@@ -48,16 +48,16 @@ public abstract class Piece {
         return legalMoves;
     }
 
-    boolean isInCheck(BoardCoordinate from, BoardCoordinate to, Board board) {
+    boolean isInCheck(Move move, Board board) {
         boolean isInCheck = false;
-        board.move(from, to);
+        board.move(move);
         outer: for (int y = 0; y < Board.BOARD_SIZE; y++) {
             for (int x = 0; x < Board.BOARD_SIZE; x++) {
                 Piece piece = board.get(x, y);
                 if (piece == null || piece.black == black) continue;
                 ArrayList<BoardCoordinate> legalMoves = piece.getLegalMoves(new BoardCoordinate(x, y), board, false);
-                for (BoardCoordinate move : legalMoves) {
-                    Piece pieceAtMove = board.get(move);
+                for (BoardCoordinate legalMove : legalMoves) {
+                    Piece pieceAtMove = board.get(legalMove);
                     if (pieceAtMove instanceof King) {
                         isInCheck = true;
                         break outer;