chess_dragon/engine/src/mcts.rs

// use crate::networks::PolicyNetwork;
use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
use crate::net::encoding::{encode_board_state_perspective, encode_move};
use crate::net::model::ChessModel;
use burn::prelude::Device;
use burn::Tensor;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;

pub struct Node {
    pub prior: f32,
    pub children: Vec<usize>,
    pub visit_count: u32,
    pub value_sum: f32,
    pub board_state: BoardState,
    pub last_move: Option<ChessMove>, // move that produced this node from its parent
}

impl Node {
    pub fn new(prior: f32, board_state: BoardState, last_move: Option<ChessMove>) -> Node {
        Node {
            prior,
            children: vec![],
            visit_count: 0,
            value_sum: 0.0,
            board_state,
            last_move,
        }
    }

    pub fn value(&self) -> f32 {
        self.value_sum / self.visit_count as f32
    }

    pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
        self.children
            .iter()
            .copied()
            .max_by(|&a, &b| {
                ucb_score(self, &arena[a], c_puct)
                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
                    .unwrap()
            })
            .expect("select_child on leaf")
    }
}

impl Clone for Node {
    fn clone(&self) -> Self {
        Self::new(self.prior, self.board_state.clone(), self.last_move.clone())
    }
}

#[derive(Clone, Debug)]
pub struct MctsResults {
    pub board_state: BoardState,
    // Compact encoded move distribution: (encoded_move_index, probability)
    pub move_dist: Vec<(usize, f32)>,
    pub value: f32,
}

impl MctsResults {
    pub fn new(board_state: BoardState, move_dist: Vec<(usize, f32)>, value: f32) -> MctsResults {
        MctsResults {
            board_state,
            move_dist,
            value,
        }
    }
}

#[derive(Debug)]
pub struct MctsConfig {
    pub num_simulations: usize,
    pub c_puct: f32,
    pub dirichlet_alpha: f32,
    pub dirichlet_epsilon: f32,
    pub batch_max: usize,
}

impl MctsConfig {
    pub fn new(
        num_simulations: usize,
        c_puct: f32,
        dirichlet_alpha: f32,
        dirichlet_epsilon: f32,
        batch_max: usize,
    ) -> MctsConfig {
        MctsConfig {
            num_simulations,
            c_puct,
            dirichlet_alpha,
            dirichlet_epsilon,
            batch_max,
        }
    }
}

impl Default for MctsConfig {
    fn default() -> MctsConfig {
        MctsConfig::new(400, 1.0, 0.05, 0.25, 64)
    }
}

pub struct Mcts {
    pub config: MctsConfig,
}

fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
    match status {
        BoardStateStatus::Ongoing => None,
        BoardStateStatus::Stalemate
        | BoardStateStatus::Threefold
        | BoardStateStatus::FiftyMove => Some(0.0),
        BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }),
        BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }),
    }
}

fn softmax(logits: &[f32]) -> Vec<f32> {
    if logits.is_empty() {
        return Vec::new();
    }
    let mut max_logit = f32::NEG_INFINITY;
    for &v in logits {
        if v > max_logit {
            max_logit = v;
        }
    }
    if !max_logit.is_finite() {
        max_logit = 0.0;
    }
    let mut exps_sum = 0.0f32;
    let mut probs = vec![0.0f32; logits.len()];
    for (i, &v) in logits.iter().enumerate() {
        let x = (v - max_logit).exp();
        if x.is_finite() {
            probs[i] = x;
            exps_sum += x;
        }
    }
    if exps_sum > 0.0 && exps_sum.is_finite() {
        for p in &mut probs {
            *p /= exps_sum;
        }
    } else {
        let uniform = 1.0 / logits.len() as f32;
        for p in &mut probs {
            *p = uniform;
        }
    }
    probs
}

