path: root/src/grossmeister/search.rs

use std::f32::INFINITY;

use crate::{moves::{Move, MoveKind}, board::io::IO};

use super::{Grossmeister, ttable::{NodeType, TranspositionTableItem}, evaluation::Score};

const SCORE_MATE: Score = 20_000.0;

#[derive(Debug, Default, PartialEq)]
pub struct PerftResult {
    leaf_nodes: u64,
    captures: u64,
    en_passants: u64,
    castles: u64,
    checks: u64,
}

impl Grossmeister {
    // Mate distance pruning
    // If a non-zero value (mating score) is returned, branch could be safely pruned
    pub fn MDP(alpha: &mut Score, beta: &mut Score, root_distance: u8) -> Score {
        { // Update bound in case we are winning
            let mating_score = SCORE_MATE - root_distance as Score;
            if mating_score < *beta {
                *beta = mating_score;
                if *alpha >= mating_score {
                    return mating_score
                }
            }
        }
        { // Update bound in case we are losing
            let mating_score = -SCORE_MATE + root_distance as Score;
            if mating_score > *alpha {
                *alpha = mating_score;
                if *beta <= mating_score {
                    return mating_score
                }
            }
        }
        0.0
    }

    pub fn negamax_search(&mut self, mut alpha: Score, mut beta: Score, depth_left: u8, root_distance: u8) -> Score {
        let color = self.board.color();

        if self.board.threefold_repetition() {
            return 0.0
        }

        if let Some(transposition) = self.transposition() {
            if transposition.depth >= depth_left {
                match transposition.node_type {
                    NodeType::PV => {
                        return transposition.score.clamp(alpha, beta);
                    }
                    NodeType::Cut => {
                        if transposition.score >= beta {
                            return beta;
                        }
                        alpha = transposition.score.clamp(alpha, beta);
                    }
                    NodeType::All => {
                        if transposition.score <= alpha {
                            return alpha;
                        }
                        beta = transposition.score.clamp(alpha, beta);
                    }
                }
            }
        }

        let mut transposition = TranspositionTableItem {
            hash: self.board.hash,
            depth: depth_left,
            mov: None,
            score: alpha,
            node_type: NodeType::All,
        };

        // Mate distance pruning
        let mating_score = Grossmeister::MDP(&mut alpha, &mut beta, root_distance);
        if mating_score != 0.0 {
            return mating_score
        }

        if depth_left == 0 {
            return self.quiscence(alpha, beta, root_distance)
        }

        let mut full_window_search = true;
        let mut legal_move_found = false;

        self.cleanup_selector();
        while let Some(mov) = self.next_move() {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);

            if !self.board.is_king_in_check(color) {
                legal_move_found = true;

                let score = if full_window_search {
                    // Assume PV-node is high in the list (if move ordering is good)
                    -self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 1)
                } else {
                    // After we have PV-node (that raised alpha) all other nodes will be searched
                    // with zero-window first to confirm this assumption
                    let score = -self.negamax_search(-(alpha + 0.01), -alpha, depth_left - 1, root_distance + 1);
                    // In case some of the other nodes raises alpha, then it's true PV node now,
                    // let's research with full window to find its exact value
                    if score > alpha {
                        -self.negamax_search(-beta, -alpha, depth_left - 1, root_distance + 1)
                    } else {
                        score
                    }
                };
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);

