"""
learn.py -- learn evaluation weights purely from self-play with TD-Leaf(lambda).

Stage 1 ("material"):  all five piece values start EQUAL (1.0) and are learned.
Stage 2 ("positional"): the 12 positional weights start at 0 and are learned,
                        with the stage-1 material values frozen (or jointly,
                        with --joint).

The algorithm (one self-play game, then one update):

  * At every ply the side to move scores each legal move with a shallow
    search (default: 1 ply + quiescence).  It plays the best one, except with
    probability epsilon it plays a uniformly random move (exploration).  The
    first few plies of each game are random, so games differ.
  * For the move actually played we keep the *leaf* of its principal
    variation (the quiet position where the static evaluation was applied)
    and that leaf's evaluation v_t = w . f(leaf_t)   (White's view).
    This is the "Leaf" in TD-Leaf: learn the evaluation of the position the
    search actually trusted, not of the root.
  * Values are squashed to a win probability  J_t = sigmoid(beta * v_t).
    The final "value" J_N is the game result (1 White win, 0.5 draw, 0 loss).
  * Temporal differences d_t = J_{t+1} - J_t and the TD(lambda) update
        w <- w + alpha * sum_t  grad_w J_t * sum_{j>=t} lambda^(j-t) d_j
    with grad_w J_t = beta * J_t (1 - J_t) * f(leaf_t).

Because the evaluation is used by both sides (features are White - Black),
each game teaches both colours at once.

Games are capped (default 300 plies) and adjudicated as a DRAW at the cap
(no material-based adjudication in training: that would leak the textbook
piece values into the learner).

Usage:
    python learn.py material   [--games N] [--quick]
    python learn.py positional [--games N] [--quick] [--joint]
"""
from __future__ import annotations

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

import chess

from engine import (ALL_FEATURES, MATERIAL_FEATURES, N_FEATURES, N_MATERIAL,
                    Evaluator, Searcher, equal_material_weights, features,
                    is_mate_score, leaf_board, load_weights,
                    normalized_piece_values, plot_heatmaps, save_weights,
                    weights_to_dict)

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


@dataclass
class TrainConfig:
    games: int = 1000            # number of self-play games
    depth: int = 1               # search depth during self-play (plus quiescence)
    qdepth: int = 4              # quiescence depth cap during self-play
    epsilon: float = 0.10        # probability of a random move
    random_opening_max: int = 8  # first k ~ U{0..max} plies are random
    max_plies: int = 300         # game is a draw after this many plies
    alpha: float = 0.05          # initial learning rate
    alpha_decay_games: float = 300.0  # alpha_g = alpha / sqrt(1 + g / alpha_decay_games)
    average_last: float = 0.25   # report the average of the last 25% of checkpoints
    lam: float = 0.7             # TD lambda
    beta: float = 0.5            # sigmoid slope:  J = sigmoid(beta * eval)
    workers: int = 2             # parallel self-play processes (<= 2 on this box)
    checkpoint_every: int = 20   # save weights every this many games
    seed: int = 0
    start_fen: Optional[str] = None   # start every game here instead of the initial position


def sigmoid(x: float) -> float:
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    z = math.exp(x)
    return z / (1.0 + z)


# --------------------------------------------------------------------------
# One self-play game
# --------------------------------------------------------------------------

