neural_engine/src/polyglot.rs

use std::fs::File;
use std::io::{self, Read};
use rand::Rng;

#[derive(Debug, Clone, Copy)]
pub struct PolyglotEntry {
    pub key: u64,
    pub mov: u16,
    pub weight: u16,
    pub learn: u32,
}

pub struct PolyglotBook {
    entries: Vec<PolyglotEntry>,
}

impl PolyglotBook {
    pub fn new(path: &str) -> io::Result<Self> {
        let mut file = File::open(path)?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer)?;
        
        let mut entries = Vec::with_capacity(buffer.len() / 16);
        for chunk in buffer.chunks_exact(16) {
            let key = u64::from_be_bytes(chunk[0..8].try_into().unwrap());
            let mov = u16::from_be_bytes(chunk[8..10].try_into().unwrap());
            let weight = u16::from_be_bytes(chunk[10..12].try_into().unwrap());
            let learn = u32::from_be_bytes(chunk[12..16].try_into().unwrap());
            
            entries.push(PolyglotEntry { key, mov, weight, learn });
        }
        
        Ok(PolyglotBook { entries })
    }
    
    pub fn lookup(&self, key: u64) -> Option<String> {
        // Binary search the entries for the first match
        let start_idx = self.entries.binary_search_by_key(&key, |e| e.key).ok()?;
        
        // The binary search might hit any entry in a block of identical keys.
        // Find the start of the block.
        let mut first = start_idx;
        while first > 0 && self.entries[first - 1].key == key {
            first -= 1;
        }
        
        // Find all moves for this key
        let mut valid_moves = Vec::new();
        let mut i = first;
        while i < self.entries.len() && self.entries[i].key == key {
            if self.entries[i].weight > 0 {
                valid_moves.push(self.entries[i]);
            }
            i += 1;
        }
        
        if valid_moves.is_empty() {
            return None;
        }
        
        // Weighted random choice
        let total_weight: u32 = valid_moves.iter().map(|e| e.weight as u32).sum();
        if total_weight == 0 {
            // Fallback to random if all weights are somehow 0 (already handled above but just in case)
            let chosen = valid_moves[rand::thread_rng().gen_range(0..valid_moves.len())];
            return Some(polyglot_to_uci(chosen.mov));
        }
        
        let mut rng = rand::thread_rng();
        let mut random_weight = rng.gen_range(0..total_weight);
        
        for entry in valid_moves {
            if random_weight < entry.weight as u32 {
                return Some(polyglot_to_uci(entry.mov));
            }
            random_weight -= entry.weight as u32;
        }
        
        None
    }
}

pub fn polyglot_to_uci(mov: u16) -> String {
    let to_file = mov & 7;
    let to_rank = (mov >> 3) & 7;
    let from_file = (mov >> 6) & 7;
    let from_rank = (mov >> 9) & 7;
    let promotion = (mov >> 12) & 7;
    
    let from_sq = format!("{}{}", (b'a' + from_file as u8) as char, (b'1' + from_rank as u8) as char);
    let to_sq = format!("{}{}", (b'a' + to_file as u8) as char, (b'1' + to_rank as u8) as char);
    
    let prom_str = match promotion {
        1 => "n",
        2 => "b",
        3 => "r",
        4 => "q",
        _ => "",
    };
    
    // Polyglot encodes castling as King captures Rook.
    // e1h1 becomes O-O, which in standard UCI is e1g1.
    // Let's normalize it to standard UCI.
    let is_castling = match (from_sq.as_str(), to_sq.as_str()) {
        ("e1", "h1") => Some("e1g1"),
        ("e1", "a1") => Some("e1c1"),
        ("e8", "h8") => Some("e8g8"),
        ("e8", "a8") => Some("e8c8"),
        _ => None,
    };
    
    if let Some(castling_move) = is_castling {
        castling_move.to_string()
    } else {
        format!("{}{}{}", from_sq, to_sq, prom_str)
    }
}