impl Mcts {
    pub fn search(
        &mut self,
        board_state: &BoardState,
        model: &ChessModel,
        device: &Device,
    ) -> MctsResults {
        let mut nodes = Vec::<Node>::new();

        let root = 0;
        nodes.push(Node::new(0.0, board_state.clone(), None));

        // Expand root to create initial children and priors
        self.expand(root, &mut nodes, model, device);

        // Apply Dirichlet noise to root children
        self.add_dirichlet_noise(root, &mut nodes);

        let mut sims_done: usize = 0;
        let num_sims = self.config.num_simulations;

        while sims_done < num_sims {
            let mut leaf_nodes: Vec<usize> = Vec::new();
            let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
            let mut leaf_states: Vec<Tensor<4>> = Vec::new();

            let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);

            while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims {
                let mut path = vec![root];
                let mut current = root;

                // Apply virtual loss during selection path
                while !nodes[current].children.is_empty() {
                    nodes[current].visit_count += 1;
                    nodes[current].value_sum -= 1.0;

                    current = nodes[current].select_child(&nodes, &self.config.c_puct);
                    path.push(current);
                }

                // Check terminal status of leaf
                let stm = nodes[current].board_state.board.side_to_move();
                if let Some(term_val) = terminal_value(nodes[current].board_state.status, stm) {
                    // Revert virtual loss for path except leaf itself
                    for &idx in &path[..path.len() - 1] {
                        nodes[idx].visit_count -= 1;
                        nodes[idx].value_sum += 1.0;
                    }
                    self.backpropagate(&mut nodes, &path, term_val, stm);
                    sims_done += 1;
                    continue;
                }

                // Prevent expanding the exact same node twice in a single batch
                if leaf_nodes.contains(&current) {
                    // Revert virtual loss for this duplicate path
                    for &idx in &path[..path.len() - 1] {
                        nodes[idx].visit_count -= 1;
                        nodes[idx].value_sum += 1.0;
                    }
                    break;
                }

                leaf_nodes.push(current);
                leaf_paths.push(path);

                let state: Tensor<4> =
                    encode_board_state_perspective(&nodes[current].board_state, device)
                        .reshape([1, 18, 8, 8]);
                leaf_states.push(state);
            }

            if leaf_nodes.is_empty() {
                if sims_done >= num_sims {
                    break;
                }
                continue;
            }

            // Batch evaluate non-terminal leaf states
            let batch = Tensor::cat(leaf_states, 0);
            let (policy_batch, value_batch) = model.forward(batch);

            let policy_data = policy_batch.into_data().to_vec::<f32>().unwrap();
            let value_data = value_batch.into_data().to_vec::<f32>().unwrap();

            let num_moves = policy_data.len() / leaf_nodes.len();

            for (i, &node_idx) in leaf_nodes.iter().enumerate() {
                let path = &leaf_paths[i];

                // Revert virtual loss along this path before applying true expansion & backpropagation
                for &idx in &path[..path.len() - 1] {
                    nodes[idx].visit_count -= 1;
                    nodes[idx].value_sum += 1.0;
                }

                let start = i * num_moves;
                let end = start + num_moves;
                let logits = &policy_data[start..end];
                let probs = softmax(logits);

                let stm = nodes[node_idx].board_state.board.side_to_move();
                let legal_moves: Vec<ChessMove> =
                    MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();

                let mut move_priors = Vec::with_capacity(legal_moves.len());
                let mut prior_sum = 0.0f32;

                for &mv in &legal_moves {
                    if let Some(enc_idx) = encode_move(mv, stm) {
                        let p = probs[enc_idx];
                        move_priors.push((mv, p));
                        prior_sum += p;
                    }
                }

                for (mv, p) in move_priors {
                    let normalized_prior = if prior_sum > 0.0 {
                        p / prior_sum
                    } else {
                        1.0 / legal_moves.len() as f32
                    };

                    let mut new_board = nodes[node_idx].board_state.clone();
                    new_board.apply_move(mv);

                    let child_idx = nodes.len();
                    nodes.push(Node::new(normalized_prior, new_board, Some(mv)));
                    nodes[node_idx].children.push(child_idx);
                }

                let value = value_data[i];
                self.backpropagate(&mut nodes, path, value, stm);
                sims_done += 1;
            }
        }