def play_selfplay_game(weights: Sequence[float], cfg: TrainConfig, seed: int,
                       positional: bool) -> Dict:
    """Play one game of the current evaluation against itself.

    Returns a dict with, for every recorded (non-random-opening) ply,
    the leaf evaluation score (White's view) and the leaf feature vector
    (None when the score is not a plain function of w, e.g. a mate or a
    repetition draw found by the search), plus the final result."""
    rng = random.Random(seed)
    ev = Evaluator(weights, positional=positional)
    searcher = Searcher(ev, use_tt=False, qdepth=cfg.qdepth)   # no TT: exact PVs
    board = chess.Board(cfg.start_fen) if cfg.start_fen else chess.Board()
    n_random = board.ply() + rng.randint(0, cfg.random_opening_max)
    scores: List[float] = []
    leaf_feats: List[Optional[List[float]]] = []
    termination = "move cap"

    while True:
        outcome = board.outcome(claim_draw=False)
        if outcome is not None:
            termination = outcome.termination.name.lower()
            break
        if board.is_repetition(3):
            termination = "threefold repetition"
            break
        if board.ply() >= cfg.max_plies:
            break

        if board.ply() < n_random:
            board.push(rng.choice(list(board.legal_moves)))
            continue

        scored = searcher.score_root_moves(board, cfg.depth)
        if rng.random() < cfg.epsilon:
            move, score, pv = rng.choice(scored)
        else:
            best = max(s for _, s, _ in scored)
            # break ties randomly so equal-valued moves don't always pick the same one
            move, score, pv = rng.choice([t for t in scored if t[1] == best])

        white_score = score if board.turn == chess.WHITE else -score
        f = None
        if not is_mate_score(score):
            leaf = leaf_board(board, pv)
            f_leaf = features(leaf, positional=positional)
            v = sum(a * b for a, b in zip(ev.w, f_leaf))
            if abs(v - white_score) < 1e-6:     # the score really is w . f(leaf)
                f = f_leaf
        scores.append(white_score)
        leaf_feats.append(f)
        board.push(move)

    if board.is_checkmate():
        result = 0.0 if board.turn == chess.WHITE else 1.0
    else:
        result = 0.5
    return {"scores": scores, "feats": leaf_feats, "result": result,
            "plies": board.ply(), "termination": termination}


# --------------------------------------------------------------------------
# TD-Leaf(lambda) update
# --------------------------------------------------------------------------

def td_leaf_update(weights: Sequence[float], game: Dict, cfg: TrainConfig,
                   learn_mask: Sequence[bool], alpha: Optional[float] = None) -> List[float]:
    """Return the weights after one TD(lambda) update from one game.

    J_t = sigmoid(beta * score_t)  (mate scores give J = 0 or 1),
    J_N = game result.  Only weights with learn_mask[i] True change.
    `alpha` overrides cfg.alpha (used for the learning-rate schedule)."""
    alpha = cfg.alpha if alpha is None else alpha
    scores, feats = game["scores"], game["feats"]
    n = len(scores)
    if n == 0:
        return list(weights)
    J = []
    for s in scores:
        if is_mate_score(s):
            J.append(1.0 if s > 0 else 0.0)
        else:
            J.append(sigmoid(cfg.beta * s))
    J.append(game["result"])
    d = [J[t + 1] - J[t] for t in range(n)]

    delta = [0.0] * N_FEATURES
    trace_err = 0.0                       # sum_{j>=t} lambda^(j-t) d_j, built backwards
    for t in range(n - 1, -1, -1):
        trace_err = d[t] + cfg.lam * trace_err
        f = feats[t]
        if f is None:
            continue
        g = cfg.beta * J[t] * (1.0 - J[t])        # dJ/dv
        for i in range(N_FEATURES):
            if learn_mask[i] and f[i] != 0.0:
                delta[i] += g * f[i] * trace_err
    return [w + (alpha * dl if learn_mask[i] else 0.0)
            for i, (w, dl) in enumerate(zip(weights, delta))]


# --------------------------------------------------------------------------
# Training loop
# --------------------------------------------------------------------------

def _worker(args):
    weights, cfg, seed, positional = args
    return play_selfplay_game(weights, cfg, seed, positional)