                if score >= beta {
                    transposition.mov = Some(mov);
                    transposition.score = score;
                    transposition.node_type = NodeType::Cut;
                    self.store_transposition(transposition);

                    if mov.kind == MoveKind::Quiet {
                        self.register_killer(mov);
                    }

                    return beta
                }
                if score > alpha {
                    alpha = score;
                    full_window_search = false; // Once we have PV-node we can start zero-window searching
                    transposition.mov = Some(mov);
                    transposition.score = score;
                    transposition.node_type = NodeType::PV;
                }
            } else {
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
            }

            // Could not finish in time, return what we have so far
            if root_distance > 0 && self.should_halt.load(std::sync::atomic::Ordering::SeqCst) {
                return alpha;
            }
        }

        if !legal_move_found {
            if self.board.is_king_in_check(color) {
                return -SCORE_MATE + root_distance as Score
            } else {
                transposition.score = 0.0;
                self.store_transposition(transposition);
                return 0.0
            }
        }

        self.store_transposition(transposition);
        alpha
    }

    pub fn quiscence(&mut self, mut alpha: Score, mut beta: Score, root_distance: u8) -> Score {
        let color = self.board.color();

        if self.board.threefold_repetition() {
            return 0.0
        }

        if let Some(transposition) = self.transposition() {
            match transposition.node_type {
                NodeType::PV => {
                    return transposition.score.clamp(alpha, beta);
                }
                NodeType::Cut => {
                    if transposition.score >= beta {
                        return beta;
                    }
                    alpha = transposition.score.clamp(alpha, beta);
                }
                NodeType::All => {
                    if transposition.score <= alpha {
                        return alpha;
                    }
                    beta = transposition.score.clamp(alpha, beta);
                }
            }
        }

        let mut transposition = TranspositionTableItem {
            hash: self.board.hash,
            depth: 0,
            mov: None,
            score: alpha,
            node_type: NodeType::All,
        };

        // Mate distance pruning
        let mating_score = Grossmeister::MDP(&mut alpha, &mut beta, root_distance);
        if mating_score != 0.0 {
            transposition.score = mating_score;
            return mating_score
        }


        let mut tactical_only = false;
        if !self.board.is_king_in_check(color) {
            // If we are not in check, we can evaluate stand pat
            let stand_pat = self.evaluate();

            if stand_pat >= beta {
                return beta;
            }
            if alpha < stand_pat {
                alpha = stand_pat;
            }

            // If we are not in check, we can only search tactical moves
            tactical_only = true;
        }

        let mut legal_move_found = false;
        self.cleanup_selector();
        while let Some(mov) = if tactical_only { self.next_tactical() } else { self.next_move() } {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);

            if !self.board.is_king_in_check(color) {
                legal_move_found = true;
                let evaluation = -self.quiscence(-beta, -alpha, root_distance + 1);
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);

                if evaluation >= beta {
                    transposition.mov = Some(mov);
                    transposition.score = evaluation;
                    transposition.node_type = NodeType::Cut;
                    self.store_transposition(transposition);

                    return beta; // Fail-hard beta-cutoff
                }
                if evaluation > alpha {
                    transposition.mov = Some(mov);
                    transposition.score = evaluation;
                    transposition.node_type = NodeType::Cut;
                    alpha = evaluation;
                } else if transposition.node_type == NodeType::All {
                    transposition.mov = None;
                    transposition.score = evaluation;
                    transposition.node_type = NodeType::All;
                }
            } else {
                self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
            }
        }

        if !legal_move_found && self.board.is_king_in_check(color) {
            return -SCORE_MATE + root_distance as Score
        }

        self.store_transposition(transposition);
        alpha
    }

    fn reconstruct_pv(&mut self, depth: u8) -> Vec<Move> {
        let mut pv = Vec::with_capacity(depth as usize);

        if depth > 0 {
            if let Some(transposition) = self.transposition().cloned() {
                if let Some(mov) = transposition.mov {
                    if transposition.depth < depth {
                        println!("info string found lower-depth {:?}-node ({}) than expected ({}). Probably bad", transposition.node_type, transposition.depth, depth);
                    }
                    if transposition.depth > depth {
                        println!("info string found higher-depth {:?}-node ({}) than expected ({}). Good for us?", transposition.node_type, transposition.depth, depth);
                    }
                    let ep_target_before = self.board.ep_target;
                    let castling_rights_before = self.board.castling_rights;
                    let hash_before = self.board.hash;
                    let captured = self.board.make_move(mov);

                    pv.push(mov);
                    let mut subtree_pv = self.reconstruct_pv(transposition.depth - 1);
                    pv.append(&mut subtree_pv);

                    self.board.unmake_move(mov, captured, ep_target_before, castling_rights_before, hash_before);

                }
            }
        }

        pv
    }

    /// Memory-enhanced Test Driver
    /// Given a guess score (from previous iteration), use it for series of small-window searches
    #[allow(unused)]
    fn mtdf(&mut self, score_guess: Score, depth: u8) -> Score {
        todo!("In order to use, (NodeType::PV ? NodeType::Cut) must be Ordering::Equal, try your luck with mate_in_3 test");
        let window_size = 1.00;
        let mut score = score_guess;
        let mut lowerbound = -INFINITY;
        let mut upperbound = INFINITY;

        while lowerbound < upperbound {
            let beta = score + if score == lowerbound {
                window_size
            } else {
                0.0
            };

            score = self.negamax_search(beta - window_size, beta, depth, 0);

            if score < beta {
                upperbound = score
            } else {
                lowerbound = score
            }
        }

        score
    }

    pub fn iterative_deepening(&mut self, max_depth: u8) -> (Score, Vec<Move>) {
        let mut best_score = 0.0;
        let mut pv = Vec::new();
        let mut depth = 1;

        while depth <= max_depth {
            let score = self.negamax_search(-INFINITY, INFINITY, depth, 0);

            if depth > 1 && self.should_halt.load(std::sync::atomic::Ordering::SeqCst) {
                println!("info string halting search");
                break;
            }

            best_score = score;
            pv = self.reconstruct_pv(depth);

            // Print UCI info
            print!("info depth {} score ", depth);
            if score.abs() >= SCORE_MATE - 128f32 { // TODO: VALUE_WIN - MAX_PLY
                let mate_distance = (SCORE_MATE - score.abs()) * if score > 0.0 {
                    1.
                } else {
                    -1. // -N means we are mated in N plies
                };
                print!("mate {:.0} ", (mate_distance / 2.0).ceil());
            } else {
                print!("cp {:.0} ", score * 100.0)
            }
            print!("pv ");
            for mov in &pv {
                print!("{} ", mov);
            }
            println!();

            depth += 1;

            // If our score is mate in N, we break at depth N
            if depth as Score >= SCORE_MATE - score.abs() {
                break;
            }
        }

        println!("info hashfull {}", self.table_full());
        if !pv.is_empty() {
            print!("bestmove {}", pv[0]);
            if pv.len() > 1 {
                print!(" ponder {}", pv[1])
            }
            println!();
        } else {
            println!("info no PV found");
        }

        (best_score, pv)
    }

    pub fn perft(&mut self, depth: u8, print: bool) -> PerftResult {
        let mut result = PerftResult::default();

        if depth == 0 {
            result.leaf_nodes = 1;
            return result;
        }
        let color = self.board.color();

        if print {
            // println!("Running perft for depth {}. Color to move is {:?}\n{} moves available", depth, color, moves.len());
            // println!("{} moves available", moves.len());
        }

        self.cleanup_selector();
        while let Some(mov) = self.next_move() {
            let ep_target_before = self.board.ep_target;
            let castling_rights_before = self.board.castling_rights;
            let hash_before = self.board.hash;
            let captured_piece = self.board.make_move(mov);
            // King can not be in check after our own move
            if !self.board.is_king_in_check(color) {
                if depth == 1 {
                    match mov.kind {
                        MoveKind::Capture => {
                            result.captures += 1;
                        }
                        MoveKind::EnPassant => {
                            result.en_passants += 1;
                            result.captures += 1;
                        }
                        MoveKind::Castle => {
                            result.castles += 1;
                        }
                        _ => {}
                    }
                    if self.board.is_king_in_check(color.flip()) {
                        result.checks += 1;
                    }
                }

                if print {
                    println!("{:?}", mov);
                    self.board.print();
                }
                let subtree_result = self.perft(depth - 1, print);

                result.leaf_nodes += subtree_result.leaf_nodes;
                result.captures += subtree_result.captures;
                result.checks += subtree_result.checks;
                result.castles += subtree_result.castles;
                result.en_passants += subtree_result.en_passants;

            }
            self.board.unmake_move(mov, captured_piece, ep_target_before, castling_rights_before, hash_before);
        }

        if print {
            println!("Found {} leaf nodes in this subtree (depth {})", result.leaf_nodes, depth);
        }

        result
    }
}


