#!/usr/bin/env python3
"""
Can the model mate AT ALL, when the mate is within reach?

This separates three explanations of the conversion failure:

  H-conv-3 CAPABILITY      : the network cannot represent/execute a mating net
  H-conv-2 CREDIT          : it could, but outcome-only returns over ~200-ply
                             games give the mating moves almost no credit
  H-conv-4 EXPLORATION     : it could, but rarely samples mating lines

DESIGN
  Build positions with a KNOWN forced mate in N (N = 1, 2, 3), verified by
  exhaustive search using the rules engine only (no engine evaluation is used
  as a training signal anywhere -- this is measurement, not supervision).

  Then measure, for each N and each recurrent depth:
    * mate-in-1 accuracy: does the policy's top legal move deliver mate?
    * play-out conversion: from the position, does the model actually win?

  A model that scores near 0 on mate-in-1 is capability-limited (H-conv-3).
  A model that solves mate-in-1 but fails play-out conversion from won middlegame
  positions is credit/exploration-limited (H-conv-2 / H-conv-4).

Positions are generated by random self-play-like walks and then filtered with a
depth-limited search over LEGAL MOVES ONLY. The search is used to *label* test
positions, never to train and never to select the model's move.
"""
from __future__ import annotations

import argparse
import json
import os
import random
import sys

import chess
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from tinychess.arena import wilson_interval
from tinychess.batching import encode_batch
from tinychess.checkpoint import load_model_only


def mate_in_one_moves(board: chess.Board):
    """All legal moves that deliver immediate checkmate."""
    out = []
    for mv in board.legal_moves:
        board.push(mv)
        if board.is_checkmate():
            out.append(mv)
        board.pop()
    return out


def forced_mate_in(board: chess.Board, n: int) -> bool:
    """True if side to move has a forced mate in <= n moves (n plies pairs).

    Pure rules-engine minimax over legal moves. Used ONLY to label evaluation
    positions; never used as a training target or to choose the model's move.
    """
    if board.is_checkmate():
        return False
    if n <= 0:
        return False
    for mv in board.legal_moves:
        board.push(mv)
        if board.is_checkmate():
            board.pop()
            return True
        if board.is_stalemate() or board.is_insufficient_material():
            board.pop()
            continue
        # opponent must have no escape
        escaped = False
        for reply in board.legal_moves:
            board.push(reply)
            ok = forced_mate_in(board, n - 1)
            board.pop()
            if not ok:
                escaped = True
                break
        board.pop()
        if not escaped:
            return True
    return False


def gen_mate_positions(n_mate: int, count: int, seed=0, max_tries=400000):
    """Random walks filtered to positions with a forced mate in exactly n_mate."""
    rng = random.Random(seed)
    found, tries = [], 0
    while len(found) < count and tries < max_tries:
        tries += 1
        b = chess.Board()
        for _ in range(rng.randint(10, 70)):
            ms = list(b.legal_moves)
            if not ms:
                break
            b.push(rng.choice(ms))
        if b.is_game_over() or b.legal_moves.count() > 30:
            continue
        if n_mate == 1:
            if mate_in_one_moves(b):
                found.append(b.fen())
        else:
            if mate_in_one_moves(b):
                continue                      # want *exactly* n
            if forced_mate_in(b, n_mate):
                found.append(b.fen())
    return found, tries


def top_move(model, board, steps):
    pb = encode_batch([board])
    with torch.no_grad():
        out = model(**pb.model_args(), steps=steps)
    logits = out["logits"][0]
    idx = int(logits.argmax())
    moves = pb.real_moves[0]
    return moves[idx] if idx < len(moves) else None


def random_control(fens, n, rng):
    """Baseline: pick a uniformly random legal move. Gives the chance rate."""
    hit = 0
    for f in fens:
        b = chess.Board(f)
        mv = rng.choice(list(b.legal_moves))
        if n == 1:
            hit += int(mv in mate_in_one_moves(b))
        else:
            b.push(mv)
            ok = False
            if b.is_checkmate():
                ok = True
            elif not b.is_game_over():
                ok = all(not _escape(b, n) for _ in [0])
            b.pop()
            hit += int(ok)
    return hit


