1 files changed, 56 insertions, 7 deletions

diff --git a/src/bot.cpp b/src/bot.cpp
index 908b418..c3458bb 100644
--- a/src/bot.cpp
+++ b/src/bot.cpp
@@ -163,6 +163,9 @@ Move EngineGetBestMove(Game *b, int depth) {
 float minimax(int depth, Game *b, float alpha, float beta) {
   uint64_t gameHash = GenerateZobristKey(b);
 
+  if (b->ThreeFoldMap[gameHash] >= 2) {
+    return b->turn ? -500 : 500;
+  }
   if (b->Transpositions->contains(gameHash)) {
     TranspositionsEntry data = b->Transpositions->at(gameHash);
     if (data.depth >= depth) {
@@ -172,7 +175,7 @@ float minimax(int depth, Game *b, float alpha, float beta) {
 
   auto moves = GetSortedLegalMoves(b);
   if (depth == 0 || moves.size() == 0) {
-    return EvaluateBoardForWhite(b, depth);
+    return EvaluateBoardForWhite(b);
   }
 
   bool shouldStore = true;
@@ -218,16 +221,53 @@ float minimax(int depth, Game *b, float alpha, float beta) {
 }
 int PSTIndex(int square, bool white) { return white ? square : (56 ^ square); }
 
-float EvaluateBoardForWhite(Game *g, int depth) {
+static int ForceKingToEdgeBonus(Position enemyKing, Position myKing) {
+  int bonus = 0;
+
+  // Push enemy king toward edge
+  int distToCenter =
+      std::abs(enemyKing.file - 3) + std::abs(enemyKing.rank - 3);
+
+  bonus += distToCenter * 10;
+
+  // Bring own king closer
+  int kingDistance = std::abs(myKing.file - enemyKing.file) +
+                     std::abs(myKing.rank - enemyKing.rank);
+
+  bonus += (14 - kingDistance) * 5;
+
+  return bonus;
+}
+static bool IsEndgame(Game *g) {
+  int queens = 0;
+  int rooks = 0;
+
+  for (int i = 0; i < 64; i++) {
+    switch (g->pieces[i].type) {
+    case QUEEN:
+      queens++;
+      break;
+    case ROOK:
+      rooks++;
+      break;
+    default:
+      break;
+    }
+  }
+
+  return queens == 0 || (queens == 2 && rooks <= 1);
+}
+
+float EvaluateBoardForWhite(Game *g) {
   float score = 0;
 
   GameState state = GetNewGameState(g);
   switch (state) {
   case WHITE_WON:
-    return MATE + depth;
+    return MATE;
     break;
   case BLACK_WON:
-    return -MATE - depth;
+    return -MATE;
     break;
   case TURN:
     break;
@@ -236,7 +276,6 @@ float EvaluateBoardForWhite(Game *g, int depth) {
   case DRAW:
     return 0;
   }
-
   for (int i = 0; i < 64; i++) {
     Piece piece = g->pieces[i];
     if (piece.type == NONEPIECE)
@@ -272,10 +311,20 @@ float EvaluateBoardForWhite(Game *g, int depth) {
       break;
     }
 
-    if (piece.color)
+    if (piece.color) {
       score += value;
-    else
+    } else {
       score -= value;
+    }
+    Position whiteKing = FindKing(g, true);
+    Position blackKing = FindKing(g, false);
+
+    if (IsEndgame(g)) {
+      score +=
+          ForceKingToEdgeBonus(blackKing, whiteKing); // White attacking black
+      score -=
+          ForceKingToEdgeBonus(whiteKing, blackKing); // Black attacking white
+    }
   }
   return score;
 }