model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from .attn_map import apm_map, apm_out
import math
from .encoding_simple import encode_fen_to_tensor, encode_moves_to_tensor
from .vocab import policy_index
from typing import Union, List, Optional
import bulletchess
import numpy as np


from transformers import PretrainedConfig, PreTrainedModel

class Gating(nn.Module):
    def __init__(self, features_shape, additive=True, init_value=None):
        super(Gating, self).__init__()
        self.additive = additive
        if init_value is None:
            init_value = 0 if self.additive else 1
        
        self.gate = nn.Parameter(torch.full(features_shape, float(init_value)))
        if not self.additive:
            self.gate.register_hook(lambda grad: torch.clamp(grad, min=0))

    def forward(self, x):
        if self.additive:
            return x + self.gate
        else:
            return x * self.gate

def ma_gating(x, in_features):
    x = Gating(in_features, additive=False)(x)
    x = Gating(in_features, additive=True)(x)
    return x

class RMSNorm(nn.Module):
    def __init__(self, in_features, scale=True):
        super(RMSNorm, self).__init__()
        self.scale = scale
        if self.scale:
            self.gamma = nn.Parameter(torch.ones(in_features))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-5)
        x_normalized = x / rms
        if self.scale:
            return x_normalized * self.gamma
        return x_normalized

class ApplyAttentionPolicyMap(nn.Module):
    def __init__(self):
        super(ApplyAttentionPolicyMap, self).__init__()
        # Register as buffers so they move with the model when .to(device) is called
        # Use same names as before for backward compatibility with saved models
        self.register_buffer('fc1', torch.from_numpy(apm_map).float())
        self.register_buffer('idx', torch.from_numpy(apm_out).long())

    def forward(self, logits, pp_logits):
        logits = torch.cat([logits.reshape(-1, 64 * 64),
                            pp_logits.reshape(-1, 8 * 24)],
                           dim=1)
        
        batch_size = logits.size(0)
        idx = self.idx.unsqueeze(0).expand(batch_size, -1)
        
        return torch.gather(logits, 1, idx)

class Mish(nn.Module):
    def __init__(self):
        super(Mish, self).__init__()

    def forward(self, x):
        return x * torch.tanh(F.softplus(x))