def _escape(b, n):
    for reply in b.legal_moves:
        b.push(reply)
        r = forced_mate_in(b, n - 1)
        b.pop()
        if not r:
            return True
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpts", required=True,
                    help="comma-separated name=path pairs")
    ap.add_argument("--depths", default="1,2,4,8")
    ap.add_argument("--n-positions", type=int, default=60)
    ap.add_argument("--mate-in", default="1,2")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="runs/mate_capability.json")
    a = ap.parse_args()

    depths = [int(x) for x in a.depths.split(",")]
    mates = [int(x) for x in a.mate_in.split(",")]

    suites = {}
    for n in mates:
        fens, tries = gen_mate_positions(n, a.n_positions, seed=a.seed + n)
        suites[n] = fens
        print(f"[suite] mate-in-{n}: {len(fens)} positions (from {tries} tries)")

    rng = random.Random(12345)
    controls = {}
    for n, fens in suites.items():
        if fens:
            h = random_control(fens, n, rng)
            ci = wilson_interval(h, 0, len(fens))
            controls[f"mate_in_{n}"] = {"solved": h, "n": len(fens),
                                        "acc": round(h / len(fens), 4),
                                        "ci95": [round(ci[0], 3), round(ci[1], 3)],
                                        "mean_legal": round(sum(
                                            chess.Board(f).legal_moves.count()
                                            for f in fens) / len(fens), 1)}
            print(f"[control] random-legal mate-in-{n}: {h}/{len(fens)} "
                  f"(mean {controls[f'mate_in_{n}']['mean_legal']} legal moves)")

    res = {"args": vars(a), "suite_sizes": {k: len(v) for k, v in suites.items()},
           "random_control": controls, "models": {}}

    for spec in a.ckpts.split(","):
        name, path = spec.split("=", 1)
        model, cfg = load_model_only(path)
        trained = int(cfg.max_steps)
        if max(depths) > trained and hasattr(model, "core"):
            import torch.nn as nn
            old = model.core.step_emb.data
            new = torch.zeros(max(depths) + 1, old.shape[1])
            new[:old.shape[0]] = old
            new[old.shape[0]:] = old[-1]
            model.core.step_emb = nn.Parameter(new)
            model.cfg.max_steps = max(depths)
            model.core.cfg.max_steps = max(depths)
        entry = {"path": path, "params": model.param_report()["total"],
                 "trained_max_steps": trained, "by_depth": {}}
        print(f"\n=== {name} ({entry['params']:,} params, trained {trained} steps) ===")
        for d in depths:
            row = {}
            for n, fens in suites.items():
                if not fens:
                    continue
                hit = 0
                for f in fens:
                    b = chess.Board(f)
                    mv = top_move(model, b, d)
                    if mv is None:
                        continue
                    if n == 1:
                        if mv in mate_in_one_moves(b):
                            hit += 1
                    else:
                        b.push(mv)
                        # after our move, is a forced mate in n-1 still on?
                        still = False
                        if not b.is_game_over():
                            ok = True
                            for reply in b.legal_moves:
                                b.push(reply)
                                r = forced_mate_in(b, n - 1)
                                b.pop()
                                if not r:
                                    ok = False
                                    break
                            still = ok
                        elif b.is_checkmate():
                            still = True
                        b.pop()
                        hit += int(still)
                ci = wilson_interval(hit, 0, len(fens))
                row[f"mate_in_{n}"] = {"solved": hit, "n": len(fens),
                                       "acc": round(hit / len(fens), 4),
                                       "ci95": [round(ci[0], 3), round(ci[1], 3)]}
                print(f"  depth={d:2d} mate-in-{n}: {hit}/{len(fens)} "
                      f"= {hit/len(fens):.3f} CI[{ci[0]:.2f},{ci[1]:.2f}]")
            entry["by_depth"][str(d)] = row
        res["models"][name] = entry

    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    json.dump(res, open(a.out, "w"), indent=2)
    print(f"\n[saved] {a.out}")


if __name__ == "__main__":
    main()