2 files changed, 38 insertions, 15 deletions

diff --git a/src/board/engine.rs b/src/board/engine.rs
index 61d9fda..9b0e36c 100644
--- a/src/board/engine.rs
+++ b/src/board/engine.rs
@@ -17,7 +17,7 @@ impl Board {
             result.leaf_nodes = 1;
             return result;
         }
-        let color = self.color_to_move();
+        let color = self.color();
 
         let moves = self.generate_pseudolegal_moves(color);
 
@@ -73,24 +73,37 @@ impl Board {
 
         result
     }
-    pub fn evaluate(&self) -> f32 {
+
+    /// Compute material advantage relative to the current player
+    pub fn material_advantage(&self) -> f32 {
         let mut eval = 0f32;
-        let pieces = self.pieces_by_color(self.color());
-        eval += pop_count(pieces[PieceType::Pawn as usize]) as f32;
-        eval += pop_count(pieces[PieceType::Bishop as usize]) as f32 * 3.;
-        eval += pop_count(pieces[PieceType::Knight as usize]) as f32 * 3.;
-        eval += pop_count(pieces[PieceType::Rook as usize]) as f32 * 4.5;
-        eval += pop_count(pieces[PieceType::Queen as usize]) as f32 * 9.;
+        for (piece_index, bitboard) in self.pieces.iter().enumerate() {
+            let piece_type = PieceType::from(piece_index);
+            let sign = if Color::from_piece(piece_type) == self.color() {
+                1.
+            } else {
+                -1.
+            };
+            let multiplier = match piece_type.without_color() {
+                PieceType::Pawn => 1.,
+                PieceType::Bishop => 3.,
+                PieceType::Knight => 3.,
+                PieceType::Rook => 4.5,
+                PieceType::Queen => 9.,
+                PieceType::King => 0.,
+                _ => panic!("Piece should be without color"),
+            };
+
+            eval += sign * multiplier * pop_count(*bitboard) as f32;
+        }
         eval
     }
 
     pub fn negamax_search(&mut self, mut alpha: f32, beta: f32, depth_left: u8) -> f32 {
         let color = Color::from(self.ply as u8 % 2);
 
-        self.print();
-
         if depth_left == 0 {
-            return self.evaluate();
+            return self.material_advantage();
         }
         let moves = self.generate_pseudolegal_moves(color);
 
@@ -156,4 +169,11 @@ mod tests {
         let eval = board.negamax_search(-INFINITY, INFINITY, 4);
         println!("{}", eval);
     }
+
+    #[test]
+    fn material_advantage() {
+        let board = Board::new();
+        assert_eq!(board.material_advantage(), 0.0, "Material advantage should be 0 at starting position");
+
+    }
 }
diff --git a/src/board/mod.rs b/src/board/mod.rs
index e2d4d60..5e2f4e3 100644
--- a/src/board/mod.rs
+++ b/src/board/mod.rs
@@ -44,6 +44,13 @@ pub enum PieceType {
     KingBlack,
 }
 
+impl PieceType {
+    pub fn without_color(&self) -> Self {
+        let index = *self as usize;
+        Self::from(index % 6)
+    }
+}
+
 const PIECE_CHARS: [&str; 12] = [
     "♟︎", "♞", "♝", "♜", "♛", "♚",
     "♙", "♘", "♗", "♖", "♕", "♔",
@@ -147,10 +154,6 @@ impl Board {
         occupancy
     }
 
-    pub fn color_to_move(&self) -> Color {
-        Color::from((self.ply % 2) as u8)
-    }
-
     pub fn print(&self) {
         println!();
         for rank in (0..8).rev() {