jja 0.9.1

//
// jja: swiss army knife for chess file formats
// src/abk.rs: Arena book constants and utilities
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
// Based in part upon ChessX's src/database/*.cpp which is:
//     Copyright (C) 2015, 2017 by Jens Nissen jens-chessx@gmx.net
// Based in part upon DroidFish's AbkBook.java which is
//     Copyright (C) 2020  Peter Österlund, peterosterlund2@gmail.com
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    cmp::PartialEq,
    collections::{BTreeMap, VecDeque},
};

use once_cell::sync::Lazy;
use shakmaty::{uci::Uci, Chess, Position, Role, Square};

use crate::{
    abkbook::AbkBook,
    hash::zobrist_hash,
    polyglot::{from_uci, is_king_on_start_square, CompactBookEntry},
    tr,
};

/// Offset of Arena Book move entries in Arena opening books
pub const ABK_INDEX_OFFSET: usize = 900;
/// Size of a single Arena Book move entry
pub const ABK_ENTRY_LENGTH: usize = std::mem::size_of::<SBookMoveEntry>();

/// Chess squares
pub const FIELDS: [&str; 64] = [
    "a1", "b1", "c1", "d1", "e1", "f1", "g1", "h1", "a2", "b2", "c2", "d2", "e2", "f2", "g2", "h2",
    "a3", "b3", "c3", "d3", "e3", "f3", "g3", "h3", "a4", "b4", "c4", "d4", "e4", "f4", "g4", "h4",
    "a5", "b5", "c5", "d5", "e5", "f5", "g5", "h5", "a6", "b6", "c6", "d6", "e6", "f6", "g6", "h6",
    "a7", "b7", "c7", "d7", "e7", "f7", "g7", "h7", "a8", "b8", "c8", "d8", "e8", "f8", "g8", "h8",
];

/// A static string containing comments about the CSV file format used for Arena book entries.
///
/// This string contains detailed information about the file format, UCI moves, and the role of
/// priority, ngames, nwon and nlost fields in the book entries. It can be used as a comment in
/// a CSV file or displayed to users when editing Arena book entries.
pub static ABK_EDIT_COMMENT: Lazy<String> = Lazy::new(|| {
    tr!(
        "#
# The file is in CSV (comma-separated-values) format.
# Lines starting with `#' are comments and ignored.
#
# Moves are given in UCI (universal chess interface) format
# which is a variation of a long algebraic format for chess
# moves commonly used by chess engines.
#
# Examples:
#   e2e4, e7e5, e1g1 (white short castling), e7e8q (for promotion)
#
# Priority is a measure for the quality of the move.
# Ngames, Nwon and Nlost are game statistics about the move.
#
# All of priority, ngames, nwon and nlost are positive integers.
# Priority has a maximum of {}.
# Ngames, Nwon and Nlost has a maximum of {}.
#
# Edit the file as you like, the moves you deleted will be removed from the book.
# If you delete all the moves the position will be removed from the book.
# Exit without saving to abort the action.
",
        u8::MAX,
        i32::MAX
    )
});

/// `NagPriority` is a structure to store priorities associated with
/// different types of moves in a chess game, using Numeric Annotation Glyphs (NAGs).
/// Each field is a u8 representing the priority associated with the corresponding move type.
pub struct NagPriority {
    /// Represents a good move, also denoted by "!".
    pub good: u8,
    /// Represents a mistake in the move, also denoted by "?".
    pub mistake: u8,
    /// Represents a hard-to-find good move, also denoted by "!!".
    pub hard: u8,
    /// Represents a significantly bad move, a major mistake or howler, also denoted by "??".
    pub blunder: u8,
    /// Represents a speculative or interesting move, traditionally denoted by "!?", Nunn Convention.
    pub interesting: u8,
    /// Represents a questionable or dubious move, traditionally denoted by "?!", Nunn Convention.
    pub dubious: u8,
    /// Represents a forced move, a move that is the only reasonable option in a given position.
    pub forced: u8,
    /// Represents a singular move, a move that has no reasonable alternatives.
    pub only: u8,
}

/// `ColorPriority` is a structure to store priorities associated with different types of moves in
/// a chess game, using CTG colored move recommendations. Each field is a u8 representing the
/// priority associated with the corresponding move color.
pub struct ColorPriority {
    /// Represents a green color move.
    pub green: u8,
    /// Represents a blue color move.
    pub blue: u8,
    /// Represents a red color move.
    pub red: u8,
}

