app.py

"""
ChessSLM — Play against FlameF0X/ChessSLM
Hugging Face Space  |  Gradio app

Move selection: ChessSLM generates 10–20 candidate moves via sampling,
a fast PST-based evaluator scores each and selects the best one.
"""

import re
import random
import chess
import chess.svg
import chess.pgn
import gradio as gr
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# ──────────────────────────────────────────────────────────────────────────────
# Model loading (once at startup)
# ──────────────────────────────────────────────────────────────────────────────

MODEL_ID = "FlameF0X/ChessSLM-RL"

print(f"Loading {MODEL_ID}...")
device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = GPT2Tokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
model     = GPT2LMHeadModel.from_pretrained(MODEL_ID)
model.to(device)
model.eval()
model.config.use_cache = True
print(f"✓ Model ready on {device}")

# ──────────────────────────────────────────────────────────────────────────────
# Fast PST-based position evaluator (replaces Stockfish)
# ──────────────────────────────────────────────────────────────────────────────

# Centipawn piece values
_PIECE_VALUES = {
    chess.PAWN:   100,
    chess.KNIGHT: 320,
    chess.BISHOP: 330,
    chess.ROOK:   500,
    chess.QUEEN:  900,
    chess.KING:   20000,
}

# Piece-square tables — index 0 = a1, from White's point of view.
# For Black pieces the table is mirrored vertically.
_PST: dict[int, list[int]] = {
    chess.PAWN: [
         0,  0,  0,  0,  0,  0,  0,  0,
        50, 50, 50, 50, 50, 50, 50, 50,
        10, 10, 20, 30, 30, 20, 10, 10,
         5,  5, 10, 25, 25, 10,  5,  5,
         0,  0,  0, 20, 20,  0,  0,  0,
         5, -5,-10,  0,  0,-10, -5,  5,
         5, 10, 10,-20,-20, 10, 10,  5,
         0,  0,  0,  0,  0,  0,  0,  0,
    ],
    chess.KNIGHT: [
        -50,-40,-30,-30,-30,-30,-40,-50,
        -40,-20,  0,  0,  0,  0,-20,-40,
        -30,  0, 10, 15, 15, 10,  0,-30,
        -30,  5, 15, 20, 20, 15,  5,-30,
        -30,  0, 15, 20, 20, 15,  0,-30,
        -30,  5, 10, 15, 15, 10,  5,-30,
        -40,-20,  0,  5,  5,  0,-20,-40,
        -50,-40,-30,-30,-30,-30,-40,-50,
    ],
    chess.BISHOP: [
        -20,-10,-10,-10,-10,-10,-10,-20,
        -10,  0,  0,  0,  0,  0,  0,-10,
        -10,  0,  5, 10, 10,  5,  0,-10,
        -10,  5,  5, 10, 10,  5,  5,-10,
        -10,  0, 10, 10, 10, 10,  0,-10,
        -10, 10, 10, 10, 10, 10, 10,-10,
        -10,  5,  0,  0,  0,  0,  5,-10,
        -20,-10,-10,-10,-10,-10,-10,-20,
    ],
    chess.ROOK: [
         0,  0,  0,  0,  0,  0,  0,  0,
         5, 10, 10, 10, 10, 10, 10,  5,
        -5,  0,  0,  0,  0,  0,  0, -5,
        -5,  0,  0,  0,  0,  0,  0, -5,
        -5,  0,  0,  0,  0,  0,  0, -5,
        -5,  0,  0,  0,  0,  0,  0, -5,
        -5,  0,  0,  0,  0,  0,  0, -5,
         0,  0,  0,  5,  5,  0,  0,  0,
    ],
    chess.QUEEN: [
        -20,-10,-10, -5, -5,-10,-10,-20,
        -10,  0,  0,  0,  0,  0,  0,-10,
        -10,  0,  5,  5,  5,  5,  0,-10,
         -5,  0,  5,  5,  5,  5,  0, -5,
          0,  0,  5,  5,  5,  5,  0, -5,
        -10,  5,  5,  5,  5,  5,  0,-10,
        -10,  0,  5,  0,  0,  0,  0,-10,
        -20,-10,-10, -5, -5,-10,-10,-20,
    ],
    chess.KING: [
        -30,-40,-40,-50,-50,-40,-40,-30,
        -30,-40,-40,-50,-50,-40,-40,-30,
        -30,-40,-40,-50,-50,-40,-40,-30,
        -30,-40,-40,-50,-50,-40,-40,-30,
        -20,-30,-30,-40,-40,-30,-30,-20,
        -10,-20,-20,-20,-20,-20,-20,-10,
         20, 20,  0,  0,  0,  0, 20, 20,
         20, 30, 10,  0,  0, 10, 30, 20,
    ],
}