def train(init_weights: Sequence[float], cfg: TrainConfig, learn_mask: Sequence[bool],
          positional: bool, log_every: int = 20, verbose: bool = True) -> Dict:
    """Self-play training.  Games are played in small batches (one per
    worker) with the same weights; updates are applied game by game after
    each batch.  Returns {"final": weights, "checkpoints": [...], "stats": ...}."""
    w = list(init_weights)
    checkpoints = [{"games": 0, "weights": weights_to_dict(w)}]
    results = {"1-0": 0, "0-1": 0, "1/2": 0}
    terminations: Dict[str, int] = {}
    total_plies = 0
    t0 = time.time()
    pool = Pool(cfg.workers) if cfg.workers > 1 else None
    try:
        g = 0
        while g < cfg.games:
            batch = min(max(1, cfg.workers), cfg.games - g)
            jobs = [(w, cfg, cfg.seed * 1_000_003 + g + k, positional) for k in range(batch)]
            games = pool.map(_worker, jobs) if pool else [_worker(j) for j in jobs]
            for game in games:
                alpha_g = cfg.alpha / math.sqrt(1.0 + g / cfg.alpha_decay_games)
                w = td_leaf_update(w, game, cfg, learn_mask, alpha_g)
                g += 1
                total_plies += game["plies"]
                key = {1.0: "1-0", 0.0: "0-1"}.get(game["result"], "1/2")
                results[key] += 1
                terminations[game["termination"]] = terminations.get(game["termination"], 0) + 1
                if g % cfg.checkpoint_every == 0 or g == cfg.games:
                    checkpoints.append({"games": g, "weights": weights_to_dict(w)})
                if verbose and (g % log_every == 0 or g == cfg.games):
                    nv = normalized_piece_values(w)
                    raw = " ".join(f"{n[0].upper()}={w[i]:.3f}" for i, n in enumerate(MATERIAL_FEATURES))
                    norm = " ".join(f"{n[0].upper()}={nv[n]:.2f}" for n in MATERIAL_FEATURES)
                    print(f"[{g:5d} games, {time.time() - t0:6.0f}s] raw: {raw} | pawn=1: {norm} "
                          f"| W/B/D {results['1-0']}/{results['0-1']}/{results['1/2']}",
                          flush=True)
                    if positional:
                        pos = " ".join(f"{ALL_FEATURES[i]}={w[i] / w[0]:+.2f}"
                                       for i in range(N_MATERIAL, N_FEATURES))
                        print(f"         positional (pawns): {pos}", flush=True)
    finally:
        if pool:
            pool.close()
            pool.join()
    stats = {"results": results, "terminations": terminations,
             "mean_plies": total_plies / max(1, cfg.games),
             "seconds": time.time() - t0}
    return {"final": w, "averaged": average_checkpoints(checkpoints, cfg.average_last),
            "checkpoints": checkpoints, "stats": stats}


def average_checkpoints(checkpoints: List[Dict], frac: float) -> List[float]:
    """Polyak-style average of the last `frac` of the checkpoints.  Self-play
    TD with a constant-ish step size keeps jittering around its fixed point;
    averaging the tail gives a much less noisy estimate."""
    tail = checkpoints[1:] if len(checkpoints) > 1 else checkpoints
    k = max(1, int(round(len(tail) * frac)))
    tail = tail[-k:]
    return [sum(c["weights"][n] for c in tail) / len(tail) for n in ALL_FEATURES]


# --------------------------------------------------------------------------
# Plots
# --------------------------------------------------------------------------

def plot_piece_values(checkpoints: List[Dict], path: str, title: str = "") -> None:
    """Learned piece values (pawn = 1) versus number of self-play games."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    xs = [c["games"] for c in checkpoints]
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
    styles = {"pawn": "-", "knight": "-", "bishop": "--", "rook": "-", "queen": "-"}
    for name in MATERIAL_FEATURES:
        raw = [c["weights"][name] for c in checkpoints]
        norm = [c["weights"][name] / c["weights"]["pawn"] if c["weights"]["pawn"] > 1e-9
                else float("nan") for c in checkpoints]
        axes[0].plot(xs, raw, styles[name], label=name)
        axes[1].plot(xs, norm, styles[name], label=name)
    for ref in (1, 3, 5, 9):
        axes[1].axhline(ref, color="grey", lw=0.5, ls=":")
    axes[0].set_title("raw learned weights")
    axes[1].set_title("normalised (pawn = 1); dotted = 1/3/5/9")
    for ax in axes:
        ax.set_xlabel("self-play games")
        ax.legend()
    fig.suptitle(title or "Piece values learned from self-play (all started equal)")
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)


def plot_positional(checkpoints: List[Dict], path: str) -> None:
    """Positional weights (in pawns) versus games."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    xs = [c["games"] for c in checkpoints]
    names = ALL_FEATURES[N_MATERIAL:]
    fig, ax = plt.subplots(figsize=(8, 5))
    for name in names:
        ys = [c["weights"][name] / c["weights"]["pawn"] for c in checkpoints]
        ax.plot(xs, ys, label=name)
    ax.axhline(0, color="grey", lw=0.5)
    ax.set_xlabel("self-play games")
    ax.set_ylabel("weight / pawn weight")
    ax.set_title("Positional weights learned from self-play")
    ax.legend(fontsize=7, ncol=2)
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)