/// `CompactSBookMoveEntry` represents an entry in the opening book with minimal information.
///
/// The struct contains information about a move, such as the origin square, the target square, the
/// promotion piece, and other statistics (priority, number of games, wins, losses, etc.).
#[derive(Clone, Copy, Debug)]
#[repr(C, packed)]
pub struct CompactSBookMoveEntry {
    /// The origin square of the move, from 0 to 63 (a1 to h8).
    pub from: u8, // a1 0, b1 1, ..., h1 7, ... h8 63
    /// The target square of the move, from 0 to 63 (a1 to h8).
    pub to: u8, // a1 0, b1 1, ..., h1 7, ... h8 63
    /// The promotion piece, if any.
    /// Values: 0 for none, ±1 for rook, ±2 for knight, ±3 for bishop, ±4 for queen.
    pub promotion: i8, // 0 none, +-1 rook, +-2 knight, +-3 bishop, +-4 queen
    /// The priority of the move.
    pub priority: u8, // 0 = bad, >0 better, 9 best
    /// The number of games in which the move was played.
    pub ngames: u32, // Number of times games in which move was played
    /// The number of games in which the move led to a win.
    pub nwon: u32, // Number of won games for white
    /// The number of games in which the move led to a loss.
    pub nlost: u32, // Number of lost games for white
}

/// `SBookMoveEntry` represents an entry in the opening book.
///
/// The struct contains information about a move, such as the origin square, the target square, the
/// promotion piece, and other statistics (priority, number of games, wins, losses, etc.).
#[derive(Clone, Copy, Debug)]
#[repr(C, packed)]
pub struct SBookMoveEntry {
    /// The origin square of the move, from 0 to 63 (a1 to h8).
    pub from: u8, // a1 0, b1 1, ..., h1 7, ... h8 63
    /// The target square of the move, from 0 to 63 (a1 to h8).
    pub to: u8, // a1 0, b1 1, ..., h1 7, ... h8 63
    /// The promotion piece, if any.
    /// Values: 0 for none, ±1 for rook, ±2 for knight, ±3 for bishop, ±4 for queen.
    pub promotion: i8, // 0 none, +-1 rook, +-2 knight, +-3 bishop, +-4 queen
    /// The priority of the move.
    pub priority: u8, // 0 = bad, >0 better, 9 best
    /// The number of games in which the move was played.
    pub ngames: u32, // Number of times games in which move was played
    /// The number of games in which the move led to a win.
    pub nwon: u32, // Number of won games for white
    /// The number of games in which the move led to a loss.
    pub nlost: u32, // Number of lost games for white
    /// Move flags, value is 0x01000000 if move has been deleted
    pub flags: u32,
    /// The index of the next move in the book.
    pub next_move: u32, // First following move (by opposite color)
    /// The index of the next sibling move in the book.
    pub next_sibling: u32, // Next alternative move (by same color)
}

impl From<SBookMoveEntry> for CompactSBookMoveEntry {
    fn from(entry: SBookMoveEntry) -> Self {
        CompactSBookMoveEntry {
            from: entry.from,
            to: entry.to,
            promotion: entry.promotion,
            priority: entry.priority,
            ngames: entry.ngames,
            nwon: entry.nwon,
            nlost: entry.nlost,
        }
    }
}

impl PartialEq for SBookMoveEntry {
    fn eq(&self, other: &Self) -> bool {
        self.from == other.from && self.to == other.to && self.promotion == other.promotion
    }
}

/// Implement conversion from `SBookMoveEntry` to `Uci`.
///
/// This implementation allows converting an opening book move entry into a UCI format move.
impl From<SBookMoveEntry> for Uci {
    fn from(entry: SBookMoveEntry) -> Uci {
        if entry.from >= 64 || entry.to >= 64 {
            // There're entries such as the one below in HS-Book.abk:
            // SBookMoveEntry { from: 255, to: 255, promotion: 0, priority: 5, ngames: 1, nwon: 0,
            // nlost: 0, flags: 0, next_move: 4294967295, next_sibling: 4294967295 }
            // page end?
            return Uci::Null;
        }
        /* SAFETY: The check in the if condition above guarantees
         * `entry.from' and `entry.to' are within limits.
         */
        let (from, to) = unsafe {
            (
                Square::new_unchecked(u32::from(entry.from)),
                Square::new_unchecked(u32::from(entry.to)),
            )
        };
        let promotion = match entry.promotion {
            1 | -1 => Some(Role::Rook),
            2 | -2 => Some(Role::Knight),
            3 | -3 => Some(Role::Bishop),
            4 | -4 => Some(Role::Queen),
            _ => None,
        };
        Uci::Normal {
            from,
            to,
            promotion,
        }
    }
}