def _pst_index(sq: int, color: chess.Color) -> int:
    """Map a board square to a PST row index (always from White's orientation)."""
    rank, file = divmod(sq, 8)
    if color == chess.WHITE:
        return rank * 8 + file
    else:
        return (7 - rank) * 8 + file


def _evaluate_position(board: chess.Board) -> int:
    """
    Return a centipawn score from White's point of view.
    Combines material, piece-square tables, mobility, bishop pair,
    and basic pawn-structure penalties.
    """
    if board.is_checkmate():
        # The side that just moved delivered checkmate → opponent (to move) lost
        return -20000 if board.turn == chess.WHITE else 20000
    if board.is_stalemate() or board.is_insufficient_material():
        return 0

    score = 0

    # --- Material + PST ---
    for sq in chess.SQUARES:
        piece = board.piece_at(sq)
        if piece is None:
            continue
        val = _PIECE_VALUES[piece.piece_type] + _PST[piece.piece_type][_pst_index(sq, piece.color)]
        score += val if piece.color == chess.WHITE else -val

    # --- Mobility (legal moves count is a cheap activity proxy) ---
    # We already know it's not game-over, so legal_moves is non-empty.
    own_mobility  = board.legal_moves.count()
    board.push(chess.Move.null())           # pass turn
    opp_mobility  = board.legal_moves.count()
    board.pop()
    score += (own_mobility - opp_mobility) * 5  # 5 cp per extra legal move
    if board.turn == chess.BLACK:
        score = -score   # flip so it's always from White's POV after this
        score = -score   # undo the flip (we do it below uniformly)

    # Re-derive cleanly — mobility from white's side
    # (simpler: score already accumulated from white's pov above)

    # --- Bishop pair bonus ---
    w_bishops = len(board.pieces(chess.BISHOP, chess.WHITE))
    b_bishops = len(board.pieces(chess.BISHOP, chess.BLACK))
    if w_bishops >= 2:
        score += 30
    if b_bishops >= 2:
        score -= 30

    # --- Pawn structure: doubled & isolated pawns ---
    for color, sign in ((chess.WHITE, 1), (chess.BLACK, -1)):
        pawns = board.pieces(chess.PAWN, color)
        files = [chess.square_file(sq) for sq in pawns]
        for f in range(8):
            cnt = files.count(f)
            if cnt > 1:
                score -= sign * 20 * (cnt - 1)   # doubled-pawn penalty
            if cnt > 0:
                # isolated pawn: no friendly pawn on adjacent files
                neighbours = [files.count(f - 1) if f > 0 else 0,
                              files.count(f + 1) if f < 7 else 0]
                if sum(neighbours) == 0:
                    score -= sign * 15

    return score

# ──────────────────────────────────────────────────────────────────────────────
# Chess / model logic
# ──────────────────────────────────────────────────────────────────────────────

