chess_mcvs.py

# chess_mcvs.py
"""
chess_mcvs.py - Full implementation analogous to breakthrough_mcvs.py

Features:
- Chess game logic using python-chess (full rules)
- Fixed 64�64 padded weighted adjacency matrix (Hilbert-ordered pieces)
- Piece-specific kappa values
- Queen-only promotions for flat 4096 policy head compatibility
- Policy and Value CNNs (lightweight with pooling for feasibility)
- HilbertOrderedZoneDatabase with add methods and proper k-NN zone score
- MCVSSearcher (guided ?-PUCT with policy/value/zone)
- UCTSearcher (baseline)
- train_networks function
- Ready for chess_mcvs_vs_uct.py tournament

Requires: pip install chess torch numpy
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import chess
import random
import time
import os
import bisect
from collections import Counter
from typing import List, Dict, Tuple

# ============================================================================
# HILBERT CURVE UTILITIES
# ============================================================================

def xy2d(n: int, x: int, y: int) -> int:
    d = 0
    s = n // 2
    while s > 0:
        rx = 1 if (x & s) > 0 else 0
        ry = 1 if (y & s) > 0 else 0
        d += s * s * ((3 * rx) ^ ry)
        if ry == 0:
            if rx == 1:
                x = n - 1 - x
                y = n - 1 - y
            x, y = y, x
        s //= 2
    return d

def matrix_to_hilbert_index(W: np.ndarray) -> int:
    if W.size == 0:
        return 0
    W_flat = W.flatten()
    max_idx = np.argmax(W_flat)
    x = max_idx % 64
    y = max_idx // 64
    return xy2d(64, x, y)

# ============================================================================
# ABC MODEL & WEIGHTED MATRIX
# ============================================================================

class ABCModelDynamic:
    """ABC Model with displacement-based B matrices."""
    
    def __init__(self, n: int = 2, t: float = 1.0, T: float = 1.41):
        """
        Initialize for 2D board.
        
        Args:
            n: Dimension (2 for 2D board)
            t: Minimum distance threshold for adjacency
            T: Maximum distance threshold for adjacency
        """
        self.n = n
        self.t = t
        self.T = T
        self.c0 = np.array([0.0, 0.0, 1.0])
        self.B_blocks = []        # Displacement matrices
        self.piece_positions = [] # Actual positions on SB
        self.a_product = np.eye(3) # Accumulated product
        self.delta = {}           # Delta values for each piece
        self.kappa = {}           # Tokenized vectors for each piece (ADDED)
        self.stage = 0
        
        # MB position tracking
        self.MB_current = self.c0.copy()  # Current MB center
        self.MB_previous = self.c0.copy() # Previous MB center
        self.history = []
    
    def create_B_displacement(self, dx: float, dy: float) -> np.ndarray:
        """
        Create displacement matrix B.
        
        Moves MB by (dx, dy) displacement.
        
        B � current_MB = new_MB
        
        B = [1 0 dx]
            [0 1 dy]
            [0 0 1]
        
        Right column: (dx, dy, 1)?
        """
        B = np.array([
            [1.0, 0.0, dx],
            [0.0, 1.0, dy],
            [0.0, 0.0, 1.0]
        ])
        return B
    
    def add_piece(self, position: tuple, delta_values: tuple, kappa_vector: np.ndarray = None):
        """
        Add a piece at position with delta values and optional kappa vector.
        
        Args:
            position: (x, y) tuple
            delta_values: (??, ??, ??) tuple
            kappa_vector: Tokenized feature vector (optional, defaults to delta_values)
        """
        x, y = position
        
        # Current MB position
        current_MB = self.MB_current.copy()
        
        # Calculate displacement
        dx = x - current_MB[0]
        dy = y - current_MB[1]
        
        # Create displacement matrix
        B_i = self.create_B_displacement(dx, dy)
        
        # Store
        idx = len(self.B_blocks)
        self.B_blocks.append(B_i)
        self.piece_positions.append((x, y))
        
        # Update accumulated product: a_t = a_{t-1} � B_t
        self.a_product = self.a_product @ B_i
        
        # Track MB positions
        self.MB_previous = self.MB_current.copy()
        self.MB_current = self.a_product @ self.c0  # Should equal (x, y, 1)
        
        # Store differential vector
        self.delta[idx] = delta_values
        
        # Store tokenized vector (or use delta_values as default)
        if kappa_vector is None:
            kappa_vector = np.array(delta_values)
        self.kappa[idx] = kappa_vector
        
        # Increment stage
        self.stage += 1
        
        # Record move
        self.history.append({
            'stage': self.stage,
            'SB_position': (x, y),
            'displacement': (dx, dy),
            'B_i': B_i.copy(),
            'a_t': self.a_product.copy(),
            'MB_prev': self.MB_previous.copy(),
            'MB_curr': self.MB_current.copy(),
            'delta': delta_values,
            'kappa': kappa_vector.copy()
        })
    
    def add_move(self, x: float, y: float, player: int):
        """
        Convenience method: Add a move with automatic delta values.
        
        Args:
            x, y: Position coordinates
            player: 1 for X (??=1.0), 2 for O (??=1.1)
        """
        delta_1 = 1.0  # Occupied
        delta_2 = 1.0 if player == 1 else 1.1  # Player color
        delta_3 = 1.0  # Same value for all pieces
        
        self.add_piece((x, y), (delta_1, delta_2, delta_3))
    
    def get_board_state(self) -> str:
        """Visual board state."""
        pos_map = {
            (-1, 1): 0, (0, 1): 1, (1, 1): 2,
            (-1, 0): 3, (0, 0): 4, (1, 0): 5,
            (-1, -1): 6, (0, -1): 7, (1, -1): 8
        }
        
        board = [' '] * 9
        for i, pos in enumerate(self.piece_positions):
            if pos in pos_map:
                player = 'X' if self.delta[i][1] == 1.0 else 'O'
                board[pos_map[pos]] = player
        
        s = f" {board[0]} | {board[1]} | {board[2]} \n"
        s += "---|---|---\n"
        s += f" {board[3]} | {board[4]} | {board[5]} \n"
        s += "---|---|---\n"
        s += f" {board[6]} | {board[7]} | {board[8]} \n"
        return s
    
    def describe_move(self, stage: int):
        """Describe a specific move."""
        if stage < 1 or stage > len(self.history):
            return "Invalid stage"
        
        move = self.history[stage - 1]
        desc = "=" * 70 + "\n"
        desc += f"MOVE {stage}: {move['SB_position']}\n"
        desc += "=" * 70 + "\n\n"
        
        idx = stage - 1
        B_i = self.B_blocks[idx]
        
        desc += f"Mobile Board (MB) movement:\n"
        desc += f" Previous MB: {move['MB_prev']}\n"
        desc += f" Target (SB): {move['SB_position']}\n"
        desc += f" Displacement: ? = {move['displacement']}\n"
        desc += f" New MB: {move['MB_curr']}\n\n"
        desc += f"Displacement matrix B_{stage}:\n{B_i}\n"
        desc += f"Right column: {B_i[:, 2]} = (?x={move['displacement'][0]}, ?y={move['displacement'][1]}, 1)\n\n"
        
        desc += f"Accumulated product a_{stage}:\n{move['a_t']}\n"
        desc += f"Right column: {move['a_t'][:, 2]}\n\n"
        
        desc += f"Delta values: {move['delta']}\n"
        desc += f"Kappa vector: {move['kappa']}\n\n"
        
        desc += f"Verification:\n"
        desc += f" a_{stage} � c? = {move['a_t'] @ self.c0}\n"
        desc += f" Should equal: {np.array([move['SB_position'][0], move['SB_position'][1], 1.0])}\n\n"
        
        desc += "Board state:\n"
        desc += self.get_board_state()
        
        return desc


class WeightedMatrixABC:
    """
    Compute weighted adjacency matrix from ABCModelDynamic.
    
    CRITICAL: Converts (x,y) tuples to homogeneous coords [x,y,1] 
    to match algebraic structure of B matrices from abc_model.py
    
    Implements Definition from Section 6:
        W[i,j](t) = A[i,j](t) ? S[i,j](t) ? F[i,j](t)
    """
    
    def __init__(self, abc_model: ABCModelDynamic, sigma: float = 1.0):
        """
        Initialize with ABC model instance.
        
        Args:
            abc_model: ABCModelDynamic instance
            sigma: Standard deviation for Gaussian spatial kernel
        """
        self.abc = abc_model
        self.n = abc_model.n
        self.t = abc_model.t
        self.T = abc_model.T
        self.sigma = sigma
        self.positions = None  # Will store homogeneous coords [x, y, 1]
        self.D = None
        self.A = None
        self.S = None
        self.F = None
        self.W = None
    
    def compute_piece_positions(self) -> np.ndarray:
        """
        Extract positions from ABC model and convert to homogeneous coordinates.
        
        In abc_model.py:
            - piece_positions are (x, y) tuples
            - MB_current = a_product @ c0 produces [x, y, 1]
        
        For matrix computations, convert to homogeneous coords to match
        the algebraic structure of B matrices.
        
        Returns: Array of shape (num_pieces, 3) with homogeneous coords [x, y, 1]
        """
        positions = []
        for pos in self.abc.piece_positions:
            # pos is (x, y) from piece_positions
            x, y = pos
            # Convert to homogeneous coordinates to align with B matrix algebra
            positions.append(np.array([x, y, 1.0]))
        
        self.positions = np.array(positions)
        return self.positions
    
    def compute_distance_matrix(self) -> np.ndarray:
        """
        Compute pairwise Euclidean distances using 2D components only.
        
        Even though positions are [x, y, 1], we compute distance using only [x, y]
        D[i,j] = ||[x_i, y_i] - [x_j, y_j]||?
        """
        if self.positions is None:
            self.compute_piece_positions()
        
        num = len(self.positions)
        D = np.zeros((num, num))
        
        for i in range(num):
            for j in range(num):
                # Use only 2D coordinates for distance
                pos_i = self.positions[i][:2]
                pos_j = self.positions[j][:2]
                D[i, j] = np.linalg.norm(pos_i - pos_j)
        
        self.D = D
        return D
    
    def compute_adjacency_matrix(self) -> np.ndarray:
        """
        Compute Adjacency Matrix (Definition 6.2).
        
        A[i,j] = 1 if:
            - k ? D[i,j] ? K (grid-adjacent)
            - i == j AND i is isolated (no neighbors in [k,K])
        
        A[i,j] = 0 otherwise
        """
        if self.D is None:
            self.compute_distance_matrix()
        
        num = len(self.positions)
        A = np.zeros((num, num))
        
        # Identify isolated pieces
        isolated = set()
        for i in range(num):
            has_neighbor = False
            for j in range(num):
                if i != j and self.t <= self.D[i, j] <= self.T:
                    has_neighbor = True
                    break
            if not has_neighbor:
                isolated.add(i)
        
        # Fill adjacency matrix
        for i in range(num):
            for j in range(num):
                # Check adjacency distance
                if i != j and self.t <= self.D[i, j] <= self.T:
                    A[i, j] = 1.0
                # Check if isolated piece
                elif i == j and i in isolated:
                    A[i, j] = 1.0
        
        self.A = A
        return A
    
    def compute_spatial_matrix(self) -> np.ndarray:
        """
        Compute Spatial Matrix (Definition 6.3).
        
        S[i,j](t) = exp(-||c_i(t) - c_j(t)||� / (2?�))
        
        Gaussian kernel based on Euclidean distance.
        """
        if self.D is None:
            self.compute_distance_matrix()
        
        num = len(self.positions)
        S = np.zeros((num, num))
        
        for i in range(num):
            for j in range(num):
                S[i, j] = np.exp(-self.D[i, j]**2 / (2 * self.sigma**2))
        
        self.S = S
        return S
    
    def compute_feature_matrix(self) -> np.ndarray:
        """
        Compute Feature Matrix (Definition 6.4).
        
        F[i,j](t) = <?(B_i), ?(B_j)> / (||?(B_i)|| � ||?(B_j)||)
        
        Cosine similarity between tokenized vectors.
        """
        num = len(self.positions)
        F = np.zeros((num, num))
        
        # Extract tokenized vectors
        kappas = []
        for i in range(num):
            if i in self.abc.kappa:
                kappas.append(self.abc.kappa[i])
            else:
                # Default: use delta as feature vector
                if i in self.abc.delta:
                    kappas.append(np.array(self.abc.delta[i]))
                else:
                    kappas.append(np.array([1.0, 1.0, 1.0]))
        
        # Compute cosine similarity
        for i in range(num):
            for j in range(num):
                kappa_i = kappas[i]
                kappa_j = kappas[j]
                
                dot_product = np.dot(kappa_i, kappa_j)
                norm_i = np.linalg.norm(kappa_i)
                norm_j = np.linalg.norm(kappa_j)
                
                if norm_i > 0 and norm_j > 0:
                    F[i, j] = dot_product / (norm_i * norm_j)
                else:
                    F[i, j] = 0.0
        
        self.F = F
        return F
    
    def compute_weighted_matrix(self) -> np.ndarray:
        """
        Compute Weighted Matrix (Definition from Section 6).
        
        W[i,j](t) = A[i,j](t) ? S[i,j](t) ? F[i,j](t)
        
        Returns: W - the full weighted matrix
        """
        if self.A is None:
            self.compute_adjacency_matrix()
        if self.S is None:
            self.compute_spatial_matrix()
        if self.F is None:
            self.compute_feature_matrix()
        
        # Hadamard product (element-wise multiplication)
        W = self.A * self.S * self.F
        
        self.W = W
        return W
    
    def compute_manhattan_distance(self, other: 'WeightedMatrixABC') -> float:
        """
        Compute Manhattan distance between two weighted matrices.
        
        distance = ||W_1 - W_2||?
        """
        if self.W is None:
            self.compute_weighted_matrix()
        if other.W is None:
            other.compute_weighted_matrix()
        
        W1 = self.W
        W2 = other.W
        
        # Pad to same size
        size = max(len(W1), len(W2))
        W1_padded = np.zeros((size, size))
        W2_padded = np.zeros((size, size))
        W1_padded[:len(W1), :len(W1)] = W1
        W2_padded[:len(W2), :len(W2)] = W2
        
        return np.linalg.norm(W1_padded - W2_padded, ord=1)

# ============================================================================
# CHESS GAME LOGIC
# ============================================================================

class Chess:
    def __init__(self):
        self.board = chess.Board()
        self.move_count = 0
        self.position_history = Counter()
        self._update_position_key()

    def _update_position_key(self):
        key = (self.board.fen(), self.board.turn)
        self.position_history[key] += 1

    def copy(self):
        new_game = Chess()
        new_game.board = self.board.copy()
        new_game.move_count = self.move_count
        new_game.position_history = self.position_history.copy()
        return new_game

    def undo_move(self) -> None:
        """
        Undo the last move (pop) and revert bookkeeping (move_count, position_history).

        This complements `apply_move` which increments move_count and updates
        the position_history for the new state after a push. `undo_move` decrements
        the same position record and pops the board to restore the previous state.
        """
        # The current board fen/turn corresponds to the state that was last pushed;
        # decrement its count before popping to fully revert the update.
        key = (self.board.fen(), self.board.turn)
        if key in self.position_history:
            # decrement safely (don't remove key to avoid KeyError on repeated ops)
            self.position_history[key] = max(0, self.position_history[key] - 1)
        # pop the board and update counters
        self.board.pop()
        self.move_count = max(0, self.move_count - 1)

    def get_legal_moves(self) -> List[chess.Move]:
        moves = []
        for move in self.board.legal_moves:
            if move.promotion is None or move.promotion == chess.QUEEN:
                moves.append(move)
        return moves

    def apply_move(self, move: chess.Move) -> None:
        self.board.push(move)
        self.move_count += 1
        self._update_position_key()

    def is_terminal(self) -> bool:
        return self.board.is_game_over(claim_draw=True)

    def check_winner(self) -> int:
        if not self.is_terminal():
            return 0
        outcome = self.board.outcome(claim_draw=True)
        if outcome is None or outcome.winner is None:
            return 0
        return 1 if outcome.winner == chess.WHITE else -1

    def get_weighted_adjacency_matrix(self) -> np.ndarray:
        occupied = []
        for sq in chess.SQUARES:
            piece = self.board.piece_at(sq)
            if piece:
                file = chess.square_file(sq)
                rank = chess.square_rank(sq)
                x = float(file) - 3.5
                y = float(rank) - 3.5
                color_val = 1.0 if piece.color == chess.WHITE else 1.1
                type_val = {
                    chess.PAWN:   1.0,
                    chess.KNIGHT: 3.0,
                    chess.BISHOP: 3.2,
                    chess.ROOK:   5.0,
                    chess.QUEEN:  9.0,
                    chess.KING:  20.0
                }[piece.piece_type]
                kappa = np.array([1.0, color_val, type_val])
                hilbert_d = xy2d(8, file, rank)
                occupied.append((hilbert_d, (x, y), kappa))

        occupied.sort(key=lambda t: t[0])
    
        abc = ABCModelDynamic(n=2, t=1.0, T=1.41)
        for _, pos, kappa_vec in occupied:
            abc.add_piece(pos, delta_values=(1.0, kappa_vec[1], kappa_vec[2]), kappa_vector=kappa_vec)

        w_calc = WeightedMatrixABC(abc, sigma=1.0)
        W_var = w_calc.compute_weighted_matrix()

        num = len(occupied)
        W = np.zeros((64, 64))
        if num > 0:
            W[:num, :num] = W_var
        return W

# ============================================================================
# NEURAL NETWORKS
# ============================================================================

class PolicyNetworkCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=5, padding=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=5, padding=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=2),
            nn.ReLU(),
        )
        self.fc = nn.Linear(128 * 8 * 8, 4096)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

class ValueNetworkCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=5, padding=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=5, padding=2, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=1, stride=2),
            nn.ReLU(),
        )
        self.fc = nn.Sequential(
            nn.Linear(128 * 8 * 8, 512),
            nn.ReLU(),
            nn.Linear(512, 1),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ============================================================================
# ZONE DATABASE
# ============================================================================

class HilbertOrderedZoneDatabase:
    def __init__(self, filepath: str = "chess_zone_db.npz", max_size: int = 100000):
        self.filepath = filepath
        self.max_size = max_size  # Maximum matrices per zone
        self.winning_matrices = []
        self.winning_indices = []
        self.losing_matrices = []
        self.losing_indices = []
        self.draw_matrices = []
        self.draw_indices = []
        self.load()
    
    def add_winning_matrix(self, W: np.ndarray):
        if W.shape != (64, 64):
            return
        
        hilbert_idx = matrix_to_hilbert_index(W)
        pos = bisect.bisect_left(self.winning_indices, hilbert_idx)
        self.winning_indices.insert(pos, hilbert_idx)
        self.winning_matrices.insert(pos, W.copy())
    
    def add_losing_matrix(self, W: np.ndarray):
        if W.shape != (64, 64):
            return
        
        hilbert_idx = matrix_to_hilbert_index(W)
        pos = bisect.bisect_left(self.losing_indices, hilbert_idx)
        self.losing_indices.insert(pos, hilbert_idx)
        self.losing_matrices.insert(pos, W.copy()) 
    
    def add_draw_matrix(self, W: np.ndarray):
        if W.shape != (64, 64):
            return
        
        hilbert_idx = matrix_to_hilbert_index(W)
        pos = bisect.bisect_left(self.draw_indices, hilbert_idx)
        self.draw_indices.insert(pos, hilbert_idx)
        self.draw_matrices.insert(pos, W.copy())
    
    def add_game_record(self, trajectory: List[Chess], result: int, sample_rate: float = 0.3):
        """Add game record with sampling to control database size."""
        # Sample states: always keep first and last, random subset of middle
        sampled_states = []
        
        for i, state in enumerate(trajectory):
            if i == 0 or i == len(trajectory) - 1:
                # Always keep start and end positions
                sampled_states.append(state)
            elif random.random() < sample_rate:
                # Randomly sample middle positions
                sampled_states.append(state)
        
        for state in sampled_states:
            W = state.get_weighted_adjacency_matrix()
            
            if result == 1:
                self.add_winning_matrix(W)
            elif result == -1:
                self.add_losing_matrix(W)
            else:
                self.add_draw_matrix(W)
    
    def save(self):
        """Save database with pruning if too large."""
        # Prune if any zone exceeds 80% of max_size
        threshold = int(self.max_size * 0.8)
        
        if (len(self.winning_matrices) > threshold or 
            len(self.losing_matrices) > threshold or 
            len(self.draw_matrices) > threshold):
            print(f"Pruning database before save...")
            self.prune_database(target_size=int(self.max_size * 0.7))
        
        try:
            np.savez_compressed(
                self.filepath,
                winning=np.array(self.winning_matrices, dtype=object),
                winning_indices=np.array(self.winning_indices),
                losing=np.array(self.losing_matrices, dtype=object),
                losing_indices=np.array(self.losing_indices),
                draw=np.array(self.draw_matrices, dtype=object),
                draw_indices=np.array(self.draw_indices),
            )
            print(f"Zone DB saved: W={len(self.winning_matrices)}, L={len(self.losing_matrices)}, D={len(self.draw_matrices)}")
        except MemoryError:
            print(f"MemoryError during save! Pruning more aggressively...")
            self.prune_database(target_size=1000)
            # Try again with much smaller size
            np.savez_compressed(
                self.filepath,
                winning=np.array(self.winning_matrices, dtype=object),
                winning_indices=np.array(self.winning_indices),
                losing=np.array(self.losing_matrices, dtype=object),
                losing_indices=np.array(self.losing_indices),
                draw=np.array(self.draw_matrices, dtype=object),
                draw_indices=np.array(self.draw_indices),
            )
    
    def load(self):
        if os.path.exists(self.filepath):
            try:
                data = np.load(self.filepath, allow_pickle=True)
                self.winning_matrices = list(data['winning_matrices'])
                self.winning_indices = list(data['winning_indices'])
                self.losing_matrices = list(data['losing_matrices'])
                self.losing_indices = list(data['losing_indices'])
                self.draw_matrices = list(data['draw_matrices']) 
                self.draw_indices = list(data['draw_indices'])
                print(f"Loaded zone DB: W={len(self.winning_matrices)}, L={len(self.losing_matrices)}, D={len(self.draw_matrices)}")
            except Exception as e:
                print(f"Failed to load zone database: {e}")
    
    def prune_database(self, target_size: int = 5000):
        """Keep only evenly-spaced diverse samples along Hilbert curve."""
        def prune_zone(matrices, indices, target):
            if len(matrices) <= target:
                return matrices, indices
            
            # Keep evenly spaced samples
            step = max(1, len(matrices) // target)
            keep_idx = list(range(0, len(matrices), step))[:target]
            
            new_matrices = [matrices[i] for i in keep_idx]
            new_indices = [indices[i] for i in keep_idx]
            
            return new_matrices, new_indices
        
        old_sizes = (len(self.winning_matrices), len(self.losing_matrices), len(self.draw_matrices))
        
        self.winning_matrices, self.winning_indices = prune_zone(
            self.winning_matrices, self.winning_indices, target_size
        )
        self.losing_matrices, self.losing_indices = prune_zone(
            self.losing_matrices, self.losing_indices, target_size
        )
        self.draw_matrices, self.draw_indices = prune_zone(
            self.draw_matrices, self.draw_indices, target_size
        )
        
        new_sizes = (len(self.winning_matrices), len(self.losing_matrices), len(self.draw_matrices))
        print(f"Pruned: W {old_sizes}?{new_sizes}, L {old_sizes}?{new_sizes}, D {old_sizes}?{new_sizes}")
    
    def compute_zone_score(self, W: np.ndarray, k: int = 5, beta: float = 0.5) -> float:
        """Compute zone guidance score Z(x(t), a) ? [-1, 1]."""
        if W.shape != (64, 64):
            return 0.0
        
        def knn_similarity(matrices, k_val):
            if len(matrices) == 0:
                return 0.0
            
            k_actual = min(k_val, len(matrices))
            distances = []
            
            for mat in matrices[:k_actual]:
                dist = np.sum(np.abs(W - mat)) / (64.0 * 64.0)
                distances.append(dist)
            
            similarities = [1.0 - d for d in distances]
            return np.mean(similarities)
        
        zone_win = knn_similarity(self.winning_matrices, k)
        zone_loss = knn_similarity(self.losing_matrices, k)
        zone_draw = knn_similarity(self.draw_matrices, k)
        
        Z = zone_win - zone_loss + beta * zone_draw
        return float(np.clip(Z, -1.0, 1.0))

# ============================================================================
# TRAINING FUNCTION
# ============================================================================

def train_networks(policy_net, value_net, W_states, policies, values, epochs=5, batch_size=32, device='cpu', lr=0.001):
    policy_net.train()
    value_net.train()
    opt_p = optim.Adam(policy_net.parameters(), lr=lr)
    opt_v = optim.Adam(value_net.parameters(), lr=lr)

    W_tensor = torch.tensor(np.array(W_states)[:, np.newaxis, :, :], dtype=torch.float32)
    pi_tensor = torch.tensor(np.array(policies), dtype=torch.float32)
    v_tensor = torch.tensor(np.array(values), dtype=torch.float32).unsqueeze(1)

    dataset = TensorDataset(W_tensor, pi_tensor, v_tensor)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    for epoch in range(epochs):
        total_loss = 0.0
        for batch_W, batch_pi, batch_v in loader:
            batch_W = batch_W.to(device)
            batch_pi = batch_pi.to(device)
            batch_v = batch_v.to(device)

            logits = policy_net(batch_W)
            # Fixed policy loss: KL divergence between network policy and MCTS visit distribution
            policy_loss = F.kl_div(F.log_softmax(logits, dim=1), batch_pi, reduction='batchmean')

            pred_v = value_net(batch_W)
            value_loss = F.mse_loss(pred_v, batch_v)

            loss = policy_loss + value_loss
            total_loss += loss.item()

            opt_p.zero_grad()
            opt_v.zero_grad()
            loss.backward()

            # Optional but recommended: gradient clipping
            torch.nn.utils.clip_grad_norm_(policy_net.parameters(), max_norm=1.0)
            torch.nn.utils.clip_grad_norm_(value_net.parameters(), max_norm=1.0)

            opt_p.step()
            opt_v.step()

        # Optional: print epoch loss for monitoring
        print(f"Epoch {epoch+1}/{epochs} - Loss: {total_loss/len(loader):.4f}")

# ============================================================================
# MCVS SEARCHER
# ============================================================================

class MCVSSearcher:
    def __init__(self, policy_net, value_net, zone_db, device='cpu',
                 cpuct=1.41, lambda_zone=0.0, k_zone=5, use_nets=True,
                 dirichlet_alpha: float = 0.3, dirichlet_noise_fraction: float = 0.25):
        self.policy_net = policy_net.to(device) if policy_net and use_nets else None
        self.value_net = value_net.to(device) if value_net and use_nets else None
        self.zone_db = zone_db
        self.device = device
        self.cpuct = cpuct
        self.lambda_zone = lambda_zone
        self.k_zone = k_zone
        self.use_nets = use_nets
        # Dirichlet noise hyperparameters (alpha for the Dirichlet, and fraction)
        self.dirichlet_alpha = float(dirichlet_alpha)
        self.dirichlet_noise_fraction = float(dirichlet_noise_fraction)

    def _rollout(self, game: Chess) -> float:
        """Simple random rollout for leaf evaluation when no value network.

        Uses push/pop (apply_move / undo_move) on the provided `game` to avoid
        allocating a full copy per rollout. The rollout will undo its own
        moves before returning.
        """
        depth = 0
        max_depth = 60  # reduced depth for faster rollouts (configurable)
        moves_played = []
        while not game.is_terminal() and depth < max_depth:
            moves = game.get_legal_moves()
            if not moves:
                break
            m = random.choice(moves)
            game.apply_move(m)
            moves_played.append(m)
            depth += 1

        winner = game.check_winner()

        # Undo rollout moves
        for _ in range(len(moves_played)):
            game.undo_move()

        return float(winner)  # +1, -1, or 0

    class Node:
        def __init__(self, prior: float = 0.0, zone: float = 0.0):
            self.prior = prior
            self.zone = zone
            self.visit_count = 0
            self.value_sum = 0.0
            self.children: Dict[chess.Move, 'MCVSSearcher.Node'] = {}

    def search_with_time_budget(self, game: Chess, time_budget: float) -> Tuple[Dict[chess.Move, int], int]:
        root = self.Node()
        start_time = time.time()
        simulations = 0

        while time.time() - start_time < time_budget:
            current_game = game
            node = root
            path = []
            moves_made = []

            # Selection
            while node.children:
                path.append(node)
                total_n = sum(c.visit_count for c in node.children.values()) + 1
                best_score = -float('inf')
                best_move = None
                best_child = None
                for move, child in node.children.items():
                    q = child.value_sum / child.visit_count if child.visit_count > 0 else 0.0
                    p_lambda = max(child.prior + self.lambda_zone * child.zone, 0.001)
                    u = self.cpuct * p_lambda * np.sqrt(total_n) / (1 + child.visit_count)
                    score = q + u
                    if score > best_score:
                        best_score = score
                        best_move = move
                        best_child = child
                current_game.apply_move(best_move)
                moves_made.append(best_move)
                node = best_child

            path.append(node)

            # Terminal node?
            if current_game.is_terminal():
                value = float(current_game.check_winner())
            else:
                # Expansion & Evaluation
                legal_moves = current_game.get_legal_moves()

                if self.use_nets and self.policy_net is not None and self.value_net is not None:
                    # === Neural network branch ===
                    W = current_game.get_weighted_adjacency_matrix()
                    W_tensor = torch.from_numpy(W)[None, None, :, :].float().to(self.device)
                    with torch.no_grad():
                        logits = self.policy_net(W_tensor)[0]
                        value = self.value_net(W_tensor)[0, 0].item()

                    probs = F.softmax(logits, dim=0).cpu().numpy()
                    prior_dict = {}
                    prior_sum = 0.0
                    for m in legal_moves:
                        idx = m.from_square * 64 + m.to_square
                        p = float(probs[idx])
                        prior_dict[m] = p
                        prior_sum += p
                    if prior_sum > 0.0:
                        for m in prior_dict:
                            prior_dict[m] /= prior_sum
                    else:
                        p = 1.0 / len(legal_moves)
                        prior_dict = {m: p for m in legal_moves}
                else:
                    # === Zone-only / no-nets branch ===
                    value = self._rollout(current_game)
                    prior_dict = {m: 1.0 / len(legal_moves) for m in legal_moves} if legal_moves else {}

                # Dirichlet noise only at root and only when using nets.
                # Use configured hyperparameters so the same alpha/fraction
                # can be applied to both priors and zone guidance.
                if len(path) == 1 and self.use_nets and legal_moves:
                    noise = np.random.dirichlet([self.dirichlet_alpha] * len(legal_moves))
                    for i, m in enumerate(legal_moves):
                        prior_dict[m] = (1 - self.dirichlet_noise_fraction) * prior_dict[m] + \
                                        self.dirichlet_noise_fraction * noise[i]
                    # Renormalize
                    prior_sum = sum(prior_dict.values())
                    if prior_sum > 0:
                        for m in prior_dict:
                            prior_dict[m] /= prior_sum

                # Create children with zone scores. Use apply/undo on the shared
                # board to avoid allocating a full copy per child. Skip heavy
                # weighted-matrix computations when zone guidance is disabled
                # (lambda_zone == 0) or when the DB is empty.
                zone_db_empty = (len(self.zone_db.winning_matrices) == 0 and
                                 len(self.zone_db.losing_matrices) == 0 and
                                 len(self.zone_db.draw_matrices) == 0)
                need_zone = (self.lambda_zone != 0.0) and (not zone_db_empty)

                for i, m in enumerate(legal_moves):
                    if need_zone:
                        current_game.apply_move(m)
                        child_W = current_game.get_weighted_adjacency_matrix()
                        child_Z = self.zone_db.compute_zone_score(child_W, k=self.k_zone)
                        current_game.undo_move()
                    else:
                        child_Z = 0.0

                    # If Dirichlet noise was added to priors at the root, apply the same
                    # noise (mapped to [-1,1]) to zone guidance so both signals share
                    # the same stochastic perturbation controlled by the same alpha.
                    if len(path) == 1 and self.use_nets and legal_moves:
                        z_noise = (noise[i] * 2.0) - 1.0  # map [0,1] -> [-1,1]
                        child_Z = (1 - self.dirichlet_noise_fraction) * child_Z + \
                                  self.dirichlet_noise_fraction * z_noise

                    child = self.Node(prior=prior_dict.get(m, 0.0), zone=child_Z)
                    node.children[m] = child

            # Backpropagation
            v = value
            for n in reversed(path):
                n.visit_count += 1
                n.value_sum += v
                v = -v

            # Undo moves that were applied during selection (the rollout
            # already undid its own moves). This restores `game` to the root
            # state for the next simulation.
            for _ in range(len(moves_made)):
                current_game.undo_move()

            simulations += 1

        visits = {move: child.visit_count for move, child in root.children.items()}
        if not visits:
            legal = game.get_legal_moves()
            visits = {random.choice(legal): 1} if legal else {}
        return visits, simulations

# ============================================================================
# UCT SEARCHER (your original from the provided file)
# ============================================================================

class UCTSearcher:
    def __init__(self, cpuct=np.sqrt(2.0)):
        self.cpuct = cpuct

    class Node:
        def __init__(self):
            self.visit_count = 0
            self.value_sum = 0.0
            self.children: Dict[chess.Move, 'UCTSearcher.Node'] = {}

    def _rollout(self, game: Chess) -> float:
        depth = 0
        max_depth = 60
        moves_played = []
        while not game.is_terminal() and depth < max_depth:
            moves = game.get_legal_moves()
            if not moves:
                break
            m = random.choice(moves)
            game.apply_move(m)
            moves_played.append(m)
            depth += 1

        if depth >= max_depth:
            # undo rollout moves
            for _ in range(len(moves_played)):
                game.undo_move()
            return 0.0

        winner = game.check_winner()

        # undo rollout moves
        for _ in range(len(moves_played)):
            game.undo_move()

        if winner == 1:
            return 1.0
        elif winner == -1:
            return -1.0
        else:
            return 0.0

    def search_with_time_budget(self, game: Chess, time_budget: float) -> Tuple[Dict[chess.Move, int], int]:
        root = self.Node()
        start_time = time.time()
        simulations = 0

        while time.time() - start_time < time_budget:
            current_game = game
            node = root
            path = [node]
            moves_made = []

            while True:
                legal_moves = current_game.get_legal_moves()

                if current_game.is_terminal():
                    winner = current_game.check_winner()
                    value = 1.0 if winner == 1 else -1.0 if winner == -1 else 0.0
                    break

                if not legal_moves:
                    value = 0.0
                    break

                unvisited_moves = [m for m in legal_moves if m not in node.children]
                if unvisited_moves:
                    move = random.choice(unvisited_moves)
                    child = self.Node()
                    node.children[move] = child
                    current_game.apply_move(move)
                    moves_made.append(move)
                    path.append(child)
                    value = self._rollout(current_game)
                    break

                best_score = -float('inf')
                best_move = None
                best_child = None
                total_visits = node.visit_count

                for move in legal_moves:
                    child = node.children[move]
                    Q = child.value_sum / (child.visit_count + 1e-8) if child.visit_count > 0 else 0.0
                    U = self.cpuct * np.sqrt(np.log(total_visits + 1) / (child.visit_count + 1e-8))
                    score = Q + U
                    if score > best_score:
                        best_score = score
                        best_move = move
                        best_child = child

                current_game.apply_move(best_move)
                moves_made.append(best_move)
                node = best_child
                path.append(node)

            for n in reversed(path):
                n.visit_count += 1
                n.value_sum += value
                value = -value

            # Undo selection moves to restore root state
            for _ in range(len(moves_made)):
                current_game.undo_move()

            simulations += 1

        visits = {move: child.visit_count for move, child in root.children.items()}
        if not visits:
            legal = game.get_legal_moves()
            return {random.choice(legal): 1} if legal else {}, simulations
        return visits, simulations

if __name__ == "__main__":
    print("chess_mcvs.py loaded successfully")


def load(ctx):
    """Platform contract: read the zone db from ctx.data_dir ourselves."""
    from pathlib import Path

    search_dirs = [d for d in (ctx.data_dir, ctx.model_dir) if d is not None]
    npz_path = None
    for d in search_dirs:
        candidates = sorted(Path(d).rglob("*.npz"))
        if candidates:
            npz_path = candidates[0]
            break
    if npz_path is None:
        raise FileNotFoundError(
            "no zone-db .npz file found under data_dir or model_dir — "
            "this model requires one to function"
        )
    return HilbertOrderedZoneDatabase(filepath=str(npz_path))


def get_move(state, fen: str, player: str, time_budget: float = 2.0) -> str:
    import chess as _chess
    import json
    from pathlib import Path

    config_path = Path(__file__).parent / "config_model.json"
    config = json.loads(config_path.read_text()) if config_path.exists() else {}

    game = Chess()
    game.board = _chess.Board(fen)
    zone_db = state
    searcher = MCVSSearcher(
        policy_net=None,
        value_net=None,
        zone_db=zone_db,
        use_nets=False,
        lambda_zone=0.0 if zone_db is None else float(config.get("lambda_zone", 1.0)),
        k_zone=int(config.get("k_zone", 5)),
        cpuct=float(config.get("cpuct", 1.5)),
        dirichlet_alpha=float(config.get("dirichlet_alpha", 0.3)),
    )
    time_budget = float(config.get("time_per_move", 2.0))
    visit_counts, _ = searcher.search_with_time_budget(game, time_budget)
    if not visit_counts:
        legal = game.get_legal_moves()
        return legal[0].uci() if legal else ""
    best_move = max(visit_counts, key=visit_counts.get)
    return best_move.uci()