use std::io::{self, BufRead}; use std::collections::HashMap; use std::fs; use std::path::Path; use shakmaty::{Chess, Position, CastlingMode}; use shakmaty::zobrist::{Zobrist64, ZobristHash}; use ort::session::Session; use ort::session::builder::GraphOptimizationLevel; use rand_distr::{Dirichlet, Distribution}; use shakmaty_syzygy::{Tablebase, Wdl, AmbiguousWdl}; mod polyglot; use polyglot::PolyglotBook; #[derive(Clone)] struct MCTSNode { visit_count: u32, value_sum: f32, prior: f32, children: HashMap, } impl MCTSNode { fn new(prior: f32) -> Self { Self { visit_count: 0, value_sum: 0.0, prior, children: HashMap::new(), } } } fn extract_pv(node: &MCTSNode) -> Vec { let mut pv = Vec::new(); let mut current = node; while !current.children.is_empty() { let mut best_action = String::new(); let mut best_visits = -1; for (action, child) in ¤t.children { if (child.visit_count as i32) > best_visits { best_visits = child.visit_count as i32; best_action = action.clone(); } } if best_action.is_empty() || best_visits == 0 { break; } pv.push(best_action.clone()); current = current.children.get(&best_action).unwrap(); } pv } fn add_dirichlet_noise(node: &mut MCTSNode, alpha: f32, epsilon: f32) { if node.children.len() <= 1 { return; } let dirichlet = Dirichlet::new(&vec![alpha; node.children.len()]).unwrap(); let noise = dirichlet.sample(&mut rand::thread_rng()); for (i, child) in node.children.values_mut().enumerate() { child.prior = (1.0 - epsilon) * child.prior + epsilon * noise[i]; } } fn evaluate_material(board: &Chess, color: shakmaty::Color) -> i32 { let b = board.board(); let c_mask = b.by_color(color); let pawns = (c_mask & b.by_role(shakmaty::Role::Pawn)).count(); let knights = (c_mask & b.by_role(shakmaty::Role::Knight)).count(); let bishops = (c_mask & b.by_role(shakmaty::Role::Bishop)).count(); let rooks = (c_mask & b.by_role(shakmaty::Role::Rook)).count(); let queens = (c_mask & b.by_role(shakmaty::Role::Queen)).count(); pawns as i32 + knights as i32 * 3 + bishops as i32 * 3 + rooks as i32 * 5 + queens as i32 * 9 } fn material_difference(board: &Chess, color: shakmaty::Color) -> i32 { let other = match color { shakmaty::Color::White => shakmaty::Color::Black, shakmaty::Color::Black => shakmaty::Color::White, }; evaluate_material(board, color) - evaluate_material(board, other) } fn get_piece_value(role: shakmaty::Role) -> i32 { match role { shakmaty::Role::Pawn => 1, shakmaty::Role::Knight => 3, shakmaty::Role::Bishop => 3, shakmaty::Role::Rook => 5, shakmaty::Role::Queen => 9, shakmaty::Role::King => 100, } } fn score_move(m: &shakmaty::Move) -> i32 { if m.is_capture() { let victim = m.capture().map(|r| get_piece_value(r)).unwrap_or(0); let attacker = get_piece_value(m.role()); 1000 + victim * 10 - attacker } else { 0 } } fn quiescence_search(board: &Chess, mut alpha: i32, beta: i32, depth: u8) -> i32 { let in_check = board.is_check(); let stand_pat = material_difference(board, board.turn()); if depth >= 4 || board.is_game_over() { return stand_pat; } if !in_check { if stand_pat >= beta { return beta; } if alpha < stand_pat { alpha = stand_pat; } } let mut legals: Vec<_> = board.legal_moves().into_iter().collect(); legals.sort_by_cached_key(|m| -score_move(m)); for m in &legals { if m.is_capture() || in_check { let mut next_board = board.clone(); next_board.play_unchecked(m); let score = -quiescence_search(&next_board, -beta, -alpha, depth + 1); if score >= beta { return beta; } if score > alpha { alpha = score; } } } alpha } fn tactical_search(board: &Chess, mut alpha: i32, beta: i32, depth: u8, qs_depth: u8) -> i32 { if depth == 0 { return quiescence_search(board, alpha, beta, qs_depth); } if board.is_game_over() { if board.is_checkmate() { return -9999; } return 0; } let mut legals: Vec<_> = board.legal_moves().into_iter().collect(); legals.sort_by_cached_key(|m| -score_move(m)); let mut best_score = -99999; let mut first = true; for m in &legals { let mut next_board = board.clone(); next_board.play_unchecked(m); let score = if first { -tactical_search(&next_board, -beta, -alpha, depth - 1, qs_depth) } else { let mut val = -tactical_search(&next_board, -alpha - 1, -alpha, depth - 1, qs_depth); if val > alpha && val < beta { val = -tactical_search(&next_board, -beta, -alpha, depth - 1, qs_depth); } val }; first = false; if score > best_score { best_score = score; } if best_score > alpha { alpha = best_score; } if alpha >= beta { break; } } best_score } fn load_vocab() -> (HashMap, HashMap) { let data = fs::read_to_string("vocab.json").expect("Unable to read vocab.json. Please run build_engine.bat first to generate it!"); let vocab: HashMap = serde_json::from_str(&data).expect("JSON format error"); let mut inv_vocab = HashMap::new(); for (k, v) in &vocab { inv_vocab.insert(*v, k.clone()); } (vocab, inv_vocab) } fn evaluate_onnx( session: &mut Session, history: &[String], vocab: &HashMap, inv_vocab: &HashMap, ) -> (HashMap, f32) { let max_length = 120; let mut ids = Vec::new(); for t in history { ids.push(*vocab.get(t).unwrap_or(&0)); } // 120-move sliding window if ids.len() > max_length { ids = ids[ids.len() - max_length..].to_vec(); } let seq_len = ids.len(); let input_value = ort::value::Tensor::from_array((vec![1, seq_len], ids)).unwrap(); // Disable logging for inputs! macro as it's safe here let inputs = ort::inputs!["input_ids" => input_value]; let outputs = session.run(inputs).unwrap(); let (_, policy_data) = outputs["policy_logits"].try_extract_tensor::().unwrap(); let (_, value_data) = outputs["value_preds"].try_extract_tensor::().unwrap(); let value = value_data[seq_len - 1]; let vocab_size = policy_data.len() / seq_len; let start_idx = (seq_len - 1) * vocab_size; let end_idx = start_idx + vocab_size; let logits = &policy_data[start_idx..end_idx]; // Softmax over policy logits with Temperature let temp = 1.0; let max_logit = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max) / temp; let mut exps = Vec::with_capacity(logits.len()); let mut sum_exp = 0.0; for l in logits.iter() { let e = (*l / temp - max_logit).exp(); exps.push(e); sum_exp += e; } let mut policy_map = HashMap::new(); for (i, p) in exps.iter().enumerate() { if let Some(uci) = inv_vocab.get(&(i as i64)) { policy_map.insert(uci.clone(), *p / sum_exp); } } (policy_map, value) } fn endgame_heuristic(board: &Chess) -> f32 { let mut score = 0.0; let turn = board.turn(); // 1. King Centralization if let Some(my_king) = board.board().king_of(turn) { let file = my_king.file() as i8; let rank = my_king.rank() as i8; let dist_f = (file - 3).abs().min((file - 4).abs()); let dist_r = (rank - 3).abs().min((rank - 4).abs()); let manhattan = dist_f + dist_r; score += (6.0 - manhattan as f32) * 0.05; } // 2. Passed Pawns (advanced pawns bonus) let my_pawns = board.board().pawns() & board.board().by_color(turn); for sq in my_pawns { let r = sq.rank() as i8; let relative_rank = if turn == shakmaty::Color::White { r } else { 7 - r }; if relative_rank >= 4 { // 5th rank or higher score += (relative_rank as f32 - 3.0) * 0.05; } } score } fn mcts_simulate( node: &mut MCTSNode, board: Chess, history: &mut Vec, history_hashes: &mut Vec, session: &mut Session, vocab: &HashMap, inv_vocab: &HashMap, nn_cache: &mut HashMap, f32)>, tablebases: &mut Tablebase, ) -> f32 { if board.is_game_over() { if board.is_checkmate() { return -1.0; // The player whose turn it is got checkmated } return 0.0; // Draw } // Draw Avoidance by 3-fold Repetition let current_hash = board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0; let mut hash_count = 0; for &h in history_hashes.iter() { if h == current_hash { hash_count += 1; } } // If the hash appears at least twice in the history path before this node, it's a 3-fold draw! if hash_count >= 2 { return 0.0; } // Probe Syzygy Endgame Tablebases for positions with <= 7 pieces if board.board().occupied().count() <= 7 { if let Ok(wdl) = tablebases.probe_wdl(&board) { let tb_value = match wdl { AmbiguousWdl::Win | AmbiguousWdl::CursedWin | AmbiguousWdl::MaybeWin => 0.99, AmbiguousWdl::Loss | AmbiguousWdl::BlessedLoss | AmbiguousWdl::MaybeLoss => -0.99, AmbiguousWdl::Draw => 0.0, }; return tb_value; } } if node.children.is_empty() { // Expand let legals = board.legal_moves(); // 1-ply tactical checkmate bypass! for m in &legals { let uci = m.to_uci(CastlingMode::Standard).to_string(); let mut next_board = board.clone(); next_board.play_unchecked(m); if next_board.is_checkmate() { node.children.insert(uci, MCTSNode::new(1.0)); node.visit_count += 1; node.value_sum += 1.0; return -0.99; // 0.99 Impatience Penalty } } // Evaluate with ONNX or Cache let cache_key = board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0; let (policy, mut value) = if let Some(cached) = nn_cache.get(&cache_key) { cached.clone() } else { let res = evaluate_onnx(session, history, vocab, inv_vocab); nn_cache.insert(cache_key, res.clone()); res }; // Quiescence Search Tactical Override let current_mat = material_difference(&board, board.turn()); let qs_val = quiescence_search(&board, -999, 999, 0); let delta = qs_val - current_mat; // If the player to move can force a massive material gain of 3+ points via captures, override! if delta >= 3 { value = 0.95; // They are winning } else if delta <= -3 { value = -0.95; // They are losing } // Phase 3: Endgame Knowledge Heuristics (8-10 pieces) if board.board().occupied().count() <= 10 { value += endgame_heuristic(&board); value = value.clamp(-0.99, 0.99); } for m in &legals { let uci = m.to_uci(CastlingMode::Standard).to_string(); let p = policy.get(&uci).copied().unwrap_or(0.0); node.children.insert(uci, MCTSNode::new(p)); } node.visit_count += 1; node.value_sum += value; return -value * 0.99; } // Select best child let mut best_score = f32::NEG_INFINITY; let mut best_action = String::new(); let total_sqrt = (node.visit_count as f32).sqrt(); let c_puct = 2.5; for (action, child) in &node.children { let q = if child.visit_count > 0 { child.value_sum / (child.visit_count as f32) } else { 0.0 }; // PUCT formula let u = c_puct * child.prior * total_sqrt / (1.0 + child.visit_count as f32); let score = q + u; if score > best_score { best_score = score; best_action = action.clone(); } } // Play move let m = shakmaty::uci::Uci::from_ascii(best_action.as_bytes()) .unwrap() .to_move(&board) .unwrap(); let mut next_board = board.clone(); next_board.play_unchecked(&m); history.push(best_action.clone()); let next_hash = next_board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0; history_hashes.push(next_hash); let child_node = node.children.get_mut(&best_action).unwrap(); let val = mcts_simulate(child_node, next_board, history, history_hashes, session, vocab, inv_vocab, nn_cache, tablebases); history.pop(); history_hashes.pop(); node.visit_count += 1; node.value_sum += val; // 0.99 Impatience Penalty during backprop -val * 0.99 } fn main() { let stdin = io::stdin(); let mut board = Chess::default(); let mut history_moves = vec!["".to_string()]; let mut history_hashes = Vec::new(); // Ensure ONNX library is set up implicitly by ort let _ = ort::init().with_name("neural_engine").commit(); let mut session = if Path::new("model_int8.onnx").exists() { Some(Session::builder().unwrap() .with_optimization_level(GraphOptimizationLevel::Level3).unwrap() .with_intra_threads(4).unwrap() .commit_from_file("model_int8.onnx").unwrap()) } else { None }; let mut tablebases = Tablebase::::new(); let _ = tablebases.add_directory("syzygy"); // Load polyglot opening book if available let polyglot_book = PolyglotBook::new("book.bin").ok(); let (vocab, inv_vocab) = load_vocab(); let mut nn_cache: HashMap, f32)> = HashMap::new(); let mut global_root = MCTSNode::new(1.0); for line in stdin.lock().lines() { let line = line.expect("Failed to read line"); let tokens: Vec<&str> = line.trim().split_whitespace().collect(); if tokens.is_empty() { continue; } match tokens[0] { "uci" => { println!("id name Neurex"); println!("id author Neural Engine Architect"); println!("uciok"); } "isready" => { println!("readyok"); } "position" => { // position startpos moves e2e4 e7e5 board = Chess::default(); let old_history = history_moves.clone(); history_moves.clear(); history_moves.push("".to_string()); history_hashes.clear(); history_hashes.push(board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0); if tokens.contains(&"moves") { let moves_idx = tokens.iter().position(|&r| r == "moves").unwrap(); for m_str in &tokens[moves_idx + 1..] { if let Ok(uci_move) = shakmaty::uci::Uci::from_ascii(m_str.as_bytes()) { if let Ok(m) = uci_move.to_move(&board) { board.play_unchecked(&m); history_moves.push(m_str.to_string()); history_hashes.push(board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0); } } } } // Persist the tree if the history matches let mut matches = true; if old_history.len() <= history_moves.len() { for i in 0..old_history.len() { if old_history[i] != history_moves[i] { matches = false; break; } } } else { matches = false; } if matches { for i in old_history.len()..history_moves.len() { let m = &history_moves[i]; if let Some(child) = global_root.children.remove(m) { global_root = child; } else { global_root = MCTSNode::new(1.0); break; } } } else { global_root = MCTSNode::new(1.0); } } "go" => { let mut movetime_ms = 5000; let mut wtime = 0; let mut btime = 0; let mut winc = 0; let mut binc = 0; if tokens.contains(&"wtime") { if let Some(idx) = tokens.iter().position(|&r| r == "wtime") { if let Ok(t) = tokens[idx + 1].parse::() { wtime = t; } } } if tokens.contains(&"btime") { if let Some(idx) = tokens.iter().position(|&r| r == "btime") { if let Ok(t) = tokens[idx + 1].parse::() { btime = t; } } } if tokens.contains(&"winc") { if let Some(idx) = tokens.iter().position(|&r| r == "winc") { if let Ok(t) = tokens[idx + 1].parse::() { winc = t; } } } if tokens.contains(&"binc") { if let Some(idx) = tokens.iter().position(|&r| r == "binc") { if let Ok(t) = tokens[idx + 1].parse::() { binc = t; } } } if tokens.contains(&"movetime") { let idx = tokens.iter().position(|&r| r == "movetime").unwrap(); if let Ok(t) = tokens[idx + 1].parse::() { movetime_ms = t; } } else if wtime > 0 || btime > 0 { // Calculate time dynamically // history_moves has "" + moves. // If length is odd (1, 3, 5), it's White's turn. let is_white = history_moves.len() % 2 != 0; if is_white && wtime > 0 { movetime_ms = (wtime as f64 / 40.0 + winc as f64 * 0.8) as u128; } else if !is_white && btime > 0 { movetime_ms = (btime as f64 / 40.0 + binc as f64 * 0.8) as u128; } } // Cap time between 0.5s and 25.0s to avoid HF timeout or instant flag if movetime_ms > 25000 { movetime_ms = 25000; } else if movetime_ms < 100 { movetime_ms = 100; } // --- 1. Polyglot Opening Book Check --- if let Some(book) = &polyglot_book { let hash = board.zobrist_hash::(shakmaty::EnPassantMode::Legal).0; if let Some(book_move) = book.lookup(hash) { println!("bestmove {}", book_move); continue; // Skip MCTS completely } } if let Some(ref mut sess) = session { let start_time = std::time::Instant::now(); // Leave a small buffer (5% or minimum 50ms) to ensure we reply in time let buffer = (movetime_ms as f64 * 0.05) as u128; let mut safe_movetime = if movetime_ms > buffer { movetime_ms - buffer } else { movetime_ms }; if global_root.children.is_empty() { mcts_simulate(&mut global_root, board.clone(), &mut history_moves, &mut history_hashes, sess, &vocab, &inv_vocab, &mut nn_cache, &mut tablebases); } add_dirichlet_noise(&mut global_root, 0.3, 0.25); let mut iter_count = 0; // Loop continuously until time expires let mut extended = false; while (start_time.elapsed().as_millis() as u128) < safe_movetime { mcts_simulate(&mut global_root, board.clone(), &mut history_moves, &mut history_hashes, sess, &vocab, &inv_vocab, &mut nn_cache, &mut tablebases); iter_count += 1; // Stream UCI info and check Dynamic Time Allocation every 100 iterations if iter_count % 100 == 0 { let mut best_eval = 0.0; let mut best_visits = -1; let mut second_best_visits = -1; let mut best_action = String::new(); for (action, child) in &global_root.children { let v = child.visit_count as i32; if v > best_visits { second_best_visits = best_visits; best_visits = v; best_action = action.clone(); if child.visit_count > 0 { best_eval = child.value_sum / (child.visit_count as f32); } } else if v > second_best_visits { second_best_visits = v; } } let time_ms = start_time.elapsed().as_millis(); let cp = (best_eval * 1000.0) as i32; // --- 2. Dynamic Time Management --- if time_ms as f64 > (safe_movetime as f64 * 0.15) && best_visits > 500 { // Early Stopping: Move is overwhelmingly obvious if second_best_visits > 0 && best_visits > 10 * second_best_visits { println!("info string Early stopping: move is obvious."); break; } } if !extended && time_ms as f64 > (safe_movetime as f64 * 0.8) { // Time Extension: Position is critical and unsure if second_best_visits > 0 && best_visits < (second_best_visits as f64 * 1.2) as i32 { let extra_time = (safe_movetime as f64 * 0.5) as u128; let max_allowed = wtime / 10; // Avoid flagging if wtime == 0 || (safe_movetime + extra_time < max_allowed) { safe_movetime += extra_time; println!("info string Extending search time for critical position. New limit: {}ms", safe_movetime); extended = true; } } } let pv_moves = extract_pv(&global_root); let pv_str = pv_moves.join(" "); let depth = std::cmp::max(1, pv_moves.len()); let nps = if time_ms > 0 { (iter_count as u128 * 1000) / time_ms } else { 0 }; println!("info depth {} nodes {} nps {} score cp {} time {} pv {}", depth, iter_count, nps, cp, time_ms, pv_str); } } // TACTICAL SAFETY NET let mut candidates: Vec<(&String, &MCTSNode)> = global_root.children.iter().collect(); // Sort candidates by visit_count descending candidates.sort_by(|a, b| b.1.visit_count.cmp(&a.1.visit_count)); let current_mat = material_difference(&board, board.turn()); let mut best_action = String::new(); let mut best_eval = 0.0; for (action, child) in candidates { let m_res = shakmaty::uci::Uci::from_ascii(action.as_bytes()).unwrap().to_move(&board); if let Ok(m) = m_res { let mut next_board = board.clone(); next_board.play_unchecked(&m); // 3 plies of full search, then 4 plies of quiescence (starting from depth 0) let tact_score = -tactical_search(&next_board, -9999, 9999, 3, 0); // If it doesn't unconditionally drop >= 2 points of material, ACCEPT IT if tact_score > current_mat - 2 { best_action = action.clone(); if child.visit_count > 0 { best_eval = child.value_sum / (child.visit_count as f32); } break; } } } // Fallback to highest visited move if ALL top moves drop material (e.g. inevitable capture) if best_action.is_empty() { let mut best_visits = -1; for (action, child) in &global_root.children { if (child.visit_count as i32) > best_visits { best_visits = child.visit_count as i32; best_action = action.clone(); if child.visit_count > 0 { best_eval = child.value_sum / (child.visit_count as f32); } } } } if !best_action.is_empty() { let cp = (best_eval * 1000.0) as i32; let time_ms = start_time.elapsed().as_millis(); let pv_moves = extract_pv(&global_root); let pv_str = pv_moves.join(" "); let depth = std::cmp::max(1, pv_moves.len()); let nps = if time_ms > 0 { (iter_count as u128 * 1000) / time_ms } else { 0 }; println!("info depth {} nodes {} nps {} score cp {} time {} pv {}", depth, iter_count, nps, cp, time_ms, pv_str); println!("bestmove {}", best_action); } else { println!("bestmove 0000"); // fallback } } else { println!("bestmove 0000"); } } "quit" => break, _ => {} } } }