mcts.py

import math
import torch
import torch.nn.functional as F
from tenerec_three import Tenerec3Chess
from game import Game
from chessGame import ChessPos, _print_board, _get_all_possible_moves
import numpy as np
from chessGame import _flip_move_mask


def ucb_score(parent, child):
    """
    The score for an action that would transition between the parent and child.
    """
    prior_score = child.prior * math.sqrt(parent.visit_count) / (child.visit_count + 1)
    if child.visit_count > 0:
        # The value of the child is from the perspective of the opposing player
        value_score = -child.value()
    else:
        value_score = 0

    return value_score + prior_score

def root_noise(root_node, alpha=0.3, epsilon=0.25):
    if not root_node.is_expanded(): return
        
    children = list(root_node.children.values())
    n_actions = len(children)
    
    if n_actions <= 1: return
   
    noise = np.random.dirichlet([alpha] * n_actions)
    
    # Modificamos los priors en caliente
    for i, child in enumerate(children):
        child.prior = (1 - epsilon) * child.prior + epsilon * noise[i]

class Node:
    def __init__(self, prior, to_play):
        self.to_play = to_play
        self.visit_count = 0
        self.prior = prior
        self.value_sum = 0
        self.children = {}
        self.state = None

    def value(self):
        if self.visit_count == 0: return 0
        return self.value_sum / self.visit_count
    
    def select_child(self):
        max_ucb = -float('inf')
        best_move = None
        best_child = None
        for move, child in self.children.items():
            ucb = ucb_score(self, child)
            if ucb > max_ucb: 
                max_ucb = ucb
                best_move = move
                best_child = child
        return best_move, best_child
    def expand(self, state, move_probs:dict):
        self.state = state
        for i, (move, prior) in enumerate(move_probs.items()):
            self.children[move] = Node(prior, -self.to_play)

    def is_expanded(self) -> bool:
        return len(self.children) > 0

class MCTS:
    def __init__(self, game:Game, model:Tenerec3Chess, num_simulations):
        self.game = game
        self.model = model
        self.num_simulations = num_simulations

    def get_move_probs_and_state_value(self, state, to_play):
        input = torch.from_numpy(self.game.get_board_input(state)).float().unsqueeze(0).to('cuda')
        mask = torch.from_numpy(self.game.get_move_mask(state)).float().to('cuda')
        logits, value = self.model(input)
        logits = logits.squeeze(0)
        masked_logits = logits.masked_fill(mask == 0, float('-inf'))
        policy = F.softmax(masked_logits.flatten(), dim=0).reshape_as(masked_logits)
        return self.game.get_move_probs(policy, to_play), value.item()

    def run(self, root:Node, state):
        with torch.no_grad():
            if not root.is_expanded():
                move_probs, _ = self.get_move_probs_and_state_value(state, root.to_play)
                root.expand(state, move_probs)
            root_noise(root)
            for _ in range(self.num_simulations - root.visit_count):
                node = root
                search_path = [root]
                try:
                    while node.is_expanded():
                        move,child = node.select_child()
                        search_path.append(child)
                        node = child
                    node_state = self.game.make_move(search_path[-2].state, move)
                    value = self.game.get_result(node_state, node.to_play)
                    if value is None:
                        move_probs, value = self.get_move_probs_and_state_value(node_state, node.to_play)
                        node.expand(node_state, move_probs)
                    self.backpropagate(search_path, value, node.to_play)
                except ValueError as e:
                    print(e)
                    return [node.state for node in search_path[:-1]]
        return None
    def backpropagate(self, search_path, value, to_play):
        for node in search_path:
            node.value_sum += value * (-1)**(to_play != node.to_play)
            node.visit_count+=1



