2️⃣ Single-Game MCTS

Action = int # 0, 1, 2, 3, 4, 5, 6

@dataclass
class Node:
    """One node of the single-game search tree = one board state. Statistics live on the **edges**
    (per-action length-7 arrays), filled in as the search runs.

    Attributes:
        obs            (1, 3, H, W) the connect4 board for this state
        is_player1     (1,) Bool[Tensor, "1"] whos turn to play
        is_terminal    bool : True if this state is game-over. Set by `expand`.
        terminal_value float : result from perspective of `is_player1` (= -reward of the move that
                       created the node). Only meaningful when `is_terminal=True`. Set by `expand`.
        P              (7,) network prior over each move. Set by `evaluate`.
        legal          (7,) bool legal-column mask. Set by `evaluate`.
        N              (7,) per-edge visit counts N(s, a). Updated by `backup`.
        W              (7,) per-edge value sums W(s, a), accumulated from the perspective of THIS
                       node's mover (`is_player1`): backup negates each child value once per step on
                       the way up, so what lands here is already in this node's convention. Updated by `backup`.
        children       dict[Action, Node] : child nodes per played action, created lazily by `expand`.
        parent         Node | None : the node we were expanded from (`None` at the root).
        parent_action  Action | None : the column that led from `parent` to this node. Updated by `expand`.

    Properties:
        Q              (7,) per-edge mean value W / max(N, 1).
    """
    obs: Float[Tensor, "1 3 H W"] #lives on GPU
    is_player1: Bool[Tensor, "1"]
    num_actions: int = 7
    is_terminal: bool = False
    terminal_value: float = 0.0
    P: Tensor | None = None
    legal: Tensor | None = None
    N: Tensor = None
    W: Tensor = None
    children: dict[Action, 'Node'] = field(default_factory=dict)
    parent: 'Node | None' = None
    parent_action: Action | None = None

    def __post_init__(self):
        if self.N is None:
            self.N = torch.zeros(self.num_actions) # N(s,a_0), ..., N(s,a_6)
        if self.W is None:
            self.W = torch.zeros(self.num_actions) # W(s,a_0), ..., W(s,a_6)

    @property
    def Q(self):
        return self.W / torch.maximum(self.N, torch.ones_like(self.N))


# Sentinel "no action" returned by `select` when the selection walk stops on a terminal node.
# We use -100 rather than None because None silently inserts a dim when used to index a tensor
# (`t[None]`), whereas -100 is always an illegal column and raises loudly if ever misused as an index.
NULL_ACTION = -100

tests.test_mcts_node(Node)

def select_child(node : Node, c_puct: float) -> Action:
    """Return the action (column index) with the highest PUCT score at `node`.

    Everything you need is already stored on the node (see the `Node` class above):
        node.Q      (7,) per-edge mean values Q(s,a) = W / max(N, 1)  — a property, kept up to date
        node.N      (7,) per-edge visit counts N(s,a); its sum is the parent visit total Σ_a' N(s,a')
        node.P      (7,) the network prior p_θ(a|s), set when the node was evaluated
        node.legal  (7,) bool mask of legal columns

    Compute the PUCT score per action, set illegal actions to -inf, and take the argmax.

    Args:
        node:   the node to pick an action from (its edge statistics are already populated)
        c_puct: the exploration/exploitation trade-off constant c

    Returns:
        int: the PUCT-maximising column 0-6 (you may assume at least one action is legal).
        Return a plain Python int — `int(scores.argmax())`, not a 0-dim Tensor. A Tensor looks
        identical in comparisons but breaks `action in node.children` dict lookups in `select`
        (tensors hash by identity, not value).
    """
    raise NotImplementedError()


tests.test_select_child(select_child, Node)

def select(root: Node, 
           c_puct: float
) -> tuple[Node, Action]:
    """Walks down the tree from root: at each node take the PUCT-best action (`select_child`);
    if that action has a child, step into it and repeat. Halt at the first node whose PUCT-best
    action has NO child yet (`action not in node.children`), or on a terminal node.

    Note the halting rule is a property of the node's best *action*, not of the node itself: a
    node may already have children on other actions (and visits on them) while its current
    PUCT-best action is still unexpanded — that node is where the walk stops.

    Args:
        root: The root node of the search tree (never terminal).
        c_puct: The PUCT constant.

    Returns:
        * `(node, action)`: the node to expand and the action to play, OR
        * `(node, NULL_ACTION)` if the walk stopped on a terminal node (nothing to expand).
    """
    assert not root.is_terminal, "select: root must not be terminal"
    raise NotImplementedError()


tests.test_select(select, Node)