/// Implement conversion from `CompactSBookMoveEntry` to `Uci`.
///
/// This implementation allows converting an opening book move entry into a UCI format move.
impl From<CompactSBookMoveEntry> for Uci {
    fn from(entry: CompactSBookMoveEntry) -> Uci {
        if entry.from >= 64 || entry.to >= 64 {
            // There're entries such as the one below in HS-Book.abk:
            // SBookMoveEntry { from: 255, to: 255, promotion: 0, priority: 5, ngames: 1, nwon: 0,
            // nlost: 0, flags: 0, next_move: 4294967295, next_sibling: 4294967295 }
            // page end?
            return Uci::Null;
        }
        /* SAFETY: The check in the if condition above guarantees
         * `entry.from' and `entry.to' are within limits.
         */
        let (from, to) = unsafe {
            (
                Square::new_unchecked(u32::from(entry.from)),
                Square::new_unchecked(u32::from(entry.to)),
            )
        };
        let promotion = match entry.promotion {
            1 | -1 => Some(Role::Rook),
            2 | -2 => Some(Role::Knight),
            3 | -3 => Some(Role::Bishop),
            4 | -4 => Some(Role::Queen),
            _ => None,
        };
        Uci::Normal {
            from,
            to,
            promotion,
        }
    }
}

impl Default for SBookMoveEntry {
    fn default() -> Self {
        Self::new()
    }
}

impl SBookMoveEntry {
    /// Creates a new `SBookMoveEntry` with default values.
    ///
    /// # Returns
    ///
    /// * `SBookMoveEntry` - A new `SBookMoveEntry` instance with the following default values:
    /// - `from: 255`
    /// - `to: 255`
    /// - `promotion: 0`
    /// - `priority: 0`
    /// - `ngames: 0`
    /// - `nwon: 0`
    /// - `nlost: 0`
    /// - `flags: 0`
    /// - `next_move: u32::MAX`
    /// - `next_sibling: u32::MAX`
    pub fn new() -> Self {
        SBookMoveEntry {
            from: 255,
            to: 255,
            promotion: 0,
            priority: 0,
            ngames: 0,
            nwon: 0,
            nlost: 0,
            flags: 0,
            next_move: u32::MAX,
            next_sibling: u32::MAX,
        }
    }

    /// Calculates the weight of a `SBookMoveEntry`.
    ///
    /// The weight is computed as follows:
    /// * If the priority is greater than 0, the weight is the product of the priority and 2500,
    /// limited by the maximum value representable by a u16.
    /// * If the priority is 0, the weight is the difference between the number of won games and
    /// the number of lost games, limited by the minimum and maximum values representable by a u16.
    pub fn weight(&self) -> u16 {
        if self.priority > 0 {
            u16::from(self.priority).saturating_mul(2500)
        } else if self.nwon > 0 {
            self.nwon
                .try_into()
                .unwrap_or(u16::MAX)
                .saturating_sub(self.nlost.try_into().unwrap_or(u16::MAX))
        } else {
            self.ngames.try_into().unwrap_or(u16::MAX)
        }
    }

    /// Merges two `SBookMoveEntry` instances.
    ///
    /// The other instance is merged into self. If the other instance has a higher
    /// number of games or priority, the corresponding fields in self will be updated.
    /// If either self or other is a deleted entry, the function will handle it accordingly.
    ///
    /// # Arguments
    ///
    /// * other: `SBookMoveEntry` - The other `SBookMoveEntry` instance to be merged with self.
    pub fn merge(&mut self, other: SBookMoveEntry) {
        if self.flags == 0x01000000 && other.flags != 0x01000000 {
            /* self is a deleted entry */
            *self = other;
            return;
        } else if other.flags == 0x01000000 {
            /* other is a deleted entry */
            return;
        }
        if other.ngames > self.ngames {
            self.ngames = other.ngames;
            self.nwon = other.nwon;
            self.nlost = other.nlost;
        }
        if other.priority > self.priority {
            self.priority = other.priority;
        }
    }

    /// Checks if the `SBookMoveEntry` is marked as deleted.
    ///
    /// # Returns
    ///
    /// * `bool` - Returns `true` if the entry is marked as deleted, otherwise `false`.
    pub fn is_deleted(&self) -> bool {
        self.flags == 1 || self.flags == 0x01000000
    }

