jja 0.9.1

//
// jja: swiss army knife for chess file formats
// src/polyglot.rs: PolyGlot constants and utilities
//
// Copyright (c) 2023 Ali Polatel <alip@chesswob.org>
// Based in part upon PolyGlot 1.4 which is:
//     Copyright 2004-2006 Fabien Letouzey.
//     Licensed under the GNU General Public License, version 2.
// Based in part upon pg_utils-0.2 which is:
//     Copyright 2008 Michel Van den Bergh <michel.vandenbergh@uhasselt.be>.
//     All rights reserved.
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    fmt::{self, Display, Formatter},
    io::{self, Read, Write},
};

use once_cell::sync::Lazy;
use serde::{
    de::{self, Deserialize, Deserializer, SeqAccess, Visitor},
    ser::{Serialize, SerializeTuple, Serializer},
};
use shakmaty::{uci::Uci, Chess, Color, File, Move, Piece, Position, Rank, Role, Square};

use crate::{
    file::{int_from_file, int_to_file},
    tr,
};

/// Size of a Polyglot book entry in bytes.
pub const BOOK_ENTRY_SIZE: usize = std::mem::size_of::<BookEntry>();

/// A static string containing comments about the CSV file format used for Polyglot book entries.
///
/// This string contains detailed information about the file format, UCI moves, and the role of
/// weight and learn fields in the book entries. It can be used as a comment in a CSV file or
/// displayed to users when editing Polyglot book entries.
pub static POLYGLOT_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)
#
# Weight is a measure for the quality of the move.
# Learn is used by some programs to save learning information,
# however this field is commonly unused.
#
# Both weight and learn are positive integers.
# Weight has a maximum of {} whereas Learn 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.
",
        u16::MAX,
        u32::MAX
    )
});

/// A struct representing a Polyglot book entry.
///
/// Polyglot book entries store information about specific positions in a chess game. They are used
/// by chess engines to store and access opening books efficiently.
#[derive(Copy, Clone, Default, Debug)]
#[repr(C, packed)]
pub struct BookEntry {
    /// A 64-bit Zobrist hash of the position.
    pub key: u64,
    /// A 16-bit compact representation of the move in UCI format.
    pub mov: u16,
    /// A 16-bit weight value representing the quality of the move.
    pub weight: u16,
    /// A 32-bit learn value used for learning information (often unused).
    pub learn: u32,
}

impl Display for BookEntry {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(
            f,
            "{}",
            serde_json::to_string(&self).map_err(|_| std::fmt::Error)?
        )
    }
}

impl Serialize for BookEntry {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut tup = serializer.serialize_tuple(4)?;

        let key = self.key;
        tup.serialize_element(&key)?;

        let mov = self.mov;
        tup.serialize_element(&mov)?;

        let weight = self.weight;
        tup.serialize_element(&weight)?;

        let learn = self.learn;
        tup.serialize_element(&learn)?;

        tup.end()
    }
}

struct BookEntryVisitor;

impl<'de> Visitor<'de> for BookEntryVisitor {
    type Value = BookEntry;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a four-element tuple")
    }

    fn visit_seq<A>(self, mut seq: A) -> Result<BookEntry, A::Error>
    where
        A: SeqAccess<'de>,
    {
        let key = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(0, &self))?;
        let mov = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(1, &self))?;
        let weight = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(2, &self))?;
        let learn = seq
            .next_element()?
            .ok_or_else(|| de::Error::invalid_length(3, &self))?;
        Ok(BookEntry {
            key,
            mov,
            weight,
            learn,
        })
    }
}

impl<'de> Deserialize<'de> for BookEntry {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_tuple(4, BookEntryVisitor)
    }
}

/// A struct representing a Polyglot book entry with minimal information.
#[derive(Copy, Clone, Default, Debug)]
#[repr(C, packed)]
pub struct CompactBookEntry {
    /// A 16-bit compact representation of the move in UCI format.
    pub mov: u16,
    /// A 16-bit weight value representing the quality of the move.
    pub weight: u16,
    /// A 32-bit learn value used for learning information (often unused).
    pub learn: u32,
}

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

