"""Chapter 10 lab, part 1: Monte Carlo tree search (UCT) for tic-tac-toe.

Fill in the four TODOs in order (Tasks 1-4). The tests in test_mcts.py check
each one. Everything marked "provided" is complete.

Run the tests:   python -m pytest -q test_mcts.py
Then play:       python run.py
"""

import math
import random

from tictactoe import winner, to_move, legal_moves, play


class Node:
    """One node of the search tree = one board position.

    Attributes
    ----------
    board : tuple of 9 ints, the position at this node.
    move : the move (0..8) that led here from the parent, None at the root.
    parent : the parent Node, None at the root.
    children : list of expanded child Nodes.
    untried : legal moves from this position that have no child node yet.
    N : visit count (number of simulations that passed through this node).
    W : total result of those simulations, from the point of view of
        `player_just_moved` (the player who made `move`): +1 per win, -1 per loss.
    player_just_moved : +1 (X) or -1 (O).
    terminal : True if the game is over at this node.
    """

    def __init__(self, board, move=None, parent=None):
        self.board = board
        self.move = move
        self.parent = parent
        self.children = []
        self.untried = legal_moves(board)
        self.N = 0
        self.W = 0.0
        self.player_just_moved = -to_move(board)
        self.terminal = winner(board) is not None

    @property
    def Q(self):
        """Mean result for the player who moved into this node (0 if unvisited)."""
        return self.W / self.N if self.N > 0 else 0.0

    def __repr__(self):
        return f"Node(move={self.move}, N={self.N}, W={self.W:+.0f}, Q={self.Q:+.3f})"


# ---------------------------------------------------------------- Task 1
def uct_score(child_W, child_N, parent_N, c):
    """UCT score of a child, as seen by the player choosing at the parent.

        UCT = W/N + c * sqrt( ln(N_parent) / N )

    Returns math.inf if the child has never been visited (child_N == 0).
    """
    # TODO: return math.inf for unvisited children, else Q + exploration bonus.
    raise NotImplementedError("Task: implement uct_score")


# ---------------------------------------------------------------- Task 2
def rollout(board, rng):
    """Play uniformly random moves (rng.choice over legal_moves) until the game
    ends. Returns the result from X's point of view: +1, -1 or 0.
    If `board` is already terminal, return its result without moving."""
    # TODO: loop: pick rng.choice(legal_moves(board)), play it, until winner(board) is not None.
    raise NotImplementedError("Task: implement rollout")


# ---------------------------------------------------------------- Task 3
def backup(path, result):
    """Propagate a simulation's result up the path [root, ..., leaf].

    Every node gets N += 1, and W += result * node.player_just_moved, so that W
    is always measured from the point of view of the player who moved into it.
    """
    # TODO: for each node on the path: N += 1, W += result * node.player_just_moved.
    raise NotImplementedError("Task: implement backup")


# ---------------------------------------------------------------- Task 4
def run_simulation(root, c, rng):
    """One MCTS simulation from `root`. Returns the rollout result (X's view).

    1. Selection: starting at the root, while the node is not terminal, has no
       untried moves and has children, move to the child with the highest
       uct_score(child.W, child.N, node.N, c).
    2. Expansion: if the node reached is not terminal and has untried moves,
       pick one with rng.randrange(len(node.untried)), remove it from
       `untried`, create the child Node, append it to `children` and to the path.
    3. Rollout: rollout(leaf.board, rng) from the last node on the path.
    4. Backup: backup(path, result).
    """
    node = root
    path = [node]
    # 1. Selection
    # TODO: while the node is fully expanded (not terminal, no untried moves, has
    #       children), step to the child with the highest UCT score; append it.

    # 2. Expansion
    # TODO: if not terminal and some moves are untried, pop one at random
    #       (rng.randrange), create Node(play(node.board, move), move, node),
    #       add it to node.children and to the path.

    # 3. Rollout and 4. Backup
    # TODO: result = rollout(...) from the last node; backup(path, result); return result.
    raise NotImplementedError("Task: implement run_simulation")


# ---------------------------------------------------------------- provided
def search(board, n_simulations, c=1.4, seed=0):
    """Run `n_simulations` simulations from `board`; return the root Node."""
    rng = random.Random(seed)
    root = Node(board)
    for _ in range(n_simulations):
        run_simulation(root, c, rng)
    return root


def best_move(root):
    """The move MCTS plays: the most visited child (ties broken by Q)."""
    return max(root.children, key=lambda ch: (ch.N, ch.Q)).move


def mcts_player(n_simulations, c=1.4, seed=0):
    """A player function board -> move that runs a fresh search each turn."""
    counter = [seed]

    def choose(board):
        counter[0] += 1
        return best_move(search(board, n_simulations, c, seed=counter[0]))

    return choose