def board_to_prompt(board: chess.Board) -> str:
    game = chess.pgn.Game()
    node = game
    for move in board.move_stack:
        node = node.add_variation(move)
    exporter = chess.pgn.StringExporter(headers=False, variations=False, comments=False)
    pgn = game.accept(exporter).strip()
    pgn = re.sub(r"\s*[\*\d][-\d/]*\s*$", "", pgn).strip()
    full_move = board.fullmove_number
    pgn += f" {full_move}." if board.turn == chess.WHITE else f" {full_move}..."
    return f"<|endoftext|>{pgn}"


def extract_move(text: str, board: chess.Board):
    text = re.sub(r"^\s*\d+\.+\s*", "", text).strip()
    for token in text.split()[:5]:
        clean = re.sub(r"[!?+#,;]+$", "", token)
        try:
            move = board.parse_san(clean)
            if move in board.legal_moves:
                return move
        except Exception:
            pass
        try:
            move = chess.Move.from_uci(clean.lower()[:4])
            if move in board.legal_moves:
                return move
        except Exception:
            pass
    return None


@torch.no_grad()
def _sample_candidate_moves(board: chess.Board, n: int = 15, max_attempts: int = 90):
    """
    Sample up to `n` distinct legal moves from ChessSLM using temperature sampling.
    Returns a list of chess.Move objects.
    """
    prompt  = board_to_prompt(board)
    inputs  = tokenizer(prompt, return_tensors="pt").to(device)

    candidates: list[chess.Move] = []
    seen: set[str] = set()
    attempts = 0

    while len(candidates) < n and attempts < max_attempts:
        outputs = model.generate(
            inputs.input_ids,
            max_new_tokens=12,
            do_sample=True,
            temperature=0.85,       # higher temp → more diverse candidates
            top_k=60,
            top_p=0.95,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        new_tokens = outputs[0][inputs.input_ids.shape[1]:]
        generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
        move       = extract_move(generated, board)
        attempts  += 1

        if move is not None and move.uci() not in seen:
            seen.add(move.uci())
            candidates.append(move)

    return candidates


def _custom_eval_best(board: chess.Board, candidates: list[chess.Move]):
    """
    Score each candidate move with the PST evaluator and return the best one.
    Returns (best_move, scores_dict, eval_available=True).
    scores_dict maps move-UCI → centipawn score (from the side to move).
    """
    if not candidates:
        return None, {}, False

    scores: dict[str, int] = {}
    side = board.turn   # WHITE or BLACK

    for move in candidates:
        test_board = board.copy()
        test_board.push(move)
        cp_white = _evaluate_position(test_board)
        # Convert to the perspective of the side that just moved (pre-push side)
        cp = cp_white if side == chess.WHITE else -cp_white
        scores[move.uci()] = cp

    best_uci  = max(scores, key=lambda u: scores[u])
    best_move = chess.Move.from_uci(best_uci)
    return best_move, scores, True


def get_model_move(board: chess.Board):
    """
    High-level entry point.
    Returns (chosen_move, was_model_legal, candidates, scores, stockfish_used).
    """
    n_samples  = random.randint(10, 20)
    candidates = _sample_candidate_moves(board, n=n_samples, max_attempts=n_samples * 6)

    if not candidates:
        # Model produced nothing valid; pure random fallback
        fallback = random.choice(list(board.legal_moves))
        return fallback, False, [], {}, False

    best, scores, sf_used = _custom_eval_best(board, candidates)

    if best is None:
        best = candidates[0]

    return best, True, candidates, scores, sf_used


# ──────────────────────────────────────────────────────────────────────────────
# Board rendering
# ──────────────────────────────────────────────────────────────────────────────

PIECE_COLORS = {
    "square light":            "#f0d9b5",
    "square dark":             "#b58863",
    "square light lastmove":   "#cdd16e",
    "square dark lastmove":    "#aaa23a",
}


def render_board_html(board: chess.Board, last_move=None, flipped=False, size=480):
    check_square = board.king(board.turn) if board.is_check() else None
    svg = chess.svg.board(
        board,
        lastmove=last_move,
        check=check_square,
        flipped=flipped,
        size=size,
        colors=PIECE_COLORS,
    )
    return f"""
    <div style="
        display:flex; justify-content:center; align-items:center;
        padding: 16px;
        background: radial-gradient(ellipse at center, #1a1208 0%, #0d0d0d 100%);
        border-radius: 12px;
        box-shadow: 0 0 60px rgba(0,0,0,0.8), inset 0 0 30px rgba(0,0,0,0.4);
    ">
        <div style="
            border-radius: 4px;
            overflow: hidden;
            box-shadow: 0 8px 32px rgba(0,0,0,0.6), 0 0 0 3px #3d2b0e, 0 0 0 5px #6b4c1e;
        ">
            {svg}
        </div>
    </div>
    """


def get_legal_moves_san(board: chess.Board):
    moves = []
    for move in board.legal_moves:
        try:
            moves.append(board.san(move))
        except Exception:
            pass
    return sorted(moves)


def format_move_history(board: chess.Board):
    if not board.move_stack:
        return "<em style='color:#666'>No moves yet.</em>"

    temp  = chess.Board()
    lines = []
    moves = list(board.move_stack)
    i     = 0

    while i < len(moves):
        move_num  = temp.fullmove_number
        white_san = temp.san(moves[i])
        temp.push(moves[i])
        i += 1

        if i < len(moves):
            black_san = temp.san(moves[i])
            temp.push(moves[i])
            i += 1
            lines.append(
                f"<span style='color:#8a7a5a;font-size:0.8em'>{move_num}.</span> "
                f"<span style='color:#e8d5a3'>{white_san}</span> "
                f"<span style='color:#c4b48a'>{black_san}</span>"
            )
        else:
            lines.append(
                f"<span style='color:#8a7a5a;font-size:0.8em'>{move_num}.</span> "
                f"<span style='color:#e8d5a3'>{white_san}</span>"
            )

    visible = lines[-10:]
    html  = "<div style='font-family:\"Courier New\",monospace; line-height:2; font-size:0.92em;'>"
    html += "<br>".join(visible)
    html += "</div>"
    return html


def format_candidates_html(board: chess.Board, candidates: list, scores: dict,
                            chosen_move, sf_used: bool) -> str:
    """
    Render a small table showing all ChessSLM candidates with their
    Stockfish evaluation scores and which one was chosen.
    """
    if not candidates:
        return ""

    rows = []
    for move in candidates:
        san     = board.san(move) if move in board.legal_moves else move.uci()
        cp      = scores.get(move.uci())
        is_best = (chosen_move is not None and move.uci() == chosen_move.uci())

        if cp is None:
            score_str = "<span style='color:#666'>—</span>"
        elif cp >= 9999:
            score_str = "<span style='color:#6cf06c'>+M</span>"
        elif cp <= -9999:
            score_str = "<span style='color:#f06c6c'>−M</span>"
        elif cp > 0:
            score_str = f"<span style='color:#a8d88a'>+{cp/100:.2f}</span>"
        elif cp < 0:
            score_str = f"<span style='color:#d88a8a'>{cp/100:.2f}</span>"
        else:
            score_str = "<span style='color:#aaa'>0.00</span>"

        star = "⭐ " if is_best else "&nbsp;&nbsp;&nbsp;"
        row  = (
            f"<tr style='background:{'#1e1a08' if is_best else 'transparent'};'>"
            f"<td style='padding:2px 6px;color:#c8a96e'>{star}</td>"
            f"<td style='padding:2px 8px;font-weight:{'700' if is_best else '400'};color:#e8d5a3'>{san}</td>"
            f"<td style='padding:2px 8px;text-align:right'>{score_str}</td>"
            f"</tr>"
        )
        rows.append(row)

    sf_label = (
        "<span style='color:#8a6;font-size:0.8em'>✔ PST Eval</span>"
        if sf_used else
        "<span style='color:#a64;font-size:0.8em'>⚠ Eval N/A — first candidate used</span>"
    )

    return (
        f"<div style='margin-top:6px'>"
        f"<span style='font-family:Cinzel,serif;font-size:0.75em;color:#8a7a5a;"
        f"letter-spacing:0.08em;text-transform:uppercase'>ChessSLM Candidates</span>"
        f"&nbsp;{sf_label}"
        f"<table style='width:100%;border-collapse:collapse;margin-top:4px;font-family:\"Courier New\",monospace;font-size:0.88em'>"
        + "".join(rows)
        + "</table></div>"
    )


def game_status(board: chess.Board, player_color: str):
    if board.is_checkmate():
        winner = "Black" if board.turn == chess.WHITE else "White"
        if (winner == "White") == (player_color == "white"):
            return "♟ Checkmate — You win! 🎉", "win"
        else:
            return "♟ Checkmate — ChessSLM wins!", "loss"
    if board.is_stalemate():
        return "½ Stalemate — Draw", "draw"
    if board.is_insufficient_material():
        return "½ Insufficient material — Draw", "draw"
    if board.is_seventyfive_moves():
        return "½ 75-move rule — Draw", "draw"
    if board.is_fivefold_repetition():
        return "½ Fivefold repetition — Draw", "draw"
    if board.is_check():
        return "⚠ Check!", "check"
    whose = "Your turn" if (board.turn == chess.WHITE) == (player_color == "white") else "ChessSLM is thinking..."
    return whose, "playing"


# ──────────────────────────────────────────────────────────────────────────────
# Gradio callbacks
# ──────────────────────────────────────────────────────────────────────────────

def new_game(player_color_choice: str):
    board        = chess.Board()
    player_color = "white" if player_color_choice == "⬜ White (move first)" else "black"
    flipped      = (player_color == "black")
    last_move    = None
    log_lines    = []
    candidates_html = ""

    if player_color == "black":
        move, legal, candidates, scores, sf_used = get_model_move(board)
        san       = board.san(move)

        # Build candidate table before pushing
        cand_html = format_candidates_html(board, candidates, scores, move, sf_used)

        board.push(move)
        last_move = move
        log_lines.append(f"ChessSLM opens with <strong>{san}</strong>")
        candidates_html = cand_html

    legal_moves = get_legal_moves_san(board)
    status_text, _ = game_status(board, player_color)
    board_html      = render_board_html(board, last_move=last_move, flipped=flipped)
    history_html    = format_move_history(board)

    if log_lines:
        log_html = (
            "<br>".join(f"<span style='color:#c8a96e'>{l}</span>" for l in log_lines)
            + candidates_html
        )
    else:
        log_html = "<em style='color:#666'>Game started.</em>"

    state = {
        "fen":           board.fen(),
        "move_stack":    [m.uci() for m in board.move_stack],
        "player_color":  player_color,
        "last_move_uci": last_move.uci() if last_move else None,
        "game_over":     False,
    }

    return (
        board_html,
        gr.Dropdown(choices=legal_moves, value=None, interactive=True, label="Your move"),
        status_text,
        history_html,
        log_html,
        state,
    )


def make_player_move(move_san: str, state: dict):
    if not state or state.get("game_over"):
        return gr.update(), gr.update(), "Game is over. Start a new game.", gr.update(), gr.update(), state
    if not move_san:
        return gr.update(), gr.update(), "Please select a move first.", gr.update(), gr.update(), state

    board = chess.Board()
    for uci in state["move_stack"]:
        board.push(chess.Move.from_uci(uci))

    player_color = state["player_color"]
    flipped      = (player_color == "black")
    log_lines    = []
    candidates_html = ""

    try:
        player_move = board.parse_san(move_san)
    except Exception:
        return gr.update(), gr.update(), f"Invalid move: {move_san}", gr.update(), gr.update(), state

    board.push(player_move)
    log_lines.append(f"You played <strong>{move_san}</strong>")
    last_move = player_move

    status_text, status_key = game_status(board, player_color)
    game_over = status_key in ("win", "loss", "draw")

    if not game_over:
        move, legal, candidates, scores, sf_used = get_model_move(board)
        model_san = board.san(move)

        # Build candidate table before pushing
        candidates_html = format_candidates_html(board, candidates, scores, move, sf_used)

        board.push(move)
        last_move = move

        flag = "" if legal else " <em style='color:#888'>(random fallback)</em>"
        log_lines.append(f"ChessSLM plays <strong>{model_san}</strong>{flag}")

        status_text, status_key = game_status(board, player_color)
        game_over = status_key in ("win", "loss", "draw")

    state = {
        "fen":           board.fen(),
        "move_stack":    [m.uci() for m in board.move_stack],
        "player_color":  player_color,
        "last_move_uci": last_move.uci() if last_move else None,
        "game_over":     game_over,
    }

    legal_moves  = [] if game_over else get_legal_moves_san(board)
    board_html   = render_board_html(board, last_move=last_move, flipped=flipped)
    history_html = format_move_history(board)
    log_html     = (
        "<br>".join(f"<span style='color:#c8a96e'>{l}</span>" for l in log_lines)
        + candidates_html
    )

    return (
        board_html,
        gr.Dropdown(choices=legal_moves, value=None, interactive=not game_over, label="Your move"),
        status_text,
        history_html,
        log_html,
        state,
    )


# ──────────────────────────────────────────────────────────────────────────────
# CSS
# ──────────────────────────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel:wght@400;600;700&family=Crimson+Text:ital,wght@0,400;0,600;1,400&display=swap');

body, .gradio-container {
    background: #0d0d0d !important;
    color: #e8d5a3 !important;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    font-family: 'Crimson Text', Georgia, serif !important;
}

h1, h2, h3 {
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.08em;
}

#title-block {
    text-align: center;
    padding: 2rem 0 1rem;
    border-bottom: 1px solid #3d2b0e;
    margin-bottom: 1.5rem;
}