/// `ColorWeight` is a structure to store weight associated with different types of colors
/// associated with CTG moves in a chess game, using the recommendation field.
pub struct ColorWeight {
    /// Represents a green color move.
    pub green: u16,
    /// Represents a blue color move.
    pub blue: u16,
    /// Represents a red color move.
    pub red: u16,
}

/// Formats a list of opening book entries as a CSV-like string.
///
/// This function takes a reference to a `Chess` position and a vector of
/// `BookEntry` 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 learn value
/// of a book entry, separated by commas.
///
/// # Arguments
///
/// * `position` - A reference to a `Chess` position.
/// * `entries` - A vector of `BookEntry` objects to format.
///
/// # Returns
///
/// A `String` containing the formatted book entries in a CSV-like format.
pub fn format_bin_entries(position: &Chess, entries: Vec<BookEntry>) -> String {
    let mut ret = String::new();

    ret.push_str("*,uci,weight,learn\n");
    let mut i = 1;
    for entry in &entries {
        if let Some(mv) = to_move(position, entry.mov) {
            let uci = Uci::from_standard(&mv);
            let weight = entry.weight;
            let learn = entry.learn;
            ret.push_str(&format!("{},{},{},{}\n", i, uci, weight, learn));
        }
        i += 1;
    }

    ret
}

/// Writes a `BookEntry` to a file in a binary format.
///
/// This function takes a mutable reference to a type implementing the `Write` trait and a reference
/// to a `BookEntry` object. It writes the `BookEntry` to the file in a binary format.
///
/// # Arguments
///
/// * `f` - A mutable reference to a type implementing the `Write` trait.
/// * `entry` - A reference to the `BookEntry` object to write to the file.
///
/// # Returns
///
/// A `Result<(), io::Error>` indicating success or failure.
pub fn bin_entry_to_file<W: Write>(mut f: W, entry: &BookEntry) -> io::Result<()> {
    int_to_file(&mut f, 8, entry.key)?;
    int_to_file(&mut f, 2, u64::from(entry.mov))?;
    int_to_file(&mut f, 2, u64::from(entry.weight))?;
    int_to_file(&mut f, 4, u64::from(entry.learn))?;
    Ok(())
}

/// Reads a `BookEntry` from a file in a binary format.
///
/// This function takes a mutable reference to a type implementing the `Read` trait and reads a
/// `BookEntry` object from the file in a binary format.
///
/// # Arguments
///
/// * `f` - A mutable reference to a type implementing the `Read` trait.
///
/// # Returns
///
/// A `Result<BookEntry, io::Error>` containing the `BookEntry` object read from the file or an
/// error if reading fails.
pub fn bin_entry_from_file<R: Read>(mut f: R) -> io::Result<BookEntry> {
    let key = int_from_file(&mut f, 8)?;
    let mov = int_from_file(&mut f, 2)? as u16;
    let weight = int_from_file(&mut f, 2)? as u16;
    let learn = int_from_file(&mut f, 4)? as u32;
    Ok(BookEntry {
        key,
        mov,
        weight,
        learn,
    })
}