class CustomMHA(nn.Module):
    def __init__(self, emb_size, d_model, num_heads, dropout=0.0, use_bias_qkv=True, use_bias_out=True,
                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
        super(CustomMHA, self).__init__()
        assert d_model % num_heads == 0
        self.emb_size = emb_size
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.wq = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
        self.wk = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
        self.wv = nn.Linear(emb_size, d_model, bias=use_bias_qkv)
        self.out_proj = nn.Linear(d_model, emb_size, bias=use_bias_out)
        self.attn_dropout = nn.Dropout(dropout)
        # Optional Smolgen components
        self.smol_compress = None
        self.smol_hidden1 = None
        self.smol_hidden1_ln = None
        self.smol_gen_from = None
        self.smol_gen_from_ln = None
        self.smol_weight_gen = None
        if use_smolgen:
            self.smol_compress = nn.Linear(emb_size, smol_hidden_channels, bias=False)
            self.smol_hidden1 = nn.Linear(64 * smol_hidden_channels, smol_hidden_sz, bias=True)
            self.smol_hidden1_ln = nn.LayerNorm(smol_hidden_sz, eps=1e-3)
            self.smol_gen_from = nn.Linear(smol_hidden_sz, num_heads * smol_gen_sz, bias=True)
            self.smol_gen_from_ln = nn.LayerNorm(num_heads * smol_gen_sz, eps=1e-3)
            self.smol_weight_gen = nn.Linear(smol_gen_sz, 64 * 64, bias=False)
        self.smol_activation = smol_activation

    def _shape(self, x):
        b, l, _ = x.shape
        return x.view(b, l, self.num_heads, self.head_dim).transpose(1, 2)

    def forward(self, x, return_attn=False):
        # x: (B, L, emb_size)
        q = self.wq(x)
        k = self.wk(x)
        v = self.wv(x)
        q = self._shape(q)  # (B, H, L, D)
        k = self._shape(k)
        v = self._shape(v)
        scale = torch.sqrt(torch.tensor(self.head_dim, dtype=x.dtype, device=x.device))
        attn_logits = torch.matmul(q, k.transpose(-2, -1)) / scale
        # Add Smolgen weights if present
        smol_w = None
        if self.smol_compress is not None:
            b, l, _ = x.shape
            compressed = self.smol_compress(x)  # (B, L, hidden_channels)
            compressed = compressed.reshape(b, l * compressed.shape[-1])  # (B, 64*hidden_channels)
            hidden_pre = self.smol_hidden1(compressed)
            hidden = F.silu(hidden_pre) if self.smol_activation == 'swish' else F.silu(hidden_pre)
            hidden_ln = self.smol_hidden1_ln(hidden)
            gen_from_pre = self.smol_gen_from(hidden_ln)
            gen_from_act = F.silu(gen_from_pre) if self.smol_activation == 'swish' else F.silu(gen_from_pre)
            gen_from = self.smol_gen_from_ln(gen_from_act)
            gen_from = gen_from.view(b, self.num_heads, -1)  # (B, H, gen_sz)
            smol_w = self.smol_weight_gen(gen_from)  # (B, H, 64*64)
            smol_w = smol_w.view(b, self.num_heads, l, l)
            attn_logits = attn_logits + smol_w
        # Numerically stable softmax matching TF (float32, subtract max)
        attn_logits = attn_logits - attn_logits.max(dim=-1, keepdim=True)[0]
        attn_weights = torch.exp(attn_logits)
        attn_weights = attn_weights / attn_weights.sum(dim=-1, keepdim=True)
        attn_weights = self.attn_dropout(attn_weights)
        attn_output = torch.matmul(attn_weights, v)  # (B, H, L, D)
        attn_output = attn_output.transpose(1, 2).contiguous().view(x.size(0), x.size(1), self.d_model)
        out = self.out_proj(attn_output)
        if return_attn:
            return out, attn_weights, smol_w, attn_logits
        return out

class FFN(nn.Module):
    def __init__(self, emb_size, dff, activation=Mish(), omit_other_biases=False):
        super(FFN, self).__init__()
        self.dense1 = nn.Linear(emb_size, dff, bias=not omit_other_biases)
        self.activation = activation
        self.dense2 = nn.Linear(dff, emb_size, bias=not omit_other_biases)

    def forward(self, x):
        x = self.dense1(x)
        x = self.activation(x)
        x = self.dense2(x)
        return x

class EncoderLayer(nn.Module):
    def __init__(self, emb_size, d_model, num_heads, dff, dropout_rate, encoder_layers, skip_first_ln=False, encoder_rms_norm=False, omit_qkv_biases=False, omit_other_biases=False,
                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
        super(EncoderLayer, self).__init__()
        self.mha = CustomMHA(emb_size, d_model, num_heads, dropout=dropout_rate, use_bias_qkv=not omit_qkv_biases, use_bias_out=not omit_other_biases,
                             use_smolgen=use_smolgen, smol_hidden_channels=smol_hidden_channels, smol_hidden_sz=smol_hidden_sz, smol_gen_sz=smol_gen_sz, smol_activation=smol_activation)
        self.ffn = FFN(emb_size, dff, omit_other_biases=omit_other_biases)
        
        self.norm1 = RMSNorm(emb_size) if encoder_rms_norm else nn.LayerNorm(emb_size, eps=0.001)
        self.norm2 = RMSNorm(emb_size) if encoder_rms_norm else nn.LayerNorm(emb_size, eps=0.001)
        
        self.dropout1 = nn.Dropout(dropout_rate)
        self.dropout2 = nn.Dropout(dropout_rate)
        
        self.alpha = (2. * encoder_layers)**-0.25
        self.skip_first_ln = skip_first_ln

    def forward(self, x):
        attn_output = self.mha(x)
        attn_output = self.dropout1(attn_output)
        
        out1 = x + attn_output * self.alpha
        if not self.skip_first_ln:
            out1 = self.norm1(out1)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output)
        
        out2 = self.norm2(out1 + ffn_output * self.alpha)
        return out2

class PolicyHead(nn.Module):
    def __init__(self, pol_embedding_size, policy_d_model, opponent=False):
        super(PolicyHead, self).__init__()
        self.opponent = opponent
        self.wq = nn.Linear(pol_embedding_size, policy_d_model)
        self.wk = nn.Linear(pol_embedding_size, policy_d_model)
        self.ppo = nn.Linear(policy_d_model, 4, bias=False)
        self.apply_map = ApplyAttentionPolicyMap()

    def forward(self, x):
        if self.opponent:
            x = torch.flip(x, [1])

        queries = self.wq(x)
        keys = self.wk(x)

        matmul_qk = torch.matmul(queries, keys.transpose(-2, -1))
        
        dk = torch.sqrt(torch.tensor(keys.shape[-1], dtype=keys.dtype, device=keys.device))
        promotion_keys = keys[:, -8:, :]
        promotion_offsets = self.ppo(promotion_keys).transpose(-2,-1) * dk
        promotion_offsets = promotion_offsets[:, :3, :] + promotion_offsets[:, 3:4, :]

        n_promo_logits = matmul_qk[:, -16:-8, -8:]
        q_promo_logits = (n_promo_logits + promotion_offsets[:, 0:1, :]).unsqueeze(3)
        r_promo_logits = (n_promo_logits + promotion_offsets[:, 1:2, :]).unsqueeze(3)
        b_promo_logits = (n_promo_logits + promotion_offsets[:, 2:3, :]).unsqueeze(3)
        promotion_logits = torch.cat([q_promo_logits, r_promo_logits, b_promo_logits], axis=3).reshape(-1, 8, 24)

        promotion_logits = promotion_logits / dk
        policy_attn_logits = matmul_qk / dk

        return self.apply_map(policy_attn_logits, promotion_logits)

class ValueHead(nn.Module):
    def __init__(self, embedding_size, val_embedding_size, default_activation=Mish()):
        super(ValueHead, self).__init__()
        self.embedding = nn.Linear(embedding_size, val_embedding_size)
        self.activation = default_activation
        self.flatten = nn.Flatten()
        self.dense1 = nn.Linear(val_embedding_size * 64, 128)
        self.dense2 = nn.Linear(128, 3)

    def forward(self, x):
        x = self.embedding(x)
        x = self.activation(x)
        x = self.flatten(x)
        x = self.dense1(x)
        x = self.activation(x)
        x = self.dense2(x)
        return x

class BT4Config(PretrainedConfig):
    """Configuration class for BT4 model."""
    model_type = "bt4"
    
    def __init__(
        self,
        embedding_size=1024,
        embedding_dense_sz=512,
        encoder_layers=15,
        encoder_d_model=1024,
        encoder_heads=32,
        encoder_dff=1536,
        dropout_rate=0.0,
        pol_embedding_size=1024,
        policy_d_model=1024,
        val_embedding_size=128,
        use_smolgen=True,
        smol_hidden_channels=32,
        smol_hidden_sz=256,
        smol_gen_sz=256,
        smol_activation="swish",
        **kwargs
    ):
        super().__init__(**kwargs)
        self.embedding_size = embedding_size
        self.embedding_dense_sz = embedding_dense_sz
        self.encoder_layers = encoder_layers
        self.encoder_d_model = encoder_d_model
        self.encoder_heads = encoder_heads
        self.encoder_dff = encoder_dff
        self.dropout_rate = dropout_rate
        self.pol_embedding_size = pol_embedding_size
        self.policy_d_model = policy_d_model
        self.val_embedding_size = val_embedding_size
        self.use_smolgen = use_smolgen
        self.smol_hidden_channels = smol_hidden_channels
        self.smol_hidden_sz = smol_hidden_sz
        self.smol_gen_sz = smol_gen_sz
        self.smol_activation = smol_activation

class BT4(PreTrainedModel):
    config_class = BT4Config
    
    def __init__(self, config=None, embedding_size=1024, embedding_dense_sz=512, encoder_layers=15, encoder_d_model=1024, encoder_heads=32, encoder_dff=1536, dropout_rate=0.0, pol_embedding_size=1024, policy_d_model=1024, val_embedding_size=128, default_activation=Mish(),
                 use_smolgen=True, smol_hidden_channels=32, smol_hidden_sz=256, smol_gen_sz=256, smol_activation='swish'):
        # Initialize PreTrainedModel with config
        if config is None:
            config = BT4Config(
                embedding_size=embedding_size,
                embedding_dense_sz=embedding_dense_sz,
                encoder_layers=encoder_layers,
                encoder_d_model=encoder_d_model,
                encoder_heads=encoder_heads,
                encoder_dff=encoder_dff,
                dropout_rate=dropout_rate,
                pol_embedding_size=pol_embedding_size,
                policy_d_model=policy_d_model,
                val_embedding_size=val_embedding_size,
                use_smolgen=use_smolgen,
                smol_hidden_channels=smol_hidden_channels,
                smol_hidden_sz=smol_hidden_sz,
                smol_gen_sz=smol_gen_sz,
                smol_activation=smol_activation,
            )
        super(BT4, self).__init__(config)
        
        # Use config values (config is now guaranteed to exist)
        embedding_size = config.embedding_size
        embedding_dense_sz = config.embedding_dense_sz
        encoder_layers = config.encoder_layers
        encoder_d_model = config.encoder_d_model
        encoder_heads = config.encoder_heads
        encoder_dff = config.encoder_dff
        dropout_rate = config.dropout_rate
        pol_embedding_size = config.pol_embedding_size
        policy_d_model = config.policy_d_model
        val_embedding_size = config.val_embedding_size
        use_smolgen = config.use_smolgen
        smol_hidden_channels = config.smol_hidden_channels
        smol_hidden_sz = config.smol_hidden_sz
        smol_gen_sz = config.smol_gen_sz
        smol_activation = config.smol_activation
        self.embedding_dense_sz = embedding_dense_sz
        # DeepNorm alpha used in embedding residual; default uses provided encoder_layers
        self.deepnorm_alpha = (2. * encoder_layers) ** -0.25
        
        self.embedding_preprocess = nn.Linear(64*12, 64*self.embedding_dense_sz)
        self.embedding = nn.Linear(112 + self.embedding_dense_sz, embedding_size)
        nn.init.xavier_uniform_(self.embedding.weight) # Explicitly set initializer
        nn.init.zeros_(self.embedding.bias)

        self.embedding_ln = nn.LayerNorm(embedding_size, eps=0.001)
        
        self.gating_mult = Gating((64, embedding_size), additive=False)
        self.gating_add = Gating((64, embedding_size), additive=True)

        self.embedding_ffn = FFN(embedding_size, encoder_dff)
        self.embedding_ffn_ln = nn.LayerNorm(embedding_size, eps=0.001)
        
        self.encoder_layers_list = nn.ModuleList([
            EncoderLayer(embedding_size, encoder_d_model, encoder_heads, encoder_dff, dropout_rate, encoder_layers,
                         use_smolgen=use_smolgen, smol_hidden_channels=smol_hidden_channels, smol_hidden_sz=smol_hidden_sz, smol_gen_sz=smol_gen_sz, smol_activation=smol_activation)
            for _ in range(encoder_layers)
        ])
        
        self.policy_embedding = nn.Linear(embedding_size, pol_embedding_size)
        self.policy_head = PolicyHead(pol_embedding_size, policy_d_model)
        self.value_head_winner = ValueHead(embedding_size, val_embedding_size)
        self.value_head_q = ValueHead(embedding_size, val_embedding_size)
        self.activation = default_activation
        
        self.apply(self._init_weights)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        """Load model from pretrained checkpoint (required by transformers)."""
        from transformers import AutoConfig
        from huggingface_hub import hf_hub_download
        from safetensors.torch import load_file
        import os
        
        # Load config
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
        
        # Create model with config
        model = cls(config=config)
        
        # Check if it's a HuggingFace Hub path or local path
        is_hf_hub = "/" in pretrained_model_name_or_path and not os.path.isdir(pretrained_model_name_or_path)
        
        if is_hf_hub:
            # Download from HuggingFace Hub
            safetensors_path = hf_hub_download(
                repo_id=pretrained_model_name_or_path,
                filename="model.safetensors",
                cache_dir=kwargs.get("cache_dir", None),
                token=kwargs.get("token", None),
            )
            state_dict = load_file(safetensors_path)
            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
            if missing_keys:
                print(f"Warning: Missing keys when loading weights: {len(missing_keys)} keys")
            if unexpected_keys:
                print(f"Warning: Unexpected keys when loading weights: {len(unexpected_keys)} keys")
        else:
            # Local path - try safetensors first
            safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors")
            if os.path.exists(safetensors_path):
                state_dict = load_file(safetensors_path)
                missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
                if missing_keys:
                    print(f"Warning: Missing keys when loading weights: {len(missing_keys)} keys")
                if unexpected_keys:
                    print(f"Warning: Unexpected keys when loading weights: {len(unexpected_keys)} keys")
            else:
                # Fall back to pytorch format
                pt_path = os.path.join(pretrained_model_name_or_path, "model.pt")
                checkpoint = torch.load(pt_path, map_location="cpu")
                if isinstance(checkpoint, dict):
                    if "state_dict" in checkpoint:
                        state_dict = checkpoint["state_dict"]
                    elif "model" in checkpoint:
                        state_dict = checkpoint["model"]
                    else:
                        state_dict = checkpoint
                else:
                    state_dict = checkpoint
                missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
                if missing_keys:
                    print(f"Warning: Missing keys when loading weights: {len(missing_keys)} keys")
                if unexpected_keys:
                    print(f"Warning: Unexpected keys when loading weights: {len(unexpected_keys)} keys")
        
        return model
    
    @classmethod
    def register_for_auto_class(cls, auto_class):
        """Register this class for auto class loading (required by transformers)."""
        # This is a no-op for custom models with trust_remote_code=True
        pass

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            # Keras' glorot_normal is equivalent to PyTorch's xavier_normal_
            nn.init.xavier_normal_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, x):
        # x shape: (batch, 112, 8, 8)
        flow = x.permute(0, 2, 3, 1).reshape(-1, 64, 112)
        
        pos_info = flow[..., :12]
        pos_info_flat = pos_info.reshape(-1, 64 * 12)
        
        pos_info_processed = self.embedding_preprocess(pos_info_flat)
        pos_info = pos_info_processed.reshape(-1, 64, self.embedding_dense_sz)
        
        flow = torch.cat([flow, pos_info], dim=-1)
        
        flow = self.embedding(flow)
        
        flow = self.activation(flow)
        
        flow = self.embedding_ln(flow)

        flow = self.gating_mult(flow)
        flow = self.gating_add(flow)
        
        ffn_dense1_pre = self.embedding_ffn.dense1(flow)
        ffn_dense1 = self.embedding_ffn.activation(ffn_dense1_pre)
        ffn_output = self.embedding_ffn.dense2(ffn_dense1)
        
        residual = flow + ffn_output * self.deepnorm_alpha
        flow = self.embedding_ffn_ln(residual)
        
        for i, layer in enumerate(self.encoder_layers_list):
            flow = layer(flow)
        
        policy_tokens = self.policy_embedding(flow)
        policy_tokens = self.activation(policy_tokens)
        
        policy_logits = self.policy_head(policy_tokens)
        
        value_winner = self.value_head_winner(flow)
        value_q = self.value_head_q(flow)

        return policy_logits, value_winner, value_q
    
    def get_move_from_history(self, fen_or_moves: Union[str, List[str]], T: float, device: str = None, **kwargs) -> str:
        """
        Predict a move from a move history or FEN position.
        
        Args:
            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
            device: Device to run the model on (if None, uses model's device)
            return_probs: If True, returns a dictionary of move probabilities instead of a single move
        
        Returns:
            UCI move string (e.g., 'e2e4') or dictionary of move probabilities if return_probs=True
        """
        # Detect device from model if not provided
        if device is None:
            device = next(self.parameters()).device
        else:
            device = torch.device(device)
        
        # Determine if input is FEN string or list of moves
        if isinstance(fen_or_moves, str):
            # FEN string input
            fen = fen_or_moves
            is_black_to_move = fen.split()[1] == 'b'
            input_tensor_112, legal_moves_mask = encode_fen_to_tensor(fen)
            castling_rights = fen.split()[2] if len(fen.split()) > 2 else ""
        elif isinstance(fen_or_moves, list):
            # List of UCI moves input
            move_history = fen_or_moves
            input_tensor_112, legal_moves_mask = encode_moves_to_tensor(move_history)
            # Create board to check if black is to move and for castling rights
            board = bulletchess.Board()
            for mv in move_history:
                move = bulletchess.Move.from_uci(mv)
                board.apply(move)
            is_black_to_move = (board.turn == bulletchess.BLACK)
            fen_parts = board.fen().split()
            castling_rights = fen_parts[2] if len(fen_parts) > 2 else ""
        else:
            raise ValueError("Input must be a FEN string or a list of UCI moves")
        
        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
        
        self.eval()
        with torch.inference_mode():
            policy_logits,_,_ = self.forward(input_tensor_112)
        
        # Apply legal moves mask without in-place ops (inference tensor)
        logits0 = policy_logits[0] + torch.from_numpy(legal_moves_mask).to(policy_logits.device)
        
        # Check if return_probs is requested
        return_probs = kwargs.get('return_probs', False)
        
        if return_probs:
            # Return probabilities dictionary
            scaled_logits = logits0 / T if T > 0 else logits0
            probs = F.softmax(scaled_logits, dim=0)
            probs_dict = {}
            for idx, move in enumerate(policy_index):
                prob_val = probs[idx].item()
                if prob_val > 1e-6:  # Only include moves with non-negligible probability
                    probs_dict[move] = prob_val
            return probs_dict
        
        if T == 0.0:
            # Deterministic: return best move
            best_move_idx = torch.argmax(logits0).item()
            uci_move = policy_index[best_move_idx]
        else:
            # Stochastic sampling with temperature
            # Apply temperature scaling
            scaled_logits = logits0 / T
            # Apply softmax to get probabilities
            probs = F.softmax(scaled_logits, dim=0)
            # Sample from the distribution
            move_idx = torch.multinomial(probs, 1).item()
            uci_move = policy_index[move_idx]
        
        # If black is to move, the board was mirrored during encoding, so we need to mirror the move back
        # Mirror ranks: 1↔8, 2↔7, 3↔6, 4↔5 (keep file letters the same)
        if is_black_to_move:
            def mirror_rank(rank_char):
                rank = int(rank_char)
                return str(9 - rank)
            
            # UCI format: e2e4, e7e8q, etc.
            if len(uci_move) >= 4:
                from_file = uci_move[0]
                from_rank = uci_move[1]
                to_file = uci_move[2]
                to_rank = uci_move[3]
                promo = uci_move[4:] if len(uci_move) > 4 else ""
                
                uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
        
        # Convert castling moves from king-to-rook-square format to standard castling format
        # Only if castling rights are available (check FEN castling rights)
        # Check and convert white castling moves
        if uci_move == "e1h1" and "K" in castling_rights:
            uci_move = "e1g1"
        elif uci_move == "e1a1" and "Q" in castling_rights:
            uci_move = "e1c1"
        # Check and convert black castling moves
        elif uci_move == "e8h8" and "k" in castling_rights:
            uci_move = "e8g8"
        elif uci_move == "e8a8" and "q" in castling_rights:
            uci_move = "e8c8"
        
        return uci_move
    
    def get_best_move_value(self, fen_or_moves: Union[str, List[str]], T: float = 0.0, device: str = None) -> tuple:
        """
        Get the best move and its value using value analysis.
        
        Args:
            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
            device: Device to run the model on (if None, uses model's device)
        
        Returns:
            Tuple of (best_move, value) where value is the position evaluation
        """
        # Detect device from model if not provided
        if device is None:
            device = next(self.parameters()).device
        else:
            device = torch.device(device)
        
        # Determine if input is FEN string or list of moves
        if isinstance(fen_or_moves, str):
            fen = fen_or_moves
            is_black_to_move = fen.split()[1] == 'b'
            input_tensor_112, legal_moves_mask = encode_fen_to_tensor(fen)
            castling_rights = fen.split()[2] if len(fen.split()) > 2 else ""
        elif isinstance(fen_or_moves, list):
            move_history = fen_or_moves
            input_tensor_112, legal_moves_mask = encode_moves_to_tensor(move_history)
            board = bulletchess.Board()
            for mv in move_history:
                move = bulletchess.Move.from_uci(mv)
                board.apply(move)
            is_black_to_move = (board.turn == bulletchess.BLACK)
            fen_parts = board.fen().split()
            castling_rights = fen_parts[2] if len(fen_parts) > 2 else ""
        else:
            raise ValueError("Input must be a FEN string or a list of UCI moves")
        
        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
        
        self.eval()
        with torch.inference_mode():
            policy_logits, _, value_q = self.forward(input_tensor_112)
        
        # Apply legal moves mask
        logits0 = policy_logits[0] + torch.from_numpy(legal_moves_mask).to(policy_logits.device)
        
        # Get best move
        if T == 0.0:
            best_move_idx = torch.argmax(logits0).item()
        else:
            scaled_logits = logits0 / T
            probs = F.softmax(scaled_logits, dim=0)
            move_idx = torch.multinomial(probs, 1).item()
            best_move_idx = move_idx
        
        uci_move = policy_index[best_move_idx]
        
        # Mirror move if black is to move
        if is_black_to_move:
            def mirror_rank(rank_char):
                rank = int(rank_char)
                return str(9 - rank)
            
            if len(uci_move) >= 4:
                from_file = uci_move[0]
                from_rank = uci_move[1]
                to_file = uci_move[2]
                to_rank = uci_move[3]
                promo = uci_move[4:] if len(uci_move) > 4 else ""
                uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
        
        # Convert castling moves
        if uci_move == "e1h1" and "K" in castling_rights:
            uci_move = "e1g1"
        elif uci_move == "e1a1" and "Q" in castling_rights:
            uci_move = "e1c1"
        elif uci_move == "e8h8" and "k" in castling_rights:
            uci_move = "e8g8"
        elif uci_move == "e8a8" and "q" in castling_rights:
            uci_move = "e8c8"
        
        # Get value (softmax over value_q)
        value_probs = F.softmax(value_q[0], dim=0)
        value = value_probs.cpu().numpy()
        
        return uci_move, value
    
    def get_position_value(self, fen_or_moves: Union[str, List[str]], device: str = None) -> np.ndarray:
        """
        Get position evaluation using value_q.
        
        Args:
            fen_or_moves: Either a FEN string representing the chess position, or a list of UCI moves
            device: Device to run the model on (if None, uses model's device)
        
        Returns:
            Array of [black_win, draw, white_win] probabilities
        """
        # Detect device from model if not provided
        if device is None:
            device = next(self.parameters()).device
        else:
            device = torch.device(device)
        
        # Determine if input is FEN string or list of moves
        if isinstance(fen_or_moves, str):
            input_tensor_112, _ = encode_fen_to_tensor(fen_or_moves)
        elif isinstance(fen_or_moves, list):
            input_tensor_112, _ = encode_moves_to_tensor(fen_or_moves)
        else:
            raise ValueError("Input must be a FEN string or a list of UCI moves")
        
        input_tensor_112 = input_tensor_112.to(device, non_blocking=True)
        
        self.eval()
        with torch.inference_mode():
            _, _, value_q = self.forward(input_tensor_112)
        
        # Apply softmax to get probabilities [black_win, draw, white_win]
        value_probs = F.softmax(value_q[0], dim=0)
        return value_probs.cpu().numpy()
    
    def batch_get_moves_from_fens(self, fens: List[str], T: float, device: str = None, use_fp16: bool = False) -> List[str]:
        """
        Get moves for multiple FEN positions using batched inference.
        
        Args:
            fens: List of FEN strings representing chess positions
            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
            device: Device to run the model on (if None, uses model's device)
        
        Returns:
            List of UCI move strings
        """
        if not fens:
            return []
        
        # Detect device from model if not provided
        if device is None:
            device = next(self.parameters()).device
        else:
            device = torch.device(device)
        
        batch_size = len(fens)
        
        # Batch encode all FENs
        input_tensors = []
        legal_moves_masks = []
        is_black_to_move_list = []
        castling_rights_list = []
        
        for fen in fens:
            input_tensor, legal_mask = encode_fen_to_tensor(fen)
            input_tensors.append(input_tensor.squeeze(0))  # Remove batch dim
            legal_moves_masks.append(legal_mask)
            is_black_to_move_list.append(fen.split()[1] == 'b')
            castling_rights_list.append(fen.split()[2] if len(fen.split()) > 2 else "")
        
        # Stack into batch tensor: (batch_size, 112, 8, 8)
        batch_tensor = torch.stack(input_tensors).to(device, non_blocking=True)
        if use_fp16 and device.type == 'cuda':
            batch_tensor = batch_tensor.half()
        
        # Run batched inference
        self.eval()
        with torch.inference_mode():
            if use_fp16 and device.type == 'cuda':
                with torch.autocast(device_type='cuda', dtype=torch.float16):
                    policy_logits,_,_ = self.forward(batch_tensor)
            else:
                policy_logits,_,_ = self.forward(batch_tensor)
        
        # Process each position in the batch
        moves = []
        for i in range(batch_size):
            # Apply legal moves mask
            logits = policy_logits[i] + torch.from_numpy(legal_moves_masks[i]).to(policy_logits.device, dtype=policy_logits.dtype)
            
            # Sample move
            if T == 0.0:
                best_move_idx = torch.argmax(logits).item()
                uci_move = policy_index[best_move_idx]
            else:
                scaled_logits = logits / T
                probs = F.softmax(scaled_logits, dim=0)
                move_idx = torch.multinomial(probs, 1).item()
                uci_move = policy_index[move_idx]
            
            # Mirror move if black is to move
            if is_black_to_move_list[i]:
                def mirror_rank(rank_char):
                    rank = int(rank_char)
                    return str(9 - rank)
                
                if len(uci_move) >= 4:
                    from_file = uci_move[0]
                    from_rank = uci_move[1]
                    to_file = uci_move[2]
                    to_rank = uci_move[3]
                    promo = uci_move[4:] if len(uci_move) > 4 else ""
                    
                    uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
            
            # Convert castling moves
            castling_rights = castling_rights_list[i]
            if uci_move == "e1h1" and "K" in castling_rights:
                uci_move = "e1g1"
            elif uci_move == "e1a1" and "Q" in castling_rights:
                uci_move = "e1c1"
            elif uci_move == "e8h8" and "k" in castling_rights:
                uci_move = "e8g8"
            elif uci_move == "e8a8" and "q" in castling_rights:
                uci_move = "e8c8"
            
            moves.append(uci_move)
        
        return moves
    
    def batch_get_moves_from_move_lists(self, move_lists: List[List[str]], T: float, device: str = None, use_fp16: bool = False, fens: Optional[List[str]] = None):
        """
        Get moves for multiple move histories using batched inference.
        
        Args:
            move_lists: List of move sequences, where each sequence is a list of UCI moves
            T: Temperature for sampling (0.0 = deterministic/argmax, >0.0 = stochastic)
            device: Device to run the model on (if None, uses model's device)
            fens: Optional list of FEN strings that represent the board state prior to 
                  applying the corresponding move list. When provided, each move history
                  is applied starting from the supplied FEN instead of the standard initial position.
        
        Returns:
            List of UCI move strings
        """
        if not move_lists:
            return []
        
        # Detect device from model if not provided
        if device is None:
            device = next(self.parameters()).device
        else:
            device = torch.device(device)
        
        batch_size = len(move_lists)
        
        if fens is not None and len(fens) != len(move_lists):
            raise ValueError("Length of fens must match length of move_lists when provided.")
        
        # Batch encode all move histories
        input_tensors = []
        legal_moves_masks = []
        is_black_to_move_list = []
        castling_rights_list = []
        
        for idx, move_history in enumerate(move_lists):
            starting_fen = fens[idx] if fens is not None else None
            input_tensor, legal_mask = encode_moves_to_tensor(move_history, starting_fen=starting_fen)
            input_tensors.append(input_tensor.squeeze(0))  # Remove batch dim
            legal_moves_masks.append(legal_mask)
            
            board = bulletchess.Board.from_fen(starting_fen) if starting_fen is not None else bulletchess.Board()
            for mv in move_history:
                move = bulletchess.Move.from_uci(mv)
                board.apply(move)
            is_black_to_move_list.append(board.turn == bulletchess.BLACK)
            fen_parts = board.fen().split()
            castling_rights_list.append(fen_parts[2] if len(fen_parts) > 2 else "")
        
        # Stack into batch tensor: (batch_size, 112, 8, 8)
        batch_tensor = torch.stack(input_tensors).to(device, non_blocking=True)
        if use_fp16 and device.type == 'cuda':
            batch_tensor = batch_tensor.half()
        
        # Run batched inference
        self.eval()
        with torch.inference_mode():
            if use_fp16 and device.type == 'cuda':
                with torch.autocast(device_type='cuda', dtype=torch.float16):
                    policy_logits,_,_ = self.forward(batch_tensor)
            else:
                policy_logits,_,_ = self.forward(batch_tensor)
        
        # Process each position in the batch
        moves = []
        for i in range(batch_size):
            # Apply legal moves mask
            logits = policy_logits[i] + torch.from_numpy(legal_moves_masks[i]).to(policy_logits.device, dtype=policy_logits.dtype)
            
            # Sample move
            if T == 0.0:
                best_move_idx = torch.argmax(logits).item()
                uci_move = policy_index[best_move_idx]
            else:
                scaled_logits = logits / T
                probs = F.softmax(scaled_logits, dim=0)
                move_idx = torch.multinomial(probs, 1).item()
                uci_move = policy_index[move_idx]
            
            # Mirror move if black is to move
            if is_black_to_move_list[i]:
                def mirror_rank(rank_char):
                    rank = int(rank_char)
                    return str(9 - rank)
                
                if len(uci_move) >= 4:
                    from_file = uci_move[0]
                    from_rank = uci_move[1]
                    to_file = uci_move[2]
                    to_rank = uci_move[3]
                    promo = uci_move[4:] if len(uci_move) > 4 else ""
                    
                    uci_move = from_file + mirror_rank(from_rank) + to_file + mirror_rank(to_rank) + promo
            
            # Convert castling moves
            castling_rights = castling_rights_list[i]
            if uci_move == "e1h1" and "K" in castling_rights:
                uci_move = "e1g1"
            elif uci_move == "e1a1" and "Q" in castling_rights:
                uci_move = "e1c1"
            elif uci_move == "e8h8" and "k" in castling_rights:
                uci_move = "e8g8"
            elif uci_move == "e8a8" and "q" in castling_rights:
                uci_move = "e8c8"
            
            moves.append(uci_move)
        return moves