"""Tests for mcts.py. Run: python -m pytest -q test_mcts.py"""

import importlib
import math
import os
import random

import pytest

from tictactoe import (EMPTY, FORK_TRAP, MUST_BLOCK, WIN_IN_ONE, play, play_game,
                       random_player)

mcts = importlib.import_module("solutions.mcts" if os.environ.get("RL_SOLUTIONS") else "mcts")
Node = mcts.Node


# ------------------------------------------------------------ Task 1: uct_score
def test_uct_score_by_hand():
    # Q = 3/5 = 0.6, bonus = 1.4 * sqrt(ln 20 / 5) = 1.4 * 0.774052... = 1.083673...
    got = mcts.uct_score(3.0, 5, 20, 1.4)
    assert got == pytest.approx(0.6 + 1.4 * math.sqrt(math.log(20) / 5), abs=1e-9), \
        f"UCT(W=3, N=5, N_parent=20, c=1.4) should be about 1.6837, got {got}"


def test_uct_chapter_table():
    # The worked example in Chapter 10: N_parent = 20, c = 1.4, children A, B, C.
    got = [mcts.uct_score(W, N, 20, 1.4) for W, N in [(6, 10), (2, 6), (1, 4)]]
    assert got == pytest.approx([1.366, 1.323, 1.462], abs=1e-3), \
        f"the chapter's table gives UCT = 1.366, 1.323, 1.462 for A, B, C; got {[round(g, 3) for g in got]}"


def test_uct_unvisited_is_infinite():
    assert mcts.uct_score(0.0, 0, 7, 1.4) == math.inf, "an unvisited child must score +inf"


def test_uct_c_zero_is_mean():
    assert mcts.uct_score(-2.0, 4, 50, 0.0) == pytest.approx(-0.5), "with c = 0 the score is just Q = W/N"


def test_uct_bonus_shrinks_with_visits():
    a = mcts.uct_score(5.0, 10, 100, 1.0)
    b = mcts.uct_score(50.0, 100, 1000, 1.0)
    assert a > b, "same Q = 0.5, but the less-visited child (relative to its parent) should get the larger bonus"


# ------------------------------------------------------------ Task 2: rollout
def test_rollout_terminal_returns_result():
    x_won = (1, 1, 1, -1, -1, 0, 0, 0, 0)
    assert mcts.rollout(x_won, random.Random(0)) == 1
    o_won = (1, 1, 0, -1, -1, -1, 1, 0, 0)
    assert mcts.rollout(o_won, random.Random(0)) == -1


def test_rollout_one_move_left():
    # X to move, only cell 8 is empty. X at 8 completes no line (2-5-8, 6-7-8
    # and 0-4-8 all contain an O), so the game ends in a draw.
    board = (1, -1, 1,
             1, -1, -1,
             -1, 1, 0)
    assert mcts.rollout(board, random.Random(0)) == 0, "the only move (8) gives a full board with no line: draw"


def test_rollout_statistics_from_empty_board():
    # Known facts about random tic-tac-toe: X wins ~58.5%, O ~28.8%, draws ~12.7%.
    rng = random.Random(123)
    n = 4000
    results = [mcts.rollout(EMPTY, rng) for _ in range(n)]
    px = results.count(1) / n
    po = results.count(-1) / n
    assert abs(px - 0.585) < 0.035, f"random play: X should win about 58.5% of games, got {px:.3f}"
    assert abs(po - 0.288) < 0.035, f"random play: O should win about 28.8% of games, got {po:.3f}"


# ------------------------------------------------------------ Task 3: backup
def test_backup_signs():
    root = Node(EMPTY)                      # nobody has moved: player_just_moved = O (-1)
    child = Node(play(EMPTY, 4), 4, root)   # X just moved
    grand = Node(play(child.board, 0), 0, child)  # O just moved
    path = [root, child, grand]
    mcts.backup(path, +1)                   # X wins the rollout
    assert [n.N for n in path] == [1, 1, 1], "every node on the path gets one more visit"
    assert [n.W for n in path] == [-1, 1, -1], \
        "X won: W should go +1 at nodes where X just moved and -1 where O just moved"
    mcts.backup(path, 0)                    # a draw
    assert [n.N for n in path] == [2, 2, 2]
    assert [n.W for n in path] == [-1, 1, -1], "a draw adds 0 to W"
    mcts.backup(path, -1)                   # O wins
    assert [n.W for n in path] == [0, 0, 0]


# ------------------------------------------------------------ Task 4: run_simulation
def test_first_nine_simulations_expand_every_root_move():
    rng = random.Random(0)
    root = Node(EMPTY)
    for k in range(1, 10):
        mcts.run_simulation(root, 1.4, rng)
        assert root.N == k, "each simulation visits the root once"
        assert len(root.children) == k, \
            "while the root has untried moves, each simulation should expand one new root child"
    assert sorted(ch.move for ch in root.children) == list(range(9))
    assert all(ch.N == 1 for ch in root.children)
    assert root.untried == []


def test_tenth_simulation_selects_by_uct():
    # After 9 simulations from the empty board every root child has N = 1, so
    # all exploration bonuses are equal and UCT must pick a child with the
    # largest W (W is +1 where X, who moved into these children, won its rollout).
    # The 10th simulation must go through such a child and expand one of its moves.
    for seed in range(5):
        rng = random.Random(seed)
        root = Node(EMPTY)
        for _ in range(9):
            mcts.run_simulation(root, 1.4, rng)
        W_before = {ch.move: ch.W for ch in root.children}
        best_W = max(W_before.values())
        mcts.run_simulation(root, 1.4, rng)
        chosen = [ch for ch in root.children if ch.N == 2]
        assert len(chosen) == 1, "the 10th simulation should pass through exactly one root child"
        ch = chosen[0]
        assert W_before[ch.move] == best_W, (
            f"selection should take a child with the highest UCT score (here: highest W = {best_W}), "
            f"but went through move {ch.move} with W = {W_before[ch.move]}. Selection should take the max "
            "of uct_score(child.W, child.N, node.N, c) over the node's children.")
        assert len(ch.children) == 1, \
            "after selecting a root child, the simulation should expand one of that child's moves"


def test_tree_bookkeeping_is_consistent():
    root = mcts.search(EMPTY, 300, c=1.4, seed=1)
    assert root.N == 300
    assert sum(ch.N for ch in root.children) == 300, \
        "every simulation passes through exactly one child of the root"

    def check(node):
        assert abs(node.W) <= node.N
        if node.children:
            inner = sum(ch.N for ch in node.children)
            assert inner <= node.N
        for ch in node.children:
            assert ch.parent is node and ch.player_just_moved == -node.player_just_moved
            check(ch)
    check(root)


def test_terminal_root_is_handled():
    x_won = (1, 1, 1, -1, -1, 0, 0, 0, 0)
    root = Node(x_won)
    r = mcts.run_simulation(root, 1.4, random.Random(0))
    assert r == 1 and root.N == 1 and root.children == [], \
        "at a finished game there is nothing to select or expand; just back up the result"


def test_finds_win_in_one():
    root = mcts.search(WIN_IN_ONE, 200, seed=0)
    assert mcts.best_move(root) == 2, "X should complete the top row at cell 2"


def test_blocks_immediate_threat():
    root = mcts.search(MUST_BLOCK, 400, seed=0)
    assert mcts.best_move(root) == 8, "O must block the diagonal 0-4-8 at cell 8"


def test_avoids_fork_trap():
    root = mcts.search(FORK_TRAP, 1500, seed=0)
    assert mcts.best_move(root) in (1, 3, 5, 7), \
        "O must play an edge here; a corner (2 or 6) lets X set up a fork"


def test_never_loses_to_random():
    # 20 games as X and 20 as O. MCTS with random rollouts is not a perfect
    # player (very rarely its noisy estimates mislead it), so we allow at most
    # one loss; a correct implementation almost always loses none.
    rng = random.Random(7)
    losses = 0
    for g in range(20):
        if play_game(mcts.mcts_player(500, seed=100 + g), random_player(rng)) == -1:
            losses += 1
        if play_game(random_player(rng), mcts.mcts_player(500, seed=200 + g)) == 1:
            losses += 1
    assert losses <= 1, f"MCTS with 500 simulations per move should (essentially) never lose to a random player; lost {losses} of 40"