/// Converts a UCI move to its compact 16-bit representation.
///
/// This function takes a `Uci` object and converts it to a 16-bit integer representation.
/// The `king_on_start_square` argument is used to distinguish between castling and non-castling
/// moves.
///
/// # Arguments
///
/// * `uci` - The UCI move to be converted.
/// * `king_on_start_square` - `true` if the side to move has their king on the starting square,
/// E1, or E8 respectively.
///
/// # Returns
///
/// A `u16` representing the compact form of the UCI move.
pub fn from_uci(uci: Uci, king_on_start_square: bool) -> u16 {
    match uci {
        // For a Null move, we return 0
        Uci::Null => 0,

        // For a Normal move, we will convert it into a 16-bit integer
        Uci::Normal {
            from,
            to,
            promotion,
        } => {
            // Get the file and rank of the 'from' square as 16-bit integers
            let ff = from.file() as u16;
            let fr = from.rank() as u16;

            // Determine the file and rank of the 'to' square
            let (tf, tr) =
                // If it's a possible white castling move, and the king is indeed on the 'from' square
                if king_on_start_square && from == Square::E1 && (to == Square::G1 || to == Square::C1) {
                    // If the 'to' square is G1, it's a kingside castling move
                    // Otherwise, it's a queenside castling move
                    if to == Square::G1 { (File::H as u16, Rank::First as u16) } else { (File::A as u16, Rank::First as u16) }
                // If it's a possible black castling move, and the king is indeed on the 'from' square
                } else if king_on_start_square && from == Square::E8 && (to == Square::G8 || to == Square::C8) {
                    // If the 'to' square is G8, it's a kingside castling move
                    // Otherwise, it's a queenside castling move
                    if to == Square::G8 { (File::H as u16, Rank::Eighth as u16) } else { (File::A as u16, Rank::Eighth as u16) }
                // For non-castling moves, we just use the file and rank of the 'to' square
                } else {
                    (to.file() as u16, to.rank() as u16)
                };

            // For the promotion piece, we convert each type to a unique integer
            let p = match promotion {
                Some(Role::Knight) => 1,
                Some(Role::Bishop) => 2,
                Some(Role::Rook) => 3,
                Some(Role::Queen) => 4,
                // If there's no promotion, we use 0
                _ => 0,
            };

            // We now combine all the pieces of information into a single 16-bit integer
            (p << 12) + (fr << 9) + (ff << 6) + (tr << 3) + tf
        }
        // Uci::Put is not supported and should not be used with this function
        Uci::Put { .. } => unreachable!(),
    }
}

/// Converts a `shakmaty::Move` to a compact 16-bit representation.
///
/// This function takes a `shakmaty::Move` object and converts it to a 16-bit integer representation.
/// It supports normal moves, en passant, and castling moves. The `Move::Put` variant is not supported
/// and should not be used with this function.
///
/// # Arguments
///
/// * `mov` - Reference to a `shakmaty::Move` to be converted.
///
/// # Returns
///
/// A `u16` representing the compact form of the `shakmaty::Move`.
pub fn from_move(mov: &Move) -> u16 {
    // Match the shakmaty::Move variant
    match mov {
        // Normal move (including promotions)
        Move::Normal {
            from,
            to,
            promotion,
            ..
        } => {
            // Get the file and rank of the from and to squares as u16
            let ff = from.file() as u16;
            let fr = from.rank() as u16;
            let tf = to.file() as u16;
            let tr = to.rank() as u16;

            // Convert the promotion piece to its corresponding value
            let p = match promotion {
                Some(Role::Knight) => 1,
                Some(Role::Bishop) => 2,
                Some(Role::Rook) => 3,
                Some(Role::Queen) => 4,
                _ => 0,
            };

            // Combine the promotion value and the file/rank values to create the compact representation
            (p << 12) + (fr << 9) + (ff << 6) + (tr << 3) + tf
        }
        // En passant move
        Move::EnPassant { from, to } => {
            // Get the file and rank of the from and to squares as u16
            let ff = from.file() as u16;
            let fr = from.rank() as u16;
            let tf = to.file() as u16;
            let tr = to.rank() as u16;

            // Combine the file/rank values to create the compact representation
            (fr << 9) + (ff << 6) + (tr << 3) + tf
        }
        // Castling move
        Move::Castle { king, rook } => {
            // Castling moves are represented somewhat unconventially as follows
            // (this convention is related to Chess960, see below).
            // white short      e1h1
            // white long       e1a1
            // black short      e8h8
            // black long       e8a8
            let kf = king.file() as u16;
            let kr = king.rank() as u16;
            let tf = rook.file() as u16;
            let tr = rook.rank() as u16;

            // Combine the file/rank values to create the compact representation
            (kr << 9) + (kf << 6) + (tr << 3) + tf
        }
        // Move::Put is not supported and should not be used with this function
        Move::Put { .. } => unreachable!(),
    }
}

