"""
=============================================================================
第 14 讲 · 共享模块：蒙特卡洛树搜索（MCTS / UCT）与 PUCT
=============================================================================

对应讲义：lectures/lecture14_mcts_alphazero.md 第 3–5 节

-----------------------------------------------------------------------------
【MCTS 的四个步骤】（每次"模拟"都重复一遍）
-----------------------------------------------------------------------------
    1. 选择（Selection）  ：从根节点出发，在已经展开的树中，按"树策略"一路往下选，直到遇到一个还有
                            没试过的着法的节点（或终局）。UCT 的树策略就是第 11 讲的 UCB：
                                选 argmax_a [ Q(s,a) + c · sqrt( ln N(s) / N(s,a) ) ]
    2. 扩展（Expansion）  ：在这个节点上，随机选一个没试过的着法，给树加一个新的子节点。
    3. 模拟（Simulation） ：从新节点出发，双方都随机落子直到终局（"随机走子 / rollout"），得到胜负。
    4. 回传（Backup）     ：把这个胜负沿着刚才走过的路径一路传回根节点，更新每个节点的访问次数 N 和总价值 W。

模拟足够多次之后，选根节点下**访问次数最多**的着法。

-----------------------------------------------------------------------------
【为什么这样能行？】
-----------------------------------------------------------------------------
- 随机走子给出一个（有噪声的）局面评估：随机下到底，赢的比例越高，局面大概越好。
- UCB 把模拟**集中**在有希望的着法上，但也不完全放弃看起来差的着法——
  于是树会在"好棋"的方向上长得很深，在"坏棋"的方向上只浅浅地看几眼。
- 树越长越深，树内部的对手也在"用 UCB 选最好的应对"，评估就从"对手随机下"逐渐变成"对手认真下"，
  最终收敛到极小极大的结果。这是它比"平坦蒙特卡洛"（只在根节点做随机走子）强的根本原因。

-----------------------------------------------------------------------------
【价值的视角】
-----------------------------------------------------------------------------
每个节点记录"走到这个节点的那一方"的平均得分（赢 1、平 0、输 -1）。
父节点在挑选子节点时，挑的正是"对自己最有利"的那个——两方都在用同一套规则最大化自己的得分。
"""

import math

import numpy as np

from tictactoe import legal_moves, play, to_move, winner

LINES_THROUGH = {i: [ln for ln in [(0, 1, 2), (3, 4, 5), (6, 7, 8), (0, 3, 6), (1, 4, 7),
                                   (2, 5, 8), (0, 4, 8), (2, 4, 6)] if i in ln] for i in range(9)}


def random_rollout(board, rng):
    """
    随机走子：双方都随机落子，直到终局。返回赢家（1 / -1），平局返回 0。

    【实现上的小优化】每落一子只检查经过这个位置的几条线，而不是全部 8 条线；
        并且一次性把空位打乱，按顺序落子——比每一步都调用 winner() 快得多。
        MCTS 的速度很大程度上取决于随机走子的速度。
    """
    w = winner(board)
    if w is not None:
        return w
    b = list(board)
    player = to_move(board)
    empties = [i for i in range(9) if b[i] == 0]
    rng.shuffle(empties)
    for m in empties:
        b[m] = player
        for x, y, z in LINES_THROUGH[m]:
            if b[x] == b[y] == b[z] == player:
                return player
        player = -player
    return 0


class Node:
    """
    搜索树中的一个节点。

    【属性】
        board    : 这个节点对应的局面
        parent   : 父节点（根节点为 None）
        children : 字典 着法 → 子节点
        untried  : 还没有展开的着法
        N        : 访问次数
        W        : 总得分（从"走到这个节点的那一方"的角度）
        prior    : 先验概率 P(s, a)（只有 PUCT 用到，见 AlphaZero）
    """
    __slots__ = ("board", "parent", "children", "untried", "N", "W", "prior")
    # 【Python 小知识】__slots__ 告诉 Python 这个类只有这些属性，能节省内存、稍微加快属性访问

    def __init__(self, board, parent=None, prior=0.0):
        self.board = board
        self.parent = parent
        self.children = {}
        self.untried = legal_moves(board) if winner(board) is None else []
        self.N = 0
        self.W = 0.0
        self.prior = prior


def uct_select(node, c):
    """
    UCT 的树策略：在已展开的子节点中选 Q + c · sqrt(ln N_parent / N_child) 最大的。
    Q = W / N 是子节点的平均得分（从"当前节点轮到的一方"的角度，因为子节点记录的是走到它的那一方）。
    【边界】c = 0 时就是纯贪心：只看平均得分，完全不探索。
    """
    log_n = math.log(node.N)
    best, best_score = None, -1e9
    for child in node.children.values():
        score = child.W / child.N + c * math.sqrt(log_n / child.N)
        if score > best_score:
            best, best_score = child, score
    return best


def mcts(board, n_sims, c, rng):
    """
    从 board 出发做 n_sims 次 MCTS 模拟（UCT 树策略 + 随机走子），返回根节点。
    调用者可以用 best_move(root) 选着法，或查看 root.children 的统计量。
    """
    root = Node(board)
    for _ in range(n_sims):
        node = root
        # 1. 选择：一直往下走，直到遇到还有没试过的着法的节点（或终局）
        while not node.untried and node.children:
            node = uct_select(node, c)
        # 2. 扩展：随机加一个子节点
        if node.untried:
            m = node.untried.pop(rng.integers(len(node.untried)))
            child = Node(play(node.board, m), parent=node)
            node.children[m] = child
            node = child
        # 3. 模拟：随机下到底
        result = random_rollout(node.board, rng)
        # 4. 回传：沿路径更新 N、W。每个节点的 W 从"走到它的那一方"的角度记分
        while node is not None:
            node.N += 1
            mover = -to_move(node.board)
            node.W += 1.0 if result == mover else (0.0 if result == 0 else -1.0)
            node = node.parent
    return root


def best_move(root, rng=None):
    """选访问次数最多的着法（最"稳健"的选择）；并列时随机。"""
    most = max(ch.N for ch in root.children.values())
    cands = [m for m, ch in root.children.items() if ch.N == most]
    return cands[0] if rng is None else cands[rng.integers(len(cands))]


def flat_monte_carlo(board, n_sims, rng):
    """
    平坦蒙特卡洛：不建树，把 n_sims 次随机走子**平均**分给根节点的每个着法，选平均得分最高的着法。

    【与 MCTS 的区别】它评估的是"我走这步之后，**双方都随机下**的胜率"。
        但真实的对手不会随机下——一个"对手随机下时很好，对手认真下时会输"的着法，平坦蒙特卡洛看不出来。
    """
    me = to_move(board)
    moves = legal_moves(board)
    scores = {m: 0.0 for m in moves}
    per = max(1, n_sims // len(moves))
    for m in moves:
        b = play(board, m)
        for _ in range(per):
            r = random_rollout(b, rng)
            scores[m] += 1.0 if r == me else (0.0 if r == 0 else -1.0)
    best = max(scores.values())
    cands = [m for m in moves if scores[m] == best]
    return cands[rng.integers(len(cands))], {m: scores[m] / per for m in moves}