jja 0.9.1

//
// jja: swiss army knife for chess file formats
// src/brainlearnbook.rs: BrainLearn experience file interface
//
// Copyright (c) 2023, 2024 Ali Polatel <alip@chesswob.org>
//
// SPDX-License-Identifier: GPL-3.0-or-later

use std::{
    cmp::Reverse,
    collections::{HashMap, HashSet},
    fs::File,
    io::Error,
    path::Path,
};

use indicatif::ProgressBar;
use memmap::Mmap;
use shakmaty::{
    fen::{Epd, Fen},
    san::San,
    Chess, Color, EnPassantMode, Position,
};
use termtree::Tree;

use crate::{
    brainlearn::*,
    hash::{ExperienceEntryHashMap, ZobristHashSet, ZobristHasherBuilder},
    stockfish::stockfish_hash,
    tr,
};

/// `BrainLearnFile` is a struct that represents a Brainlearn experience file.
pub struct BrainLearnFile {
    /// The number of entries in the Brainlearn experience file.
    pub num_entries: usize,
    /// The experience data
    pub data: ExperienceEntryHashMap,
    /// The `Mmap` handle to the Brainlearn experience file.
    /// This is `Some` if `num_entries` > 0, `None` otherwise.
    book: Option<Mmap>,
}

impl BrainLearnFile {
    /// Opens a Brainlearn experience file.
    ///
    /// # Arguments
    ///
    /// * `file_name` - The path to the experience file.
    ///
    /// # Returns
    ///
    /// * `Result<Self, Error>` - Returns a `Result` containing a `BrainLearnFile` instance if
    /// successful, or an error if there was an issue opening the file.
    pub fn open<P: AsRef<Path>>(file_name: P) -> Result<Self, Error> {
        let file = File::open(file_name)?;
        let num_bytes = file.metadata()?.len();
        let num_entries = (num_bytes / EXPERIENCE_ENTRY_SIZE as u64) as usize;

        let book = if num_bytes == 0 {
            None
        } else {
            // SAFETY: Mmap::map is unsafe because it involves file I/O which might lead to data races
            // if the underlying file is modified while the memory map is active. Here, it's safe
            // because we assume that the EXP files are not concurrently modified while they're
            // memory-mapped.
            Some(unsafe { Mmap::map(&file)? })
        };

        Ok(Self {
            book,
            num_entries,
            data: HashMap::with_capacity_and_hasher(num_entries, ZobristHasherBuilder),
        })
    }

    /// Loads a Brainlearn experience file into memory.
    ///
    /// # Arguments
    ///
    /// * `depth_cutoff` - Skip moves with depth lower than or equal to this optional value.
    /// * `progress_bar` - An optional progress bar to report progress.
    pub fn load(&mut self, depth_cutoff: Option<u16>, progress_bar: Option<&ProgressBar>) {
        if let Some(pb) = progress_bar {
            pb.println(tr!(
                "Parsing and loading BrainLearn file entries into memory."
            ));
            pb.set_message(tr!("Loading:"));
            pb.set_length(self.num_entries as u64);
            pb.set_position(0);
        }

        if let Some(book) = &self.book {
            for i in 0..self.num_entries {
                let offset = i * EXPERIENCE_ENTRY_SIZE;
                let slice = &book[offset..offset + EXPERIENCE_ENTRY_SIZE];

                let depth = i32::from_le_bytes(slice[8..12].try_into().unwrap());
                if let Some(depth_cutoff) = depth_cutoff {
                    if depth <= i32::from(depth_cutoff) {
                        if let Some(pb) = progress_bar {
                            pb.inc(1);
                        }
                        continue;
                    }
                }

                let key = u64::from_le_bytes(slice[0..8].try_into().unwrap());
                let score = i32::from_le_bytes(slice[12..16].try_into().unwrap());
                let mov = i32::from_le_bytes(slice[16..20].try_into().unwrap());
                let perf = i32::from_le_bytes(slice[20..24].try_into().unwrap());

                let entry = ExperienceEntry {
                    key,
                    depth,
                    score,
                    mov,
                    perf,
                };

                self.data.entry(key).or_default().push(entry);

                if let Some(pb) = progress_bar {
                    pb.inc(1);
                }
            }
        }

        if let Some(pb) = progress_bar {
            pb.println(tr!(
                "Success loading {} BrainLearn file entries into memory.",
                self.num_entries
            ));
            pb.finish_with_message(tr!("Loading done."));
        }
    }