class MatchMaker:
    def __init__(self, game:Game, model, max_moves, mcts_simulations):
        self.game = game
        self.model = model
        self.max_moves = max_moves
        self.mcts_simulations = mcts_simulations
        self.mcts = MCTS(self.game, self.model, self.mcts_simulations)

    def _print_move_traceback(self, pos_list):
        """Print the full move list with board after each position."""
        print(f"\n{'='*60}")
        print(f"  MOVE TRACEBACK — {len(pos_list)} positions")
        print(f"{'='*60}")
        for i, item in enumerate(pos_list, start=1):
            print()
            print(i)
            print(_print_board(item))
            print()
        print(f"{'='*60}\n")

    def run_match(self, n_step=5):
        """Run a self-play match, collecting position, policy, and value data.

        For terminal wins/losses, value targets are the actual game result
        (±1 from each player's perspective). For true draws (stalemate and
        insufficient material), the target is 0.
        For non-terminal games and 50-move rule draws, value targets use
        n-step MCTS lookahead: position k's target = MCTS root value at
        position k+n (converted to position k's player perspective). If k+n
        exceeds the game length, uses the last available position's MCTS value.
        A 50-move draw is a property of the continuation played, not of the
        position, so it gets the same n-step treatment as a non-terminal game.

        Sign logic (critical):
          MCTS values are from the perspective of the player TO MOVE at that
          position. To convert from position j's perspective to position k's:
            target_k = v_j * turn_k * turn_j
          where turn = +1 for white, -1 for black.
        """
        self.model.eval()
        current_pos = ChessPos()
        current_node = Node(0, 1)
        current_node.visit_count = 1
        pos_list = []
        mcts_values = []  # (root_value, to_play) per stored position
        policy_targets = []
        for turn in range(self.max_moves):
            temp = 1 if turn <= 60 else 0.2
            if self.game.is_game_over(current_pos): 
                break

            result = self.mcts.run(current_node, current_pos)
            if result is not None:
                self._print_move_traceback(pos_list+result)
                raise RuntimeError("El rey ha desaparecido, aquí el traceback ^")

            weights = np.zeros(len(current_node.children), dtype=np.float64)
            policy = np.zeros(shape=(73, 8, 8))
            for i, (move, c) in enumerate(current_node.children.items()):
                ch, ox, oy = self.game.encode_move(*move)
                policy[ch, oy, ox] = c.visit_count / (current_node.visit_count - 1)
                weights[i] = c.visit_count ** (1/temp)
            weights /= weights.sum()
            next_move_idx = np.random.choice(len(weights), p=weights)
            next_move = list(current_node.children.keys())[next_move_idx]
            next_node = current_node.children[next_move]
            pos_list.append(current_pos)
            mcts_values.append(current_node.value())
            if current_pos.turn == -1:
                policy = _flip_move_mask(policy)
            policy_targets.append(policy)
            current_pos = self.game.make_move(current_pos, next_move)
            del current_node.children[next_move]
            current_node.children.clear()
            del current_node
            current_node = next_node

        moves_played = len(pos_list)
        if moves_played == 0:
            return [], np.array([]), np.array([])

        last_value_white = None
        is_fifty_move_draw = False
        if self.game.is_game_over(current_pos):
            last_value_white = self.game.get_result(current_pos, 1)
            # A 50-move draw says nothing about the value of the positions:
            # the draw is a property of the continuation played, not of the
            # board. Route it through the same n-step bootstrap as
            # non-terminal games so the value gradient survives.
            # Exclude true draws: stalemate at the 50-move mark has no legal
            # moves (insufficient material can't coincide with rule50>=100
            # because material only changes via captures/pawn moves, which
            # reset rule50).
            is_fifty_move_draw = (
                last_value_white == 0.0
                and current_pos.rule50 >= 100
                and len(_get_all_possible_moves(current_pos)) > 0
            )

        if last_value_white is not None and last_value_white != 0.0:
            value_targets = np.array([last_value_white * pos.turn for pos in pos_list])
            if last_value_white == 1.0:
                result = f"White wins in {moves_played} moves"
            else:
                result = f"Black wins in {moves_played} moves"
        elif is_fifty_move_draw or last_value_white is None:
            value_targets = np.zeros(moves_played)
            for k in range(moves_played):
                target_idx = min(k + n_step, moves_played - 1)
                value_targets[k] = mcts_values[target_idx] * (1 if (target_idx - k) % 2 == 0 else -1)
            if is_fifty_move_draw:
                result = f"50-move draw (n-step={n_step}) in {moves_played} moves"
            else:
                result = f"Non-terminal (n-step={n_step}) in {moves_played} moves"
        else:
            value_targets = np.zeros(moves_played)
            result = f"Draw in {moves_played} moves"

        print(result)
        policy_targets = np.stack(policy_targets)
        return pos_list, policy_targets, value_targets