    /// Checks if the `SBookMoveEntry` is selected based on the given
    /// `AbkBook` instance.
    ///
    /// The function considers the `probability_games` and
    /// `probability_win_percent` fields of the `AbkBook` to determine if
    /// the `SBookMoveEntry` should be selected.
    ///
    /// # Arguments
    ///
    /// * `book: &AbkBook` - A reference to an `AbkBook` instance used for
    /// the selection criteria.
    ///
    /// # Returns
    ///
    /// * `bool` - Returns `true` if the entry meets the selection criteria
    /// based on the `AbkBook`, otherwise `false`.
    pub fn is_selected(&self, book: &AbkBook) -> bool {
        if book.probability_games > 0 && self.ngames < book.probability_games {
            return false;
        }

        if book.probability_win_percent > 0 {
            let win_percent = if self.ngames > 0 {
                (f64::from(self.nwon) / f64::from(self.ngames)) * 100.0
            } else {
                0.0
            };
            if win_percent < book.probability_win_percent.into() {
                return false;
            }
        }

        true
    }
}

/// Formats a list of opening book entries as a CSV-like string.
///
/// This function takes a vector of `SBookMoveEntry` objects, and returns a formatted string
/// representing the book entries in a CSV-like format. Each line in the output string contains the
/// index, UCI notation of the move, weight, and priority value of a book entry, separated by
/// commas.
///
/// # Arguments
///
/// * `entries` - A vector of `SBookMoveEntry` objects to format.
///
/// # Returns
///
/// A `String` containing the formatted book entries in a CSV-like format.
pub fn format_abk_entries(entries: Vec<SBookMoveEntry>) -> String {
    let mut ret = String::new();

    ret.push_str("*,uci,priority,ngames,nwon,nlost\n");
    for (i, entry) in entries.iter().enumerate() {
        let uci = Uci::from(*entry);
        let ngames = entry.ngames;
        let nwon = entry.nwon;
        let nlost = entry.nlost;
        ret.push_str(&format!(
            "{},{},{},{},{},{}\n",
            i + 1,
            uci,
            entry.priority,
            ngames,
            nwon,
            nlost,
        ));
    }

    ret
}

/// Traverses a tree of moves and builds a `BTreeMap` containing the Zobrist hash of the
/// position as the key and the corresponding moves as the value.
///
/// # Arguments
///
/// * `tree` - A reference to a `BTreeMap` containing the Zobrist hash of positions and
/// their corresponding moves.
/// * `pos` - A `shakmaty::Chess` position.
///
/// # Returns
///
/// Returns a `BTreeMap` containing the Zobrist hash of positions and their corresponding moves.
pub fn traverse_tree(
    tree: &BTreeMap<u64, Vec<SBookMoveEntry>>,
    pos: Chess,
) -> BTreeMap<u64, Vec<CompactBookEntry>> {
    let mut book = BTreeMap::new();
    let mut queue = VecDeque::new();

    let initial_zobrist = zobrist_hash(&pos);
    queue.push_back((pos, initial_zobrist));
    while let Some((pos, key)) = queue.pop_front() {
        if let Some(entries) = tree.get(&key) {
            let mut book_entries = Vec::with_capacity(entries.len());

            for entry in entries {
                let uci = Uci::from(*entry);
                let chess_move = match uci.to_move(&pos) {
                    Ok(move_) => move_,
                    Err(_) => {
                        /* illegal uci, mmph */
                        continue;
                    }
                };
                let mov = from_uci(uci, is_king_on_start_square(&pos));

                let book_entry = CompactBookEntry {
                    mov,
                    weight: entry.weight(),
                    learn: u32::from(entry.priority),
                };

                // Perform a binary search to find the right position for
                // the new entry based on its weight. We use
                // `std::cmp::Reverse` because we want higher weights to
                // come first, but `binary_search_by_key` finds positions
                // for ascending order by default.
                match book_entries.binary_search_by_key(
                    &std::cmp::Reverse(book_entry.weight),
                    |entry: &CompactBookEntry| {
                        // Extract and reverse the priority of the existing move for comparison.
                        std::cmp::Reverse(entry.weight)
                    },
                ) {
                    // If the exact weight is found, we get its position
                    // with Ok(pos), otherwise, Err(pos) gives the position
                    // where it should be inserted to maintain sort order.
                    Ok(pos) | Err(pos) => {
                        // Insert the new entry at the determined position.
                        book_entries.insert(pos, book_entry);
                    }
                };

                let mut pos = pos.clone();
                pos.play_unchecked(&chess_move);
                let new_zobrist = zobrist_hash(&pos);
                if !book.contains_key(&new_zobrist) {
                    queue.push_back((pos, new_zobrist));
                }
            }

            book.insert(key, book_entries);
        }
    }

    book
}