@torch.no_grad()
def expand(node: Node, 
           action: int, 
           env: Connect4Env) -> Node:
    """
    Plays the given action on the current board state, constructs a new node with the result,
    stores it as a child of the current node, and returns it.

    Side effects:
        * Mutates the current node to add a new child node.

    Args:
        node: The current node in the search tree.
        action: The action to play.
        env: The Connect4Env instance

    Returns:
        The newly created child node after the action has been played.
    """
    assert not node.is_terminal, "expand: cannot expand a terminal node"
    raise NotImplementedError()


tests.test_expand(expand)

@torch.no_grad()
def evaluate(node: Node, model: nn.Module, env: Connect4Env) -> float:
    """
    Estimates the value of a node from its mover's perspective.

    Side effects:
        If the node is non-terminal, mutates the node to 
        set the `node.P` and `node.legal` attributes.

    Args:
        node: The node to estimate the value of.
        model: The model to use.
        env: The Connect4Env instance

    Returns:
        For terminal nodes: the stored `terminal_value` (already from the mover's perspective).
        For non-terminal nodes: the critic head's value estimate, from the mover's perspective.
            Use `eval_net(model, obs, is_player1) -> (value, logits)` for the forward pass; the
            returned tensors live on the GPU, while the tree's per-node stats live in CPU memory.
    """
    raise NotImplementedError()


tests.test_evaluate(evaluate)

def backup(leaf: Node, value: float) -> None:
    """Walks from a node back up the tree to the root, updating the visit counts and value sums at each edge.

    Side effects:
        * Mutates the visit counts and value sums at each edge.

    Args:
        leaf: The leaf node to backup from.
        value: The value of the leaf node from its mover's perspective.

    Returns:
        None
    """
    raise NotImplementedError()


tests.test_backup(backup, Node)

def sample_tree_policy(visits: Float[Tensor, "7"], 
                       temperature: float = 1.0,
) -> Action:
    """Samples an action from the tree policy. `temperature` -> 0 is greedy argmax; `temperature` = 1
    samples proportional to visits; larger temperatures flattens the distribution.

    Args:
        visits: The visit counts for each action, shape `(7,)`.
        temperature: The temperature parameter.

    Returns:
        int: The action to sample.
    """
    raise NotImplementedError()


tests.test_sample_tree_policy(sample_tree_policy)

@torch.no_grad()
def mcts_search(
    root_obs: Float[Tensor, "1 3 H W"],
    root_is_player1: Bool[Tensor, "1"],
    model: nn.Module,
    env: Connect4Env,
    cfg: MCTSConfig,
    add_noise: bool = False,
    return_root: bool = False,
) -> Float[Tensor, "7"]:
    """Run `cfg.sims` MCTS simulations from the root; return the root's visit counts `(7,)` — or
    `(visit_counts, root)` when `return_root=True` (e.g. to inspect / visualise the search tree).
    """
    root = Node(root_obs, root_is_player1)
    evaluate(root, model, env) #required to set root.P and root.legal
    raise NotImplementedError()
    return (root.N, root) if return_root else root.N

# First check the search logic in isolation, with a dummy (uniform-policy, zero-value) network:
# a forced win-in-one must be found purely from the terminal reward backing up the tree.
tests.test_mcts_search(mcts_search)
# Then confirm the same search drives the real network correctly:
tests.test_mcts_search(mcts_search, Connect4Model(device).eval())

class DummyNet(nn.Module):

    def forward(self, x):
        b = x.shape[0]
        values = torch.zeros(b, device=x.device)
        logits = torch.zeros(b, 7, device=x.device)
        return values, logits

model = DummyNet()
obs, red = tests.diagonal_win_red()

print("Starting position (X = Red to move):")
print(render_board(obs, is_player1=True))

visits, root = mcts_search(obs, torch.tensor([red], device=device), model, env,
                           MCTSConfig(sims=32), return_root=True)
print("\nMCTS visit counts per column:", visits.int().tolist())
chosen = int(visits.argmax())
print(f"Most-visited column: {chosen}  ({int(visits[chosen])} of {int(visits.sum())} visits)")

obs_after = place_piece(obs, chosen, is_player1=True)
print(f"\nBoard after X plays column {chosen}  (completes the diagonal):")
print(render_board(obs_after))

# board + the visit-count policy pi(a) = N(s,a) / sum_a' N(s,a'), chosen column highlighted
plot_board_and_policy(obs, visits / visits.sum(), chosen_action=chosen,
                      title="MCTS finds the diagonal win")
# the search tree the simulations grew: edge thickness ~ visit count, terminal leaves in yellow
plot_mcts_tree(root, max_depth=2, title="MCTS search tree (after 32 sims)")
assert chosen == 4, "MCTS should find the diagonal win"