"""
ladder.py -- the bot "levels", head-to-head matches, Elo estimates, and the
training-games x search-depth "knob grid".

Levels (each is one upgrade on the previous one):
    L0  random mover
    L1  learned material values, 1-ply greedy (no quiescence)
    L2  learned material values, depth-3 alpha-beta + quiescence + TT
    L3  L2 + learned positional features
    (L4, AlphaZero-style network + MCTS, is described in the README, not implemented.)

Commands:
    python ladder.py ladder [--games-per-pair N] [--quick]
    python ladder.py grid   [--games-per-cell N] [--quick]

Games start from short random openings (each opening is played twice with
colours swapped).  Games reaching the ply cap are adjudicated with TEXTBOOK
material values (a lead of >= 3 pawns wins, otherwise draw) -- a referee
rule for the tournament only; the learner never sees these values.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import random
import time
from multiprocessing import Pool
from typing import Dict, List, Optional, Sequence, Tuple

import chess
import chess.pgn

from engine import (STANDARD_VALUES, Evaluator, Searcher, SearchResult,
                    load_weights, standard_weights, weights_from_dict)

HERE = os.path.dirname(os.path.abspath(__file__))
WEIGHTS_DIR = os.path.join(HERE, "weights")


# --------------------------------------------------------------------------
# Players
# --------------------------------------------------------------------------

class RandomPlayer:
    """L0: uniformly random legal move."""

    def __init__(self, name: str = "L0-random", seed: int = 0):
        self.name = name
        self.rng = random.Random(seed)

    def choose(self, board: chess.Board) -> Tuple[chess.Move, Optional[SearchResult]]:
        return self.rng.choice(list(board.legal_moves)), None


class EnginePlayer:
    """Alpha-beta engine with a linear evaluation."""

    def __init__(self, name: str, weights: Sequence[float], depth: int,
                 qdepth: int = 6, use_tt: bool = True,
                 node_limit: Optional[int] = None, time_limit: Optional[float] = None):
        self.name = name
        self.weights = list(weights)
        self.evaluator = Evaluator(self.weights)
        self.searcher = Searcher(self.evaluator, use_tt=use_tt, qdepth=qdepth)
        self.depth = depth
        self.node_limit = node_limit
        self.time_limit = time_limit

    def choose(self, board: chess.Board) -> Tuple[chess.Move, Optional[SearchResult]]:
        r = self.searcher.search(board, self.depth, node_limit=self.node_limit,
                                 time_limit=self.time_limit)
        return r.move, r


def _weights_path(name: str, weights_dir: str) -> str:
    return os.path.join(weights_dir, name)


def make_level(level: str, weights_dir: str = WEIGHTS_DIR, node_limit: int = 20000,
               seed: int = 0):
    """Build a player for level 'L0'..'L3' from the shipped weights."""
    level = level.upper()
    if level == "L0":
        return RandomPlayer("L0-random", seed=seed)
    material = load_weights(_weights_path("material.json", weights_dir))
    if level == "L1":
        return EnginePlayer("L1-material-1ply", material, depth=1, qdepth=0, use_tt=False)
    if level == "L2":
        return EnginePlayer("L2-material-d3", material, depth=3, qdepth=6,
                            node_limit=node_limit)
    if level == "L3":
        positional = load_weights(_weights_path("positional.json", weights_dir))
        return EnginePlayer("L3-positional-d3", positional, depth=3, qdepth=6,
                            node_limit=node_limit)
    raise ValueError(f"unknown level {level!r} (use L0, L1, L2, L3)")


# --------------------------------------------------------------------------
# Playing games
# --------------------------------------------------------------------------

def random_opening(seed: int, plies: int) -> List[chess.Move]:
    """A reproducible random opening of `plies` legal moves (no immediate mates)."""
    rng = random.Random(seed)
    while True:
        b = chess.Board()
        moves = []
        for _ in range(plies):
            m = rng.choice(list(b.legal_moves))
            b.push(m)
            moves.append(m)
            if b.is_game_over():
                break
        if not b.is_game_over():
            return moves


def standard_material(board: chess.Board) -> int:
    """Textbook material balance, White minus Black (referee use only)."""
    total = 0
    for pt, v in STANDARD_VALUES.items():
        total += v * (len(board.pieces(pt, chess.WHITE)) - len(board.pieces(pt, chess.BLACK)))
    return total


def play_game(white, black, opening: Sequence[chess.Move] = (), max_plies: int = 200,
              adjudicate_margin: Optional[int] = 3) -> Tuple[float, chess.pgn.Game, str]:
    """Play one game.  Returns (score for White in {1, 0.5, 0}, PGN, reason)."""
    board = chess.Board()
    for m in opening:
        board.push(m)
    reason = ""
    while True:
        outcome = board.outcome(claim_draw=False)
        if outcome is not None:
            reason = outcome.termination.name.lower()
            score = {True: 1.0, False: 0.0, None: 0.5}[outcome.winner]
            break
        if board.is_repetition(3):
            reason, score = "threefold repetition", 0.5
            break
        if board.ply() >= max_plies:
            bal = standard_material(board)
            if adjudicate_margin is not None and abs(bal) >= adjudicate_margin:
                reason, score = f"adjudicated (material {bal:+d})", (1.0 if bal > 0 else 0.0)
            else:
                reason, score = "move cap", 0.5
            break
        player = white if board.turn == chess.WHITE else black
        move, _ = player.choose(board)
        board.push(move)

    game = chess.pgn.Game.from_board(board)
    game.headers["White"] = white.name
    game.headers["Black"] = black.name
    game.headers["Result"] = {1.0: "1-0", 0.0: "0-1", 0.5: "1/2-1/2"}[score]
    game.headers["Termination"] = reason
    return score, game, reason


# --------------------------------------------------------------------------
# Elo estimation
# --------------------------------------------------------------------------

def estimate_elo(names: List[str], games: List[Tuple[str, str, float]],
                 anchor: str, prior_draws: float = 1.0, iters: int = 5000) -> Dict[str, float]:
    """Maximum-likelihood Elo (Bradley-Terry / logistic, base 10, scale 400),
    treating a draw as half a win and half a loss.

    `games` is a list of (player_a, player_b, score_of_a).  To keep ratings
    finite when one bot wins every game, each pair that met gets
    `prior_draws` virtual drawn games (a mild shrinkage prior).  Ratings are
    shifted so that `anchor` has Elo 0."""
    idx = {n: i for i, n in enumerate(names)}
    n = len(names)
    wins = [[0.0] * n for _ in range(n)]   # wins[i][j] = points of i vs j
    cnt = [[0.0] * n for _ in range(n)]
    for a, b, s in games:
        i, j = idx[a], idx[b]
        wins[i][j] += s
        wins[j][i] += 1.0 - s
        cnt[i][j] += 1
        cnt[j][i] += 1
    for i in range(n):
        for j in range(n):
            if i != j and cnt[i][j] > 0:
                wins[i][j] += prior_draws * 0.5
                cnt[i][j] += prior_draws
    # Minorise-maximise iterations for Bradley-Terry strengths gamma_i
    # (Hunter 2004):  gamma_i <- W_i / sum_j n_ij / (gamma_i + gamma_j).
    gamma = [1.0] * n
    for _ in range(iters):
        new = []
        for i in range(n):
            w_i = sum(wins[i])
            denom = sum(cnt[i][j] / (gamma[i] + gamma[j]) for j in range(n) if j != i)
            new.append(w_i / denom if denom > 0 else gamma[i])
        g = math.exp(sum(math.log(x) for x in new) / n)     # fix the scale
        new = [x / g for x in new]
        if max(abs(a - b) for a, b in zip(new, gamma)) < 1e-10:
            gamma = new
            break
        gamma = new
    r = [400.0 * math.log10(x) for x in gamma]
    a = r[idx[anchor]]
    return {nm: r[idx[nm]] - a for nm in names}


# --------------------------------------------------------------------------
# Parallel match runner
# --------------------------------------------------------------------------

def _play_job(job):
    """Worker: build players from specs (picklable) and play one game."""
    (wspec, bspec, opening_seed, opening_plies, max_plies, game_seed) = job
    white, black = build_player(wspec, game_seed), build_player(bspec, game_seed + 1)
    opening = random_opening(opening_seed, opening_plies)
    t0 = time.time()
    score, game, reason = play_game(white, black, opening, max_plies)
    return wspec["name"], bspec["name"], score, str(game), reason, time.time() - t0


def build_player(spec: Dict, seed: int = 0):
    """spec: {"name", "kind": "random"|"engine", "weights", "depth", "qdepth", "use_tt", "node_limit"}"""
    if spec["kind"] == "random":
        return RandomPlayer(spec["name"], seed=seed)
    return EnginePlayer(spec["name"], spec["weights"], spec["depth"], spec.get("qdepth", 6),
                        spec.get("use_tt", True), spec.get("node_limit"))


def level_specs(weights_dir: str, node_limit: int) -> List[Dict]:
    material = load_weights(os.path.join(weights_dir, "material.json"))
    positional = load_weights(os.path.join(weights_dir, "positional.json"))
    return [
        {"name": "L0-random", "kind": "random"},
        {"name": "L1-material-1ply", "kind": "engine", "weights": material, "depth": 1,
         "qdepth": 0, "use_tt": False},
        {"name": "L2-material-d3", "kind": "engine", "weights": material, "depth": 3,
         "qdepth": 6, "node_limit": node_limit},
        {"name": "L3-positional-d3", "kind": "engine", "weights": positional, "depth": 3,
         "qdepth": 6, "node_limit": node_limit},
    ]


def run_matches(pairs: List[Tuple[Dict, Dict]], openings: int, opening_plies: int,
                max_plies: int, workers: int, seed: int = 0, verbose: bool = True):
    """For each (a, b) pair play `openings` openings x 2 colours."""
    jobs = []
    for pi, (a, b) in enumerate(pairs):
        for o in range(openings):
            oseed = seed * 100003 + o            # same openings for every pair
            jobs.append((a, b, oseed, opening_plies, max_plies, 7919 * (pi * 1000 + 2 * o)))
            jobs.append((b, a, oseed, opening_plies, max_plies, 7919 * (pi * 1000 + 2 * o + 1)))
    out = []
    t0 = time.time()
    if workers > 1:
        with Pool(workers) as pool:
            for k, res in enumerate(pool.imap_unordered(_play_job, jobs)):
                out.append(res)
                if verbose and (k + 1) % max(1, len(jobs) // 10) == 0:
                    print(f"  {k + 1}/{len(jobs)} games ({time.time() - t0:.0f}s)", flush=True)
    else:
        for k, j in enumerate(jobs):
            out.append(_play_job(j))
            if verbose and (k + 1) % max(1, len(jobs) // 10) == 0:
                print(f"  {k + 1}/{len(jobs)} games ({time.time() - t0:.0f}s)", flush=True)
    return out


# --------------------------------------------------------------------------
# Experiments
# --------------------------------------------------------------------------

def run_ladder(args) -> Dict:
    specs = level_specs(args.weights, args.node_limit)
    names = [s["name"] for s in specs]
    pairs = [(specs[i], specs[j]) for i in range(len(specs)) for j in range(i + 1, len(specs))]
    openings = max(1, args.games_per_pair // 2)
    print(f"round robin: {len(pairs)} pairs x {2 * openings} games "
          f"(random {args.opening_plies}-ply openings, colours alternated)")
    t0 = time.time()
    res = run_matches(pairs, openings, args.opening_plies, args.max_plies, args.workers, args.seed)
    elapsed = time.time() - t0

    games = [(w, b, s) for (w, b, s, _, _, _) in res]
    table = {a: {b: [0.0, 0] for b in names} for a in names}
    for w, b, s in games:
        table[w][b][0] += s
        table[w][b][1] += 1
        table[b][w][0] += 1 - s
        table[b][w][1] += 1
    elo = estimate_elo(names, games, anchor=names[0])

    print(f"\nscore table (row player's points / games), {elapsed:.0f}s total")
    print(" " * 18 + "".join(f"{n[:16]:>18}" for n in names))
    for a in names:
        cells = "".join(f"{(f'{table[a][b][0]:g}/{table[a][b][1]}' if table[a][b][1] else '-'):>18}"
                        for b in names)
        print(f"{a[:18]:<18}{cells}")
    print("\nElo estimate (L0-random = 0; 1 virtual draw per pair as prior):")
    for n in names:
        print(f"  {n:<20} {elo[n]:7.0f}")
    reasons: Dict[str, int] = {}
    for r in res:
        key = r[4].split(" (")[0]
        reasons[key] = reasons.get(key, 0) + 1
    print("terminations:", reasons)

    if args.pgn:
        with open(args.pgn, "w") as fh:
            for r in res:
                fh.write(r[3] + "\n\n")
        print(f"games written to {args.pgn}")
    out = {"elo": elo, "table": table, "terminations": reasons, "seconds": elapsed,
           "games_per_pair": 2 * openings, "max_plies": args.max_plies,
           "node_limit": args.node_limit}
    with open(os.path.join(args.outdir, "ladder_results.json"), "w") as fh:
        json.dump(out, fh, indent=1)
    plot_ladder(elo, os.path.join(args.outdir, "ladder_elo.png"))
    return out


def plot_ladder(elo: Dict[str, float], path: str) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    names = list(elo)
    fig, ax = plt.subplots(figsize=(7, 3.6))
    ax.bar(range(len(names)), [elo[n] for n in names], color="#4C72B0")
    ax.set_xticks(range(len(names)), names, rotation=15)
    ax.set_ylabel("Elo (relative to L0)")
    ax.set_title("Self-play ladder")
    for i, n in enumerate(names):
        ax.text(i, elo[n], f"{elo[n]:.0f}", ha="center", va="bottom")
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)


def run_grid(args) -> Dict:
    """Training checkpoints x search depth -> score vs a fixed reference bot."""
    # Rows are checkpoints along the whole training pipeline:
    #   "m<g>" = stage-1 (material) checkpoint after g games,
    #   "p<g>" = stage-2 (positional) checkpoint after g more games.
    with open(os.path.join(args.weights, "material_checkpoints.json")) as fh:
        mat = {c["games"]: c["weights"] for c in json.load(fh)["checkpoints"]}
    pos = {}
    pos_path = os.path.join(args.weights, "positional_checkpoints.json")
    if os.path.exists(pos_path):
        with open(pos_path) as fh:
            pos = {c["games"]: c["weights"] for c in json.load(fh)["checkpoints"]}
    n_mat = max(mat)
    rows = []          # (label, total self-play games, weights)
    for tok in args.checkpoints.split(","):
        tok = tok.strip()
        stage, g = tok[0], int(tok[1:])
        src = mat if stage == "m" else pos
        if g not in src:
            print(f"skipping {tok}: not among saved checkpoints {sorted(src)}")
            continue
        total = g if stage == "m" else n_mat + g
        label = f"{total} (material)" if stage == "m" else f"{n_mat}+{g} (positional)"
        rows.append((label, total, src[g]))
    if not rows:
        raise SystemExit("no valid checkpoints selected")
    games_axis = [r[0] for r in rows]
    depths = [int(x) for x in args.depths.split(",")]
    ref = {"name": "ref-textbook-d1", "kind": "engine", "weights": standard_weights(),
           "depth": 1, "qdepth": 6, "use_tt": False}
    pairs, keys = [], []
    for g, (_, total, wdict) in zip(games_axis, rows):
        for d in depths:
            spec = {"name": f"{g}-d{d}", "kind": "engine",
                    "weights": weights_from_dict(wdict), "depth": d, "qdepth": 6,
                    "node_limit": args.node_limit}
            pairs.append((spec, ref))
            keys.append((g, d))
    openings = max(1, args.games_per_cell // 2)
    print(f"knob grid: {len(games_axis)} checkpoints x {len(depths)} depths, "
          f"{2 * openings} games per cell vs {ref['name']}")
    t0 = time.time()
    res = run_matches(pairs, openings, args.opening_plies, args.max_plies, args.workers, args.seed)
    elapsed = time.time() - t0
    score = {k: [0.0, 0] for k in keys}
    name_to_key = {p[0]["name"]: k for p, k in zip(pairs, keys)}
    for w, b, s, _, _, _ in res:
        if w in name_to_key:
            k, pts = name_to_key[w], s
        else:
            k, pts = name_to_key[b], 1 - s
        score[k][0] += pts
        score[k][1] += 1
    grid = [[score[(g, d)][0] / score[(g, d)][1] for d in depths] for g in games_axis]

    print(f"\nscore vs {ref['name']} (rows: training games, cols: depth), {elapsed:.0f}s")
    print(f"{'self-play games':<24}" + "".join(f"{'d' + str(d):>7}" for d in depths))
    for g, row in zip(games_axis, grid):
        print(f"{g:<24}" + "".join(f"{v:7.2f}" for v in row))
    out = {"games_axis": games_axis, "depths": depths, "score": grid,
           "reference": ref["name"], "games_per_cell": 2 * openings, "seconds": elapsed}
    with open(os.path.join(args.outdir, "grid_results.json"), "w") as fh:
        json.dump(out, fh, indent=1)
    plot_grid(out, os.path.join(args.outdir, "knob_grid.png"))
    return out


def plot_grid(out: Dict, path: str) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np
    grid = np.array(out["score"])
    fig, ax = plt.subplots(figsize=(6.5, 4.8))
    im = ax.imshow(grid, cmap="RdYlGn", vmin=0, vmax=1, origin="lower", aspect="auto")
    ax.set_xticks(range(len(out["depths"])), [str(d) for d in out["depths"]])
    ax.set_yticks(range(len(out["games_axis"])), [str(g) for g in out["games_axis"]])
    ax.set_xlabel("search depth (test-time compute)")
    ax.set_ylabel("self-play games (training)")
    for i in range(grid.shape[0]):
        for j in range(grid.shape[1]):
            ax.text(j, i, f"{grid[i, j]:.2f}", ha="center", va="center")
    ax.set_title(f"score vs {out['reference']}\n({out['games_per_cell']} games/cell)")
    fig.colorbar(im, ax=ax)
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("experiment", choices=["ladder", "grid"])
    ap.add_argument("--quick", action="store_true", help="fewer games, shorter cap")
    ap.add_argument("--games-per-pair", type=int, default=8)
    ap.add_argument("--games-per-cell", type=int, default=8)
    ap.add_argument("--checkpoints", default="m0,m20,m100,m1000,p20,p100,p300",
                    help="rows: m<g> = material checkpoint after g games, "
                         "p<g> = positional checkpoint after g stage-2 games")
    ap.add_argument("--depths", default="1,2,3")
    ap.add_argument("--max-plies", type=int, default=200)
    ap.add_argument("--opening-plies", type=int, default=4)
    ap.add_argument("--node-limit", type=int, default=20000, help="per move, depth>=2 bots")
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--weights", default=WEIGHTS_DIR)
    ap.add_argument("--outdir", default=os.path.join(HERE, "results"))
    ap.add_argument("--pgn", default=None, help="write all ladder games to this PGN file")
    args = ap.parse_args(argv)
    args.workers = max(1, min(2, args.workers))
    if args.quick:
        if args.outdir == os.path.join(HERE, "results"):
            args.outdir = os.path.join(HERE, "runs", "quick")   # keep shipped results intact
        args.games_per_pair = 2
        args.games_per_cell = 2
        args.max_plies = 120
        args.node_limit = 5000
        args.checkpoints = "m0,m1000,p300"
        args.depths = "1,3"
    os.makedirs(args.outdir, exist_ok=True)
    if args.experiment == "ladder":
        run_ladder(args)
    else:
        run_grid(args)


if __name__ == "__main__":
    main()