    /// Looks up the moves available in the opening book for the given chess position.
    ///
    /// # Arguments
    ///
    /// * `key: u64` - Zobrist hash of the chess position to look up.
    ///
    /// # Returns
    ///
    /// * `Option<Vec<ExperienceEntry>>` - Returns an `Option` with a vector of `ExperienceEntry`
    /// instances representing the possible moves if found, or `None` if no moves are found.
    pub fn lookup_moves(&self, key: u64) -> Option<Vec<ExperienceEntry>> {
        self.data.get(&key).cloned()
    }

    /// Writes all possible games contained in a Brainlearn experience file to a PGN file.
    ///
    /// This function traverses the Brainlearn file, which is a type of experience file, and writes
    /// all possible games to the output file in PGN format. A game is considered "possible" if it
    /// follows a path of moves in the book from the given starting position to a position with no
    /// more book moves. Each game is written as a separate round, and the rounds are numbered
    /// consecutively starting from 1.
    ///
    /// The `output` argument is a mutable reference to a `Write` trait object where the generated PGN will be written.
    /// The `event`, `site`, `date`, `white`, `black`, and `result` arguments are used to fill in the corresponding PGN tags for each game.
    /// The `max_ply` argument determines the limit of variation depth in plies.
    /// The `progress_bar` is an optional reference to a progress bar to report progress.
    ///
    /// # Errors
    ///
    /// This function will panic if writing to the output file fails.
    ///
    /// # Panics
    ///
    /// Panics if the disk is full or the file isn't writable.
    #[allow(clippy::too_many_arguments)]
    pub fn write_pgn(
        &self,
        output: &mut dyn std::io::Write,
        position: &Chess,
        event: &str,
        site: &str,
        date: &str,
        white: &str,
        black: &str,
        result: &str,
        max_ply: usize,
        progress_bar: Option<&ProgressBar>,
    ) {
        let fen_header: String;
        let fen = if *position == Chess::default() {
            None
        } else {
            fen_header = Fen::from_position(position.clone(), EnPassantMode::Legal).to_string();
            Some(&fen_header)
        };

        if let Some(progress_bar) = progress_bar {
            progress_bar.set_message(tr!("Writing:"));
            progress_bar.set_length(0);
            progress_bar.set_position(0);
        }
        self._write_pgn(
            output,
            position,
            &HashSet::with_hasher(ZobristHasherBuilder),
            &mut Vec::new(),
            fen,
            &mut 1,
            event,
            site,
            date,
            white,
            black,
            result,
            max_ply,
            position.turn(),
            progress_bar,
        );
        if let Some(progress_bar) = progress_bar {
            progress_bar.set_message(tr!("Writing done."));
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn _write_pgn(
        &self,
        output: &mut dyn std::io::Write,
        position: &Chess,
        position_set: &ZobristHashSet,
        move_history: &mut Vec<San>,
        fen: Option<&String>,
        round: &mut usize,
        event: &str,
        site: &str,
        date: &str,
        white: &str,
        black: &str,
        result: &str,
        max_ply: usize,
        initial_color: Color,
        progress_bar: Option<&ProgressBar>,
    ) {
        // Return if the maximum ply is reached
        if move_history.len() >= max_ply {
            return;
        }

        // Each recursive call gets a localized copy of visited positions, preventing global skips.
        // TODO: This is a relatively memory-intensive operation but does the right thing.
        let mut position_set = position_set.clone();

        if let Some(mut entries) = self.lookup_moves(stockfish_hash(position)) {
            // Sort the moves by their weight in reverse order.
            entries.sort_unstable_by_key(|entry| Reverse((entry.depth, entry.score, entry.perf)));

            for entry in entries {
                let mov = match to_move(position, entry.mov) {
                    Some(mov) => mov,
                    None => continue, // TODO: warn about illegal move?
                };
                let san = San::from_move(position, &mov);
                move_history.push(san);
                let mut new_position = position.clone();
                new_position.play_unchecked(&mov);

                // If the new position has been seen before, skip it to avoid infinite recursion.
                let hash = stockfish_hash(&new_position);
                if !position_set.insert(hash) {
                    // Insert returned false, the set already contained this value.
                    move_history.pop();
                    continue;
                }

                // Recursively generate all games starting from the new position.
                self._write_pgn(
                    output,
                    &new_position,
                    &position_set,
                    move_history,
                    fen,
                    round,
                    event,
                    site,
                    date,
                    white,
                    black,
                    result,
                    max_ply,
                    initial_color,
                    progress_bar,
                );

                // Undo the move and remove it from the move history.
                move_history.pop();
            }
        } else {
            // This is a leaf node.
            if !move_history.is_empty() {
                let opening = move_history
                    .iter()
                    .enumerate()
                    .map(|(i, san)| {
                        let move_number = i / 2 + 1;
                        let move_text = san.to_string();
                        match (initial_color, i, i % 2) {
                            (Color::White, _, 0) => format!("{}. {} ", move_number, move_text),
                            (Color::Black, 0, 0) => format!("{}... {} ", move_number, move_text),
                            (Color::Black, _, 1) => format!("{}. {} ", move_number + 1, move_text),
                            _ => format!("{} ", move_text),
                        }
                    })
                    .collect::<String>();

                let fen_header = if let Some(fen) = fen {
                    format!("[FEN \"{}\"]\n[Setup \"1\"]\n", fen)
                } else {
                    String::new()
                };

                writeln!(
                    output,
                    "[Event \"{}\"]\n\
                    [Site \"{}\"]\n\
                    [Date \"{}\"]\n\
                    [Round \"{}\"]\n\
                    [White \"{}\"]\n\
                    [Black \"{}\"]\n\
                    [Result \"{}\"]\n{}\
                    [Annotator \"{} v{}\"]",
                    event,
                    site,
                    date,
                    round,
                    white,
                    black,
                    result,
                    fen_header,
                    crate::built_info::PKG_NAME,
                    crate::built_info::PKG_VERSION,
                )
                .expect("write output PGN");

                writeln!(output, "\n{} {}\n", opening.trim(), result).expect("write output PGN");
                *round += 1;
                if let Some(progress_bar) = progress_bar {
                    progress_bar.inc(1);
                }
            }
        }
    }

    /// A method that generates a tree of moves from a given position using an opening book.
    pub fn tree(&self, position: &Chess, max_ply: u16) -> Tree<String> {
        fn build_tree(
            book: &BrainLearnFile,
            position: &Chess,
            parent: &mut Tree<String>,
            ply: u16,
            max_ply: u16,
            visited_keys: &ZobristHashSet,
        ) {
            if ply >= max_ply {
                return;
            }

            let moves = book.lookup_moves(stockfish_hash(position));
            if let Some(mut book_entries) = moves {
                book_entries
                    .sort_unstable_by_key(|entry| Reverse((entry.depth, entry.score, entry.perf)));
                for entry in book_entries {
                    let mov = entry.mov;
                    let m = match to_move(position, mov) {
                        Some(m) => m,
                        None => continue,
                    };
                    let mut new_position = position.clone();
                    new_position.play_unchecked(&m);

                    let key = stockfish_hash(&new_position);
                    if visited_keys.contains(&key) {
                        continue;
                    }

                    let mut new_visited_keys = visited_keys.clone(); // Clone visited_keys
                    new_visited_keys.insert(key);

                    let mut new_tree = Tree::new(San::from_move(position, &m).to_string());
                    build_tree(
                        book,
                        &new_position,
                        &mut new_tree,
                        ply + 1,
                        max_ply,
                        &new_visited_keys,
                    );

                    parent.push(new_tree);
                }
            }
        }

        let epd = format!(
            "{}",
            Epd::from_position(position.clone(), EnPassantMode::PseudoLegal)
        );
        let mut root_tree = Tree::new(epd);

        let key = stockfish_hash(position);
        let mut visited_keys: ZobristHashSet = HashSet::with_hasher(ZobristHasherBuilder);
        visited_keys.insert(key);

        build_tree(self, position, &mut root_tree, 0, max_ply, &visited_keys);

        root_tree
    }
}