/// Finds the corresponding `Move` object for a given compact book move.
///
/// This function takes a reference to a `Position` and a 16-bit integer representing a compact
/// book move. It returns the corresponding `Move` object if the move is legal in the given position.
///
/// # Arguments
///
/// * `position` - A reference to a type implementing the `Position` trait.
/// * `book_move` - A `u16` representing the compact form of the book move.
///
/// # Returns
///
/// An `Option<Move>` containing the corresponding `Move` object if it is legal in the given position,
/// otherwise `None`.
pub fn to_move(position: &dyn Position, book_move: u16) -> Option<Move> {
    let book_move = u32::from(book_move);

    let to_ = Square::from_coords(
        File::new(book_move & 0x7),
        Rank::new((book_move >> 3) & 0x07),
    );
    let from_ = Square::from_coords(
        File::new((book_move >> 6) & 0x7),
        Rank::new((book_move >> 9) & 0x07),
    );
    let promotion_role = match book_move >> 12 {
        0 => None,
        1 => Some(Role::Knight),
        2 => Some(Role::Bishop),
        3 => Some(Role::Rook),
        4 => Some(Role::Queen),
        _ => {
            return None;
        }
    };

    let moves = position.legal_moves();

    for chess_move in moves {
        match chess_move {
            Move::Normal {
                from,
                to,
                promotion,
                ..
            } => {
                if from == from_ && to == to_ && promotion == promotion_role {
                    return Some(chess_move);
                }
            }
            Move::Castle { king, rook } => {
                if king == from_ && rook == to_ {
                    return Some(chess_move);
                }
            }
            Move::EnPassant { from, to } => {
                if from == from_ && to == to_ {
                    return Some(chess_move);
                }
            }
            _ => {
                unreachable!(
                    "What's going on with this move? `{:?}' (from:{} to:{}), please report a bug.",
                    chess_move, from_, to_
                );
            } /* No Move::Put in opening book :) */
        };
    }

    None
}

/// Returns the index of a given `Piece` object.
///
/// This function takes a `Piece` object and returns its index based on its color and role.
///
/// # Arguments
///
/// * `piece` - The `Piece` object to find the index for.
///
/// # Returns
///
/// A `u8` representing the index of the given `Piece` object.
pub fn piece_index(piece: Piece) -> u8 {
    match piece {
        Piece {
            color: Color::Black,
            role: Role::Pawn,
        } => 0,
        Piece {
            color: Color::White,
            role: Role::Pawn,
        } => 1,
        Piece {
            color: Color::Black,
            role: Role::Knight,
        } => 2,
        Piece {
            color: Color::White,
            role: Role::Knight,
        } => 3,
        Piece {
            color: Color::Black,
            role: Role::Bishop,
        } => 4,
        Piece {
            color: Color::White,
            role: Role::Bishop,
        } => 5,
        Piece {
            color: Color::Black,
            role: Role::Rook,
        } => 6,
        Piece {
            color: Color::White,
            role: Role::Rook,
        } => 7,
        Piece {
            color: Color::Black,
            role: Role::Queen,
        } => 8,
        Piece {
            color: Color::White,
            role: Role::Queen,
        } => 9,
        Piece {
            color: Color::Black,
            role: Role::King,
        } => 10,
        Piece {
            color: Color::White,
            role: Role::King,
        } => 11,
    }
}

/// Takes a `shakmaty::Chess` position as argument and returns `true` if the side to move has their
/// king on the starting square, `false` otherwise.
pub fn is_king_on_start_square(position: &Chess) -> bool {
    let turn = position.turn();
    position
        .board()
        .king_of(turn)
        .map(|square| square == turn.fold_wb(Square::E1, Square::E8))
        .unwrap_or(false)
}