        // Build compact move distribution
        let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
        let stm = board_state.board.side_to_move();
        let total_root_visits: u32 = nodes[root].children.iter().map(|&c| nodes[c].visit_count).sum();
        let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 };
        for child_idx in nodes[root].children.iter() {
            let mv = nodes[*child_idx].last_move.expect("move didnt exist");
            if let Some(enc) = encode_move(mv, stm) {
                let prob = nodes[*child_idx].visit_count as f32 / denom;
                move_dist.push((enc, prob));
            }
        }

        MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
    }

    fn expand(
        &mut self,
        node_idx: usize,
        arena: &mut Vec<Node>,
        model: &ChessModel,
        device: &Device,
    ) -> f32 {
        let status = arena[node_idx].board_state.status;
        let stm = arena[node_idx].board_state.board.side_to_move();

        if let Some(term_val) = terminal_value(status, stm) {
            term_val
        } else {
            let state: Tensor<4> =
                encode_board_state_perspective(&arena[node_idx].board_state, device)
                    .reshape([1, 18, 8, 8]);
            let (policy_head, value_head) = model.forward(state);

            let legal_moves: Vec<ChessMove> =
                MoveGen::new_legal(&arena[node_idx].board_state.board).collect();

            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
            let probs = softmax(&logits);

            let mut move_priors = Vec::with_capacity(legal_moves.len());
            let mut prior_sum = 0.0f32;

            for &mv in &legal_moves {
                if let Some(enc_idx) = encode_move(mv, stm) {
                    let p = probs[enc_idx];
                    move_priors.push((mv, p));
                    prior_sum += p;
                }
            }

            for (mv, p) in move_priors {
                let normalized_prior = if prior_sum > 0.0 {
                    p / prior_sum
                } else {
                    1.0 / legal_moves.len() as f32
                };

                let mut new_board = arena[node_idx].board_state.clone();
                new_board.apply_move(mv);

                let child_idx = arena.len();
                arena.push(Node::new(normalized_prior, new_board, Some(mv)));
                arena[node_idx].children.push(child_idx);
            }

            value_head.into_data().to_vec::<f32>().unwrap()[0]
        }
    }

    fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) {
        for &idx in path {
            let node = &mut nodes[idx];

            if node.board_state.board.side_to_move() == color {
                node.value_sum += value;
            } else {
                node.value_sum -= value;
            }

            node.visit_count += 1;
        }
    }

    fn add_dirichlet_noise(&mut self, node_id: usize, nodes: &mut Vec<Node>) {
        let node_children = &nodes[node_id].children.clone();

        let n = node_children.len();
        if n == 0 {
            return;
        }

        let alpha = self.config.dirichlet_alpha;
        let epsilon = self.config.dirichlet_epsilon;

        let noise = dirichlet_sample(n, alpha);

        for (node_idx, noise_val) in node_children.iter().zip(noise) {
            let prev_prior = nodes[*node_idx].prior;
            nodes[*node_idx].prior = (1.0 - epsilon) * prev_prior + epsilon * noise_val;
        }
    }
}

fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
    use rand_distr::{Distribution, Gamma};

    let gamma = Gamma::new(alpha as f64, 1.0).unwrap();

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap();
    let seed = now.as_nanos() as u64;
    let mut rng = rand::rngs::SmallRng::seed_from_u64(seed);

    let mut samples: Vec<f32> = (0..size).map(|_| gamma.sample(&mut rng) as f32).collect();

    let sum: f32 = samples.iter().sum();

    if sum > 0.0 && sum.is_finite() {
        for x in &mut samples {
            *x /= sum;
        }
    } else if size > 0 {
        let uniform = 1.0 / size as f32;
        for x in &mut samples {
            *x = uniform;
        }
    }

    samples
}

fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
    let prior_score = c_puct * child.prior * (parent.visit_count as f32).sqrt()
        / (1.0 + child.visit_count as f32);

    let value_score: f32;
    if child.visit_count > 0 {
        value_score = -child.value(); // value from opposing side
    } else {
        value_score = 0.0;
    }

    value_score + prior_score
}

pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
    // Choose best child of root by visit count (standard MCTS behavior)
    let root_node = &nodes[root];
    if root_node.children.is_empty() {
        return None;
    }

    let mut best_visits = 0u32;
    let mut best_child_idx: Option<usize> = None;
    for &child_idx in &root_node.children {
        let child = &nodes[child_idx];
        if child.visit_count > best_visits {
            best_visits = child.visit_count;
            best_child_idx = Some(child_idx);
        }
    }

    let child_idx = match best_child_idx {
        Some(i) => i,
        None => return None,
    };

    let bn = &nodes[child_idx];

    println!("best_child visits: {}", bn.visit_count);
    let uci_move = bn.last_move?.to_string();
    Some((bn.board_state.clone(), uci_move))
}

pub fn print_tree_arena(nodes: &[Node], root: usize, depth: usize) {
    let node = &nodes[root];

    println!(
        "{}- Value: {:.2}, Visits: {}, Prior: {:.2}, Side: {:?}",
        "  ".repeat(depth),
        node.value_sum,
        node.visit_count,
        node.prior,
        node.board_state.board.side_to_move()
    );

    for &child_idx in &node.children {
        print_tree_arena(nodes, child_idx, depth + 1);
    }
}

fn material_of(piece: Piece) -> f32 {
    match piece {
        Queen => 9.0,
        Rook => 5.0,
        Bishop => 3.0,
        Knight => 3.0,
        Pawn => 1.0,
        _ => 0.0,
    }
}

const MATERIAL_VALUE_DIVISOR: f32 = 40.0;

pub fn heuristic_eval(board: &Board, perspective: Color) -> f32 {
    let mut value = 0.0;

    // material
    for color in ALL_COLORS {
        for piece in ALL_PIECES {
            let bitboard = board.color_combined(color).0 & board.pieces(piece).0;
            let total_val =
                (bitboard.count_ones() as f32 * material_of(piece)) / MATERIAL_VALUE_DIVISOR;

            if color == perspective {
                value += total_val;
            } else {
                value -= total_val;
            }
        }
    }

    value
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BoardStateStatus {
    Ongoing,
    Stalemate,
    WhiteWinner,
    BlackWinner,
    Threefold,
    FiftyMove,
}

#[derive(Clone, Debug)]
pub struct BoardState {
    pub board: Board,
    pub halfmove_clock: u8,
    pub repetition_table: HashMap<u64, u8>,
    pub status: BoardStateStatus,
}

impl BoardState {
    pub fn new(board: Board, halfmove_clock: u8, repetition_table: HashMap<u64, u8>) -> BoardState {
        BoardState {
            board,
            halfmove_clock,
            repetition_table: repetition_table,
            status: Ongoing,
        }
    }

    pub fn apply_move(&mut self, mv: ChessMove) {
        self.halfmove_clock += 1;
        if self.board.piece_on(mv.get_source()) == Some(Pawn)
            || self.board.piece_on(mv.get_dest()).is_some()
        {
            self.halfmove_clock = 0;
        }
        self.board = self.board.make_move_new(mv);

        let board_hash = self.board.get_hash();

        let current_rep = self.repetition_table.get(&board_hash);
        let mut new_rep = 1;
        if current_rep.is_some() {
            new_rep = current_rep.unwrap() + 1;
        }
        self.repetition_table.insert(board_hash, new_rep);

        if self.board.status() == Checkmate {
            if self.board.side_to_move() == White {
                // white's move after black played the mating move
                self.status = BoardStateStatus::BlackWinner;
            } else {
                self.status = BoardStateStatus::WhiteWinner;
            }
        } else if self.board.status() == Stalemate {
            self.status = BoardStateStatus::Stalemate;
        }

        if self.halfmove_clock >= 100 {
            self.status = FiftyMove;
        }

        if new_rep >= 3 {
            self.status = Threefold;
        }
    }
}

impl Default for BoardState {
    fn default() -> BoardState {
        BoardState::new(Board::default(), 0, HashMap::new())
    }
}