3 files changed, 64 insertions, 12 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;
 }
diff --git a/src/bot.hpp b/src/bot.hpp
index 2a007db..382c9be 100644
--- a/src/bot.hpp
+++ b/src/bot.hpp
@@ -4,7 +4,7 @@
 #include "board/board.hpp"
 #include <vector>
 Move EngineGetBestMove(Game *b, int depth);
-float EvaluateBoardForWhite(Game *b, int depth);
+float EvaluateBoardForWhite(Game *b);
 float minimax(int depth, Game *b, float alpha, float beta);
 int ScoreMove(const Game *board, const Move &move);
 std::vector<Move> GetSortedLegalMoves(Game *g);
diff --git a/src/moves.cpp b/src/moves.cpp
index c5c5b85..d5ad989 100644
--- a/src/moves.cpp
+++ b/src/moves.cpp
@@ -164,16 +164,19 @@ std::vector<Move> GetLegalMoves(Game *g) {
     }
   }
 
-  return legalMoves;
+  return legalMoves;    
 }
 GameState GetNewGameState(Game *g) {
-  auto legalMoves = GetLegalMoves(g);
-
   // make sure we dont override game ending states
   if (g->state == DRAW || g->state == STALEMATE || g->state == WHITE_WON ||
       g->state == BLACK_WON) {
     return g->state;
   }
+  if (g->halfMoveClock >= 50) {
+    g->state = DRAW;
+  }
+
+  auto legalMoves = GetLegalMoves(g);
 
   if (legalMoves.empty()) {
     Position kingPosition = FindKing(g, g->turn);
@@ -186,7 +189,7 @@ GameState GetNewGameState(Game *g) {
       g->state = STALEMATE;
     }
   }
-  return TURN;
+  return g->state;
 }
 bool IsSquareAttacked(Game *g, Position square, bool white) {
   int target = PositionToIndex(square);