#[cfg(test)]
mod tests {
    use crate::{board::{Board, io::IO}, square::Square, moves::{Move, MoveKind}, grossmeister::{Grossmeister, search::PerftResult}};
    use super::SCORE_MATE;

    #[test]
    fn perft() {
        let board = Board::new();
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(0, false), PerftResult { leaf_nodes: 1, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 20, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 400, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(3, false), PerftResult { leaf_nodes: 8902, captures: 34, en_passants: 0, castles: 0 , checks: 12 });
        assert_eq!(gm.perft(4, false), PerftResult { leaf_nodes: 197281, captures: 1576, en_passants: 0, castles: 0 , checks: 469 });
        // assert_eq!(board.perft(5, false), PerftResult { leaf_nodes: 4865609, captures: 82719, en_passants: 258, castles: 0, checks: 27351 });
    }

    #[test]
    fn position_perft() {
        let fen = String::from("r3k2r/p1ppqpb1/bn2pnp1/3PN3/1p2P3/2N2Q1p/PPPBBPPP/R3K2R w KQkq - ");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(0, false), PerftResult { leaf_nodes: 1, captures: 0, en_passants: 0, castles: 0 , checks: 0 });
        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 48, captures: 8, en_passants: 0, castles: 2 , checks: 0 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 2039, captures: 351, en_passants: 1, castles: 91 , checks: 3 });
        assert_eq!(gm.perft(3, false), PerftResult { leaf_nodes: 97862, captures: 17102, en_passants: 45, castles: 3162, checks: 993 });
        // assert_eq!(board.perft(4, false), PerftResult { leaf_nodes: 4085603, captures: 757163, en_passants: 1929, castles: 128013, checks: 25523 });
    }

    #[test]
    fn endgame_perft() {
        let fen = String::from("8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - ");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);

        assert_eq!(gm.perft(1, false), PerftResult { leaf_nodes: 14, captures: 1, en_passants: 0, castles: 0 , checks: 2 });
        assert_eq!(gm.perft(2, false), PerftResult { leaf_nodes: 191, captures: 14, en_passants: 0, castles: 0 , checks: 10 });
        // assert_eq!(board.perft(3, false), PerftResult { leaf_nodes: 2812, captures: 209, en_passants: 2, castles: 0 , checks: 267 });
    }

    #[test]
    fn mate_in_3() {
        let fen = String::from("2kr1b1r/pp1npppp/2p1bn2/7q/5B2/2NB1Q1P/PPP1N1P1/2KR3R w - - 0 1");
        let board = Board::from_FEN(fen);
        let mut gm = Grossmeister::new(board);
        let (score, pv) = gm.iterative_deepening(8);

        assert_eq!(score, SCORE_MATE - 3.0);
        assert_eq!(pv, vec![
            Move { source: Square::F3, target: Square::C6, kind: MoveKind::Capture },
            Move { source: Square::B7, target: Square::C6, kind: MoveKind::Capture },
            Move { source: Square::D3, target: Square::A6, kind: MoveKind::Quiet },
        ]);
    }

    #[test]
    fn stupid_knight_sac() {
        let fen = String::from("r3k1r1/pp3ppp/1q6/2ppPn2/6P1/1PPP1P2/P1N3KP/R2QR3 b - - 0 18");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        let (_, pv) = gm.iterative_deepening(6);
        assert_eq!(
            pv[0],
            Move { source: Square::F5, target: Square::H4, kind: MoveKind::Quiet },
            "You should save this poor knight from danger!"
        );
    }

    #[test]
    fn weird_bishop_sac() {
        let fen = String::from("r1b1k1nr/p4pp1/1pp1p3/4n2p/1b1qP3/1B1P3N/PPPBQPPP/RN2K2R w KQkq - 7 10");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        let (_, pv) = gm.iterative_deepening(5);
        assert_eq!(
            pv[0],
            Move { source: Square::C2, target: Square::C3, kind: MoveKind::Quiet },
            "You should fork this bastard!"
        );
    }

    #[test]
    fn tricky_mate_them_in_2() {
        let fen = String::from("7k/2P2Q2/3K1b2/p4N1p/P6P/6P1/8/8 w - - 10 61");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(SCORE_MATE - score, 3.0); // Mate in 3 plies
    }

    #[test]
    fn tricky_mate_us_in_2() {
        let fen = String::from("8/2P2Q1k/3K1b2/p4N1p/P6P/6P1/8/8 b - - 5 58");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(score + SCORE_MATE, 4.0); // Mate in 4 plies
    }

    #[test]
    fn another_mate_in_2() {
        let fen = String::from("7r/2pr1pR1/2p2p2/1p2pN2/p3Pk1P/P2P2b1/1PP1K3/R7 w - - 0 33");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert_eq!(SCORE_MATE - score, 3.0); // Mate in 3 plies
    }

    #[test]
    fn unwinnable() {
        let fen = String::from("8/8/2P1k3/1K1n4/8/8/8/8 b - - 0 55");
        let board = Board::from_FEN(fen);

        let mut gm = Grossmeister::new(board);
        gm.debug = true;
        let (score, pv) = gm.iterative_deepening(5);
        dbg!(score, pv);
        assert!(score <= 0.0);
    }

    #[test]
    fn broken_position() {
        let board = Board::new();
        let mut gm = Grossmeister::new(board);

        gm.parse_command(String::from("position startpos moves e2e4 e7e5 g1f3 b8c6 f1b5 a7a6 b5c6 d7c6 b1c3 c8g4 d2d3 g8f6 h2h3 g4f3 d1f3 f8b4 c1g5 b4c3 b2c3 d8e7 d3d4 e7e6 g5f6 g7f6 e1c1 e8c8 c1b2 h8g8 h1g1 b7b6 d4d5 c6d5 e4d5 e6d7 f3f6 d7b5 b2c1 d8d5 g2g4 b5c4 f6f5 c8b7 d1d5 c4d5 f5h7 g8d8 c1b2 d5b5 b2c1 b5e2 h7f5 d8d2 f2f3 d2c2 f5c2"), &mut None);
        gm.board.print();
        let (score, pv) = gm.iterative_deepening(8);
        assert!(score.abs() <= 10_000.0);
        assert!(!pv.is_empty())
    }
}