jja 0.9.1

//
// jja: swiss army knife for chess file formats
// src/random.rs: Random playouts using books
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{collections::HashMap, sync::Arc};

use rand::{seq::SliceRandom, SeedableRng};
use rand_xorshift::XorShiftRng;
use shakmaty::{fen::Epd, Chess, EnPassantMode, Move, Outcome, Position};

use crate::{hash::zobrist_hash, polyglot::to_move, polyglotbook::PolyGlotBook, tr};

/// An enumeration representing the different move selection strategies available.
#[derive(Copy, Clone, Debug)]
pub enum MoveSelection {
    /// Selects the best move based on the move weight.
    BestMove,
    /// Selects a move uniformly at random from the available moves.
    UniformRandom,
    /// Selects a move with a probability proportional to its weight.
    WeightedRandom,
}

/// Plays a random chess game starting from the given position with the specified `initial_move`.
///
/// # Arguments
///
/// * `position` - A mutable reference to a `Chess` object representing the current position.
/// * `initial_move` - A reference to a `Move` object representing the initial move to be played.
///
/// # Returns
///
/// Returns a `shakmaty::Outcome` indicating the outcome of the random game, such as a win, draw, or loss.
pub fn play_random_game(position: &mut Chess, initial_move: &Move) -> Outcome {
    // Set HashMap for 3-position detection.
    let mut rep = HashMap::<u64, u8>::new();
    rep.insert(zobrist_hash(position), 1);

    // Play the initial move.
    position.play_unchecked(initial_move);

    // Play the random game
    let mut rng = XorShiftRng::from_entropy();
    loop {
        match position.outcome() {
            Some(outcome) => return outcome,
            None => {
                // Step 1: Check 50-move rule.
                if position.halfmoves() >= 50 {
                    return Outcome::Draw;
                }

                // Step 2: Repetition detection.
                let key = zobrist_hash(position);
                let val = rep.entry(key).and_modify(|v| *v += 1).or_insert(1);
                if *val >= 3 {
                    return Outcome::Draw;
                }

                // Fall through, the game is not over.
            }
        }

        // Step 1: Play the random move.
        let legal_moves: Vec<Move> = position.legal_moves().as_slice().to_vec();
        let random_move = legal_moves.choose(&mut rng).unwrap();
        position.play_unchecked(random_move);
    }
}

/// Plays a random game using two `PolyGlotBook` instances, switching between them on each move.
///
/// # Arguments
///
/// * `position`: The initial `Chess` position to start the game from.
/// * `book1`: The first `PolyGlotBook` instance.
/// * `book2`: The second `PolyGlotBook` instance.
/// * `move_selection`: The `MoveSelection` strategy used to select moves from the opening books.
/// * `prefer_irreversible`: If `true`, prefer irreversible moves when out of the book.
///
/// # Returns
///
/// * `Result<Outcome, Box<dyn std::error::Error>>`: The outcome of the game or an error if the
/// chosen move cannot be converted.
pub fn play_random_book_game(
    position: &Chess,
    book1: Arc<PolyGlotBook>,
    book2: Arc<PolyGlotBook>,
    move_selection: &MoveSelection,
    prefer_irreversible: bool,
) -> Result<Outcome, Box<dyn std::error::Error>> {
    let mut pos = position.clone();
    let mut is_book1_turn = true;
    let mut book1_miss = false;
    let mut book2_miss = false;
    let mut rng = XorShiftRng::from_entropy();

    // Set HashMap for 3-position detection.
    let mut rep = HashMap::<u64, u8>::new();
    rep.insert(zobrist_hash(position), 1);

    loop {
        let book = if is_book1_turn { &book1 } else { &book2 };
        let move_ = if (is_book1_turn && !book1_miss) || (!is_book1_turn && !book2_miss) {
            match book.lookup_moves(zobrist_hash(&pos)) {
                Some(mut entries) => {
                    let chosen = match move_selection {
                        MoveSelection::BestMove => {
                            entries.sort_by(|a, b| {
                                let aw = a.weight;
                                let bw = b.weight;
                                bw.cmp(&aw) /* reverse sort by weight */
                            });
                            &entries[0]
                        }
                        MoveSelection::UniformRandom => entries.choose(&mut rng).unwrap(),
                        MoveSelection::WeightedRandom => {
                            // We add +1 to the weight here to avoid AllWeightsZero errors.
                            entries.choose_weighted(&mut rng, |entry| {
                                u32::from(entry.weight).saturating_add(1)
                            })?
                        }
                    };
                    to_move(&pos, chosen.mov).ok_or_else(|| {
                        let epd = format!(
                            "{}",
                            Epd::from_position(pos.clone(), EnPassantMode::PseudoLegal)
                        );
                        std::io::Error::new(
                            std::io::ErrorKind::InvalidData,
                            tr!(
                                "Failed to convert the chosen entry `{:?}' in position `{}'.",
                                chosen,
                                epd
                            ),
                        )
                    })?
                }
                _ => {
                    if is_book1_turn {
                        book1_miss = true;
                    } else {
                        book2_miss = true;
                    }
                    let legal_moves = pos.legal_moves().as_slice().to_vec();

                    // Filter irreversible moves if prefer_irreversible is true
                    let irreversible_moves: Option<Vec<_>> = if prefer_irreversible {
                        Some(
                            legal_moves
                                .iter()
                                .filter(|mov| pos.is_irreversible(mov))
                                .cloned()
                                .collect(),
                        )
                    } else {
                        None
                    };

                    // Choose from irreversible_moves if it exists and
                    // is not empty, otherwise choose from legal_moves
                    if let Some(irreversible_moves) = &irreversible_moves {
                        if !irreversible_moves.is_empty() {
                            irreversible_moves.choose(&mut rng).unwrap().clone()
                        } else {
                            legal_moves.choose(&mut rng).unwrap().clone()
                        }
                    } else {
                        legal_moves.choose(&mut rng).unwrap().clone()
                    }
                }
            }
        } else {
            let legal_moves = pos.legal_moves().as_slice().to_vec();
            legal_moves.choose(&mut rng).unwrap().clone()
        };

        pos.play_unchecked(&move_);

        match position.outcome() {
            Some(outcome) => return Ok(outcome),
            None => {
                // Step 1: Check 50-move rule.
                if position.halfmoves() >= 50 {
                    return Ok(Outcome::Draw);
                }

                // Step 2: Repetition detection.
                let key = zobrist_hash(position);
                let val = rep.entry(key).and_modify(|v| *v += 1).or_insert(1);
                if *val >= 3 {
                    return Ok(Outcome::Draw);
                }

                // Fall through, the game is not over.
            }
        }

        is_book1_turn = !is_book1_turn;
    }
}