.panel {
    background: #141008 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 8px !important;
    padding: 1rem !important;
}

#status-bar {
    text-align: center;
    font-family: 'Cinzel', serif;
    font-size: 1.1em;
    letter-spacing: 0.05em;
    padding: 0.6rem 1rem;
    border-radius: 6px;
    background: #1a1208;
    border: 1px solid #4a3520;
    color: #f0c060;
}

button.primary {
    background: linear-gradient(135deg, #8b6914 0%, #c4922a 50%, #8b6914 100%) !important;
    border: 1px solid #d4a843 !important;
                color: #fff8e8 !important;
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.06em !important;
    font-size: 0.9em !important;
    border-radius: 4px !important;
    transition: all 0.2s ease !important;
}
button.primary:hover {
    background: linear-gradient(135deg, #a07820 0%, #d4a843 50%, #a07820 100%) !important;
    box-shadow: 0 0 16px rgba(212,168,67,0.4) !important;
}

button.secondary {
    background: #1e1810 !important;
    border: 1px solid #5a4020 !important;
    color: #c8a96e !important;
    font-family: 'Cinzel', serif !important;
    letter-spacing: 0.04em !important;
    border-radius: 4px !important;
}

select, .gr-dropdown select {
    background: #1a1208 !important;
    border: 1px solid #5a4020 !important;
    color: #e8d5a3 !important;
    font-family: 'Crimson Text', serif !important;
    font-size: 1em !important;
}

#move-log {
    background: #0f0c06 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 6px;
    padding: 0.8rem 1rem;
    font-family: 'Crimson Text', serif;
    font-size: 0.95em;
    line-height: 1.8;
    min-height: 80px;
    color: #c8a96e;
}

#history-panel {
    background: #0f0c06 !important;
    border: 1px solid #3d2b0e !important;
    border-radius: 6px;
    padding: 0.8rem 1rem;
    min-height: 200px;
    max-height: 320px;
    overflow-y: auto;
}