# --------------------------------------------------------------------------
# CLI
# --------------------------------------------------------------------------

def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("stage", choices=["material", "positional"])
    ap.add_argument("--games", type=int, default=None)
    ap.add_argument("--quick", action="store_true", help="short run (~1 min) for a smoke test")
    ap.add_argument("--joint", action="store_true",
                    help="stage 2: also keep learning the material weights")
    ap.add_argument("--alpha", type=float, default=None)
    ap.add_argument("--lam", type=float, default=None)
    ap.add_argument("--epsilon", type=float, default=None)
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default=WEIGHTS_DIR,
                    help="output directory (default: weights/, or runs/quick/ with --quick)")
    ap.add_argument("--init", default=None,
                    help="stage 2: material weights JSON to start from "
                         "(default: <out>/material.json)")
    args = ap.parse_args(argv)
    if args.quick and args.out == WEIGHTS_DIR:
        # never overwrite the shipped weights with a smoke-test run
        args.out = os.path.join(HERE, "runs", "quick")
    os.makedirs(args.out, exist_ok=True)

    cfg = TrainConfig(workers=max(1, min(2, args.workers)), seed=args.seed)
    if args.stage == "positional":
        cfg.games = 300
        cfg.alpha = 0.02
        cfg.checkpoint_every = 20
    if args.quick:
        cfg.games = 120 if args.stage == "material" else 30
        cfg.max_plies = 200
        cfg.checkpoint_every = 10
    for k in ("games", "alpha", "lam", "epsilon"):
        if getattr(args, k) is not None:
            setattr(cfg, k, getattr(args, k))
    tag = "_quick" if args.quick else ""

    if args.stage == "material":
        init = equal_material_weights(1.0)
        mask = [i < N_MATERIAL for i in range(N_FEATURES)]
        positional = False
    else:
        init_path = args.init or os.path.join(args.out, f"material{tag}.json")
        if not os.path.exists(init_path):
            init_path = os.path.join(WEIGHTS_DIR, "material.json")
        init = load_weights(init_path)
        init = init[:N_MATERIAL] + [0.0] * (N_FEATURES - N_MATERIAL)
        mask = [args.joint or i >= N_MATERIAL for i in range(N_FEATURES)]
        positional = True
        print(f"stage 2 starting from material weights in {init_path}")

    print(f"config: {asdict(cfg)}")
    out = train(init, cfg, mask, positional)
    final = out["averaged"]          # tail-averaged weights are what we ship
    name = args.stage + tag
    meta = {"stage": args.stage, "config": asdict(cfg), "stats": out["stats"],
            "normalized_piece_values": normalized_piece_values(final),
            "note": f"weights = average of the last {cfg.average_last:.0%} of checkpoints",
            "last_iterate": weights_to_dict(out["final"])}
    save_weights(os.path.join(args.out, f"{name}.json"), final, meta)
    with open(os.path.join(args.out, f"{name}_checkpoints.json"), "w") as fh:
        json.dump({"config": asdict(cfg), "checkpoints": out["checkpoints"]}, fh, indent=1)

    if args.stage == "material":
        plot_piece_values(out["checkpoints"], os.path.join(args.out, f"{name}_piece_values.png"))
    else:
        plot_positional(out["checkpoints"], os.path.join(args.out, f"{name}_weights.png"))
    plot_heatmaps(final, os.path.join(args.out, f"{name}_heatmaps.png"),
                  title=f"Implied piece-square values after stage '{args.stage}'")

    print("\nstats:", json.dumps(out["stats"]))
    print("final normalised piece values (pawn = 1):",
          {k: round(v, 2) for k, v in normalized_piece_values(final).items()})
    if args.stage == "positional":
        print("positional weights (in pawns):",
              {n: round(final[i] / final[0], 3) for i, n in enumerate(ALL_FEATURES) if i >= N_MATERIAL})
    print(f"saved to {args.out}/{name}.json (+ checkpoints and plots)")


if __name__ == "__main__":
    main()