path: root/src/grossmeister/move_selector.rs

use smallvec::SmallVec;

use crate::{board::move_generation::MoveList, moves::Move};

use super::{ttable::NodeType, Grossmeister};

pub type ScoredMove = (Move, f32);
pub type ScoredMoveList = SmallVec<[ScoredMove; 128]>;

#[derive(Debug, Clone, Default)]
pub struct ScoredMoveIter {
    pub moves: ScoredMoveList,
    index: usize,
}

impl Iterator for ScoredMoveIter {
    type Item = ScoredMove;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < self.moves.len() {
            // Perform a selection sort
            let mut best_score = self.moves[self.index].1;
            for i in (self.index + 1)..self.moves.len() {
                if self.moves[i].1 > best_score {
                    best_score = self.moves[i].1;
                    self.moves.swap(i, self.index);
                }
            }

            self.index += 1;
            return Some(self.moves[self.index - 1]);
        }
        None
    }
}

impl FromIterator<ScoredMove> for ScoredMoveIter {
    fn from_iter<T: IntoIterator<Item = ScoredMove>>(iter: T) -> Self {
        ScoredMoveIter {
            moves: iter.into_iter().collect(),
            index: 0,
        }
    }
}

#[derive(Debug, Clone)]
pub enum MoveGenStage {
    Hash,
    WinningOrEqualTactical,
    Killer,
    Quiet,
    LosingTactical,
}

impl Default for MoveGenStage {
    fn default() -> Self {
        Self::Hash
    }
}

impl MoveGenStage {
    pub fn next_stage(&self) -> Self {
        match self {
            MoveGenStage::Hash => MoveGenStage::WinningOrEqualTactical,
            MoveGenStage::WinningOrEqualTactical => MoveGenStage::Killer,
            MoveGenStage::Killer => MoveGenStage::Quiet,
            MoveGenStage::Quiet => MoveGenStage::LosingTactical,
            MoveGenStage::LosingTactical => todo!(),
        }
    }
}

#[derive(Default, Debug, Clone)]
pub struct MoveSelector {
    tactical_moves: ScoredMoveList,
    stage_moves: ScoredMoveIter,
    pub killer_moves: SmallVec<[Move; 4]>,
    stage: MoveGenStage,
}

impl Grossmeister {
    /// Return move selector for the current ply
    pub fn move_selector(&mut self) -> &mut MoveSelector {
        // dbg!(self.board.ply, self.move_selectors.len());
        &mut self.move_selectors[self.board.ply as usize]
    }

    pub fn cleanup_selector(&mut self) {
        let selector = self.move_selector();
        selector.tactical_moves.clear();
        selector.stage_moves.moves.clear();
        selector.stage_moves.index = 0;
        selector.stage = MoveGenStage::default();
    }

    /// Register killer for ply-before
    pub fn register_killer(&mut self, killer: Move) {
        if self.board.ply > 1 {
            let parent_killers =
                &mut self.move_selectors[(self.board.ply - 2) as usize].killer_moves;
            match parent_killers.iter().find(|m| **m == killer) {
                None => {
                    parent_killers.push(killer);
                    debug_assert!(
                        !parent_killers.spilled(),
                        "Killer move list should remain on the stack"
                    );
                    // We want to have max 3 killers, so if we exceed the limit remove the oldest killer
                    if parent_killers.len() > 3 {
                        parent_killers.remove(0);
                    }
                }
                Some(..) => {}
            }
        }
    }

    /// Evaluate move for move ordering, prioritizing efficient captures
    /// where low-value pieces capture high-value pieces
    pub fn eval_move(&self, m: Move) -> f32 {
        if m.is_tactical() {
            let [source_eval, target_eval] = [m.source, m.target]
                .map(|sq| self.board.piece_by_square(sq))
                .map(|p| match p {
                    Some(p) => p.static_eval(),
                    None => 0.,
                });
            return 2. * target_eval - source_eval;
        }
        0.0
    }

    pub fn score_moves(&self, movelist: MoveList) -> ScoredMoveList {
        movelist.iter().map(|&m| (m, self.eval_move(m))).collect()
    }

    pub fn next_tactical(&mut self) -> Option<Move> {
        if self.move_selector().stage_moves.moves.is_empty() {
            let moves = self.board.generate_moves_core(true);
            let scored = self.score_moves(moves);
            self.init_stage(scored);
        }

        self.move_selector().stage_moves.next().map(|(mov, _)| mov)
    }

    fn init_stage(&mut self, moves: ScoredMoveList) {
        self.move_selector().stage_moves = ScoredMoveIter { moves, index: 0 }
    }

    /// TODO: next stage
    fn next_stage(&mut self) {
        let selector = self.move_selector();
        selector.stage = selector.stage.next_stage();
        selector.stage_moves = ScoredMoveIter::default();
        // println!("NEXT STAGE {:?}", selector.stage);
    }

    pub fn next_move(&mut self) -> Option<Move> {
        loop {
            match self.move_selector().stage {
                MoveGenStage::Hash => {
                    self.next_stage();
                    if let Some(transposition) = self.transposition() {
                        if transposition.node_type != NodeType::All {
                            if let Some(mov) = transposition.mov {
                                return Some(mov);
                            }
                        }
                    }
                }
                MoveGenStage::WinningOrEqualTactical => {
                    // Init stage moves
                    if self.move_selector().stage_moves.moves.is_empty() {
                        // Generate all tactical moves (for future re-use)
                        let moves = self.board.generate_moves_core(true);
                        self.move_selector().tactical_moves = self.score_moves(moves);

                        // But we only care about current stage now
                        let new_stage = self
                            .move_selector()
                            .tactical_moves
                            .iter()
                            .filter(|(_, score)| *score >= 0.0)
                            .copied()
                            .collect();

                        self.init_stage(new_stage);
                    }

                    match self.move_selector().stage_moves.next() {
                        Some((mov, _)) => return Some(mov),
                        None => self.next_stage(),
                    }
                }
                MoveGenStage::Killer => {
                    if self.move_selector().stage_moves.moves.is_empty() {
                        let new_stage = self
                            .move_selector()
                            .killer_moves
                            .clone()
                            .iter()
                            .filter(|&m| {
                                self.move_selector()
                                    .tactical_moves
                                    .iter()
                                    .any(|(m2, _)| m2 == m)
                            }) // Test if killer is in the movelist
                            .map(|&m| (m, 0.0))
                            .collect();
                        self.init_stage(new_stage);
                    }
                    match self.move_selector().stage_moves.next() {
                        Some((mov, _)) => return Some(mov),
                        None => self.next_stage(),
                    }
                }
                MoveGenStage::Quiet => {
                    if self.move_selector().stage_moves.moves.is_empty() {
                        let moves = self.board.generate_moves_core(false);
                        // No need to score quiet moves? TODO: doublecheck
                        let scored = moves.iter().map(|&m| (m, 0.0)).collect();
                        self.init_stage(scored);
                    }
                    match self.move_selector().stage_moves.next() {
                        Some((mov, _)) => return Some(mov),
                        None => self.next_stage(),
                    }
                }
                MoveGenStage::LosingTactical => {
                    if self.move_selector().stage_moves.moves.is_empty() {
                        let new_stage = self
                            .move_selector()
                            .tactical_moves
                            .iter()
                            .filter(|(_, score)| *score < 0.0)
                            .copied()
                            .collect();
                        self.init_stage(new_stage);
                    }
                    match self.move_selector().stage_moves.next() {
                        Some((mov, _)) => return Some(mov),
                        None => return None,
                    }
                }
            }
        }
    }
}