.gr-radio label {
    color: #e8d5a3 !important;
    font-family: 'Crimson Text', serif !important;
}

label span {
    color: #a08050 !important;
    font-family: 'Cinzel', serif !important;
    font-size: 0.8em !important;
    letter-spacing: 0.06em !important;
    text-transform: uppercase !important;
}

::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: #0d0d0d; }
::-webkit-scrollbar-thumb { background: #5a4020; border-radius: 3px; }
"""

# ──────────────────────────────────────────────────────────────────────────────
# Layout
# ──────────────────────────────────────────────────────────────────────────────

with gr.Blocks(css=CSS, title="ChessSLM — Play vs AI") as demo:

    state = gr.State({})

    gr.HTML("""
    <div id="title-block">
        <h1 style="
            font-family:'Cinzel',serif;
            font-size:2.4em;
            font-weight:700;
            color:#e8c96e;
            text-shadow: 0 0 30px rgba(232,180,80,0.4);
            margin:0 0 0.3rem;
            letter-spacing:0.12em;
        ">♛ ChessSLM</h1>
        <p style="
            font-family:'Crimson Text',serif;
            color:#8a7a5a;
            font-size:1.1em;
            font-style:italic;
            margin:0;
        ">GPT-2 generates 10–20 candidate moves · PST evaluator picks the best</p>
    </div>
    """)

    with gr.Row():

        with gr.Column(scale=3):
            board_display = gr.HTML(
                value=render_board_html(chess.Board()),
                label="Board"
            )
            status_display = gr.HTML(
                value="<div id='status-bar'>Choose your colour and press New Game</div>"
            )

        with gr.Column(scale=2):

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>NEW GAME</h3>")

            color_choice = gr.Radio(
                choices=["⬜ White (move first)", "⬛ Black (move second)"],
                value="⬜ White (move first)",
                label="Play as",
            )

            new_game_btn = gr.Button("♟  New Game", variant="primary", size="lg")

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>YOUR MOVE</h3>")

            move_dropdown = gr.Dropdown(
                choices=[],
                value=None,
                label="Select move (SAN notation)",
                interactive=False,
            )

            move_btn = gr.Button("▶  Make Move", variant="secondary")

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>MOVE LOG</h3>")

            log_display = gr.HTML(
                value="<div id='move-log'><em style='color:#555'>Start a new game to begin.</em></div>",
            )

            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")

            gr.HTML("<h3 style='font-family:Cinzel,serif;color:#c8a96e;font-size:1em;letter-spacing:0.1em;margin:0 0 0.5rem;'>GAME HISTORY</h3>")

            history_display = gr.HTML(
                value="<div id='history-panel'><em style='color:#555'>No moves yet.</em></div>",
            )

    gr.HTML("""
    <div style="
        text-align:center; margin-top:2rem; padding-top:1rem;
        border-top:1px solid #2a1e0a;
        font-family:'Crimson Text',serif; font-size:0.85em; color:#5a4a30;
    ">
        Model: <a href="https://huggingface.co/FlameF0X/ChessSLM" target="_blank"
            style="color:#8a6a30; text-decoration:none;">FlameF0X/ChessSLM</a>
        &nbsp;·&nbsp; GPT-2 samples 10–20 candidates (temp=0.85) · PST evaluator selects the best
    </div>
    """)

    new_game_btn.click(
        fn=new_game,
        inputs=[color_choice],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )

    move_btn.click(
        fn=make_player_move,
        inputs=[move_dropdown, state],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )

    move_dropdown.select(
        fn=make_player_move,
        inputs=[move_dropdown, state],
        outputs=[board_display, move_dropdown, status_display, history_display, log_display, state],
    )


if __name__ == "__main__":
    demo.launch()