"""Chapter 10 lab: see MCTS get stronger with more simulations, watch flat Monte
Carlo fall into the fork trap, and sweep beta for the KL-regularized optimum.

    python run.py                  # uses your mcts.py and gibbs.py
    RL_SOLUTIONS=1 python run.py   # uses the reference solutions

Saves landscape.png next to this file. Takes about 30 seconds.
"""

import importlib
import os
import random

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from tictactoe import FORK_TRAP, legal_moves, play, play_game, random_player

SOL = bool(os.environ.get("RL_SOLUTIONS"))
mcts = importlib.import_module("solutions.mcts" if SOL else "mcts")
gibbs = importlib.import_module("solutions.gibbs" if SOL else "gibbs")
HERE = os.path.dirname(os.path.abspath(__file__))


def strength_vs_budget(budgets, games=40):
    """Win / draw / loss rates of MCTS against a random player (half as X, half as O)."""
    out = []
    rng = random.Random(0)
    for n in budgets:
        w = d = l = 0
        for g in range(games):
            if g % 2 == 0:
                res = play_game(mcts.mcts_player(n, seed=10_000 * g + n), random_player(rng))
            else:
                res = -play_game(random_player(rng), mcts.mcts_player(n, seed=10_000 * g + n))
            w += res == 1
            d += res == 0
            l += res == -1
        out.append((w / games, d / games, l / games))
    return np.array(out)


def flat_monte_carlo_move(board, n, rng):
    """No tree: spend n random rollouts spread evenly over the legal moves, and
    pick the move with the best average result for the player to move."""
    moves = legal_moves(board)
    me = 1 if sum(v != 0 for v in board) % 2 == 0 else -1
    per = max(1, n // len(moves))
    means = [np.mean([mcts.rollout(play(board, m), rng) * me for _ in range(per)]) for m in moves]
    return moves[int(np.argmax(means))]


def fork_trap_accuracy(budgets, seeds=40):
    edges = {1, 3, 5, 7}
    tree, flat = [], []
    for n in budgets:
        tree.append(np.mean([mcts.best_move(mcts.search(FORK_TRAP, n, seed=s)) in edges for s in range(seeds)]))
        flat.append(np.mean([flat_monte_carlo_move(FORK_TRAP, n, random.Random(s)) in edges for s in range(seeds)]))
    return np.array(tree), np.array(flat)


def main():
    budgets = [1, 3, 10, 30, 100, 300]
    wdl = strength_vs_budget(budgets)
    fork_budgets = [10, 30, 100, 300, 1000, 3000]
    tree_acc, flat_acc = fork_trap_accuracy(fork_budgets)

    pi0 = np.array([0.40, 0.25, 0.15, 0.10, 0.07, 0.03])
    r = np.array([0.1, 0.3, 0.6, 0.2, 0.9, 1.0])
    betas = np.logspace(-2, 1, 60)
    er, kl, jstar = [], [], []
    for b in betas:
        p = gibbs.gibbs_policy(pi0, r, b)
        er.append(p @ r)
        kl.append(np.sum(p[p > 0] * np.log(p[p > 0] / pi0[p > 0])))
        jstar.append(gibbs.kl_objective(p, pi0, r, b))

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
    ax[0].plot(budgets, wdl[:, 0], "o-", label="win")
    ax[0].plot(budgets, wdl[:, 1], "s-", label="draw")
    ax[0].plot(budgets, wdl[:, 2], "x-", label="loss")
    ax[0].set_xscale("log"); ax[0].set_xlabel("simulations per move"); ax[0].set_ylabel("fraction of games")
    ax[0].set_title("MCTS vs a random player"); ax[0].legend()

    ax[1].plot(fork_budgets, tree_acc, "o-", label="MCTS (UCT, c = 1.4)")
    ax[1].plot(fork_budgets, flat_acc, "s--", label="flat Monte Carlo (no tree)")
    ax[1].set_xscale("log"); ax[1].set_ylim(-0.05, 1.05)
    ax[1].set_xlabel("total rollouts / simulations"); ax[1].set_ylabel("fraction choosing a safe edge")
    ax[1].set_title("Fork trap (O to move)"); ax[1].legend()

    ax[2].plot(betas, er, label="E[r] under π*")
    ax[2].plot(betas, kl, label="KL(π* ‖ π0)")
    ax[2].plot(betas, jstar, label="J(π*) = β log Z")
    ax[2].axhline(r.max(), color="gray", lw=0.8, ls=":")
    ax[2].axhline(pi0 @ r, color="gray", lw=0.8, ls="--")
    ax[2].set_xscale("log"); ax[2].set_xlabel("β"); ax[2].set_title("KL-regularized optimum vs β"); ax[2].legend()
    fig.tight_layout()
    path = os.path.join(HERE, "landscape.png")
    fig.savefig(path, dpi=120)
    print(f"MCTS loss rate vs random at {budgets[-1]} sims: {wdl[-1, 2]:.0%}; "
          f"fork trap solved: MCTS {tree_acc[-1]:.0%} vs flat MC {flat_acc[-1]:.0%} at {fork_budgets[-1]}; saved {path}")


if __name__ == "__main__":
    main()