Change PSQT at endgame

src/evaluation/evaluation.rs

@@ -1,11 +1,15 @@
 use u64 as Bitboard;
 
 use crate::{
-    board::bitboard::lsb,
-    board::board::{Board, Color, PieceType},
+    board::{
+        bitboard::{self, lsb},
+        board::{Board, Color, PieceType},
+    },
     evaluation::psqt::{mirror_index, piece_square_score},
 };
 
+use super::psqt::piece_square_score_endgame;
+
 const fn piece_score(piece_type: PieceType) -> i32 {
     match piece_type {
         PieceType::Pawn => 100,
@@ -17,6 +21,21 @@ const fn piece_score(piece_type: PieceType) -> i32 {
     }
 }
 
+// TODO: Use piece scores instead of piece count
+fn is_end_game(board: &Board) -> bool {
+    let white_pieces_num = board
+        .white_pieces
+        .iter()
+        .fold(0, |acc, p| acc + bitboard::bit_count(p.bitboard));
+
+    let black_pieces_num = board
+        .black_pieces
+        .iter()
+        .fold(0, |acc, p| acc + bitboard::bit_count(p.bitboard));
+
+    white_pieces_num < 6 && black_pieces_num < 6
+}
+
 fn evaluate_side_for(board: &Board, color: Color) -> i32 {
     let mut total_score = 0;
     let pieces = match color {
@@ -24,6 +43,12 @@ fn evaluate_side_for(board: &Board, color: Color) -> i32 {
         Color::Black => &board.black_pieces,
     };
 
+    let psqt = if !is_end_game(board) {
+        piece_square_score
+    } else {
+        piece_square_score_endgame
+    };
+
     for piece in pieces {
         let (piece_type, mut bitboard): (PieceType, Bitboard) = (piece.piece_type, piece.bitboard);
         let mut score = 0;
@@ -34,7 +59,7 @@ fn evaluate_side_for(board: &Board, color: Color) -> i32 {
             };
 
             score += piece_score(piece_type);
-            score += piece_square_score(piece_type, psqt_index);
+            score += psqt(piece_type, psqt_index);
             bitboard &= bitboard - 1;
         }
         total_score += score;
@@ -46,7 +71,10 @@ pub fn evaluate_position(board: &Board) -> i32 {
     let total_white_score = evaluate_side_for(board, Color::White);
     let total_black_score = evaluate_side_for(board, Color::Black);
 
-    -(total_white_score - total_black_score)
+    match board.state.current_player() {
+        Color::White => total_white_score - total_black_score,
+        Color::Black => total_black_score - total_white_score,
+    }
 }
 
 #[cfg(test)]
@@ -96,17 +124,9 @@ mod tests {
     }
 
     #[test]
-    fn test_evaluate() -> Result<(), String> {
-        let game = from_fen(FEN_QUIET[0])?;
-
-        assert_eq!(
-            evaluate_position(&game.board),
-            evaluate_position(&game.board)
-        );
-
-        let game_2 = from_fen(FEN_QUIET[1])?;
-
-        assert_eq!(-40, evaluate_position(&game_2.board));
+    fn test_evaluate_position() -> Result<(), String> {
+        let game = from_fen(FEN_QUIET[1])?;
+        assert_eq!(40, evaluate_position(&game.board));
 
         Ok(())
     }