Upload app.py

app.py

@@ -1,556 +1,571 @@
 """
-ChessSLM — AI vs AI  &  Model Arena
-HF Space  |  Gradio 4
+ChessSLM — Play against FlameF0X/ChessSLM
+Hugging Face Space  |  Gradio app
 """
 
-import re, random, time, torch, chess, chess.pgn, chess.svg, gradio as gr
-from typing import Optional
-from transformers import (
-    GPT2LMHeadModel, GPT2Tokenizer,
-    AutoModelForCausalLM, AutoTokenizer,
-)
+import re
+import random
+import time
+import chess
+import chess.svg
+import chess.pgn
+import gradio as gr
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Model loading (once at startup)
+# ──────────────────────────────────────────────────────────────────────────────
 
-# ══════════════════════════════════════════════════════════════════════════════
-# CONFIG
-# ══════════════════════════════════════════════════════════════════════════════
 MODEL_ID = "FlameF0X/ChessSLM"
-BOARD_SZ = 460
-
-# ══════════════════════════════════════════════════════════════════════════════
-# MODEL LOADING  (lazy — on first move, not at build time)
-# ══════════════════════════════════════════════════════════════════════════════
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-_tok  = None
-_mdl  = None
-_extra: dict = {}   # cache for custom models
-
-def _ensure_main_model():
-    global _tok, _mdl
-    if _mdl is not None:
-        return
-    print(f"Loading {MODEL_ID}…")
-    _tok = GPT2Tokenizer.from_pretrained(MODEL_ID)
-    _tok.pad_token = _tok.eos_token
-    _mdl = GPT2LMHeadModel.from_pretrained(MODEL_ID).to(device).eval()
-    _mdl.config.use_cache = True
-    print(f"✓ {MODEL_ID} ready on {device}")
-
-def load_custom(model_id: str):
-    if model_id in _extra:
-        return _extra[model_id]
-    tok = AutoTokenizer.from_pretrained(model_id)
-    if tok.pad_token is None:
-        tok.pad_token = tok.eos_token
-    mdl = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    ).to(device).eval()
-    mdl.config.use_cache = True
-    _extra[model_id] = (mdl, tok)
-    return mdl, tok
-
-# ══════════════════════════════════════════════════════════════════════════════
-# CHESS HELPERS
-# ══════════════════════════════════════════════════════════════════════════════
+
+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}")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Chess / model logic
+# ──────────────────────────────────────────────────────────────────────────────
 
 def board_to_prompt(board: chess.Board) -> str:
     game = chess.pgn.Game()
     node = game
-    for m in board.move_stack:
-        node = node.add_variation(m)
-    exp = chess.pgn.StringExporter(headers=False, variations=False, comments=False)
-    pgn = game.accept(exp).strip()
+    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()
-    n = board.fullmove_number
-    pgn += f" {n}." if board.turn == chess.WHITE else f" {n}..."
+    full_move = board.fullmove_number
+    pgn += f" {full_move}." if board.turn == chess.WHITE else f" {full_move}..."
     return f"<|endoftext|>{pgn}"
 
-def parse_model_move(text: str, board: chess.Board) -> Optional[chess.Move]:
+
+def extract_move(text: str, board: chess.Board):
     text = re.sub(r"^\s*\d+\.+\s*", "", text).strip()
-    for tok in text.split()[:5]:
-        t = re.sub(r"[!?+#,;]+$", "", tok)
+    for token in text.split()[:5]:
+        clean = re.sub(r"[!?+#,;]+$", "", token)
         try:
-            m = board.parse_san(t)
-            if m in board.legal_moves: return m
-        except: pass
+            move = board.parse_san(clean)
+            if move in board.legal_moves:
+                return move
+        except Exception:
+            pass
         try:
-            m = chess.Move.from_uci(t.lower()[:4])
-            if m in board.legal_moves: return m
-        except: pass
+            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 ai_move(board: chess.Board, mdl=None, tok=None):
-    """Ask a model for its next move. Returns (move, was_legal)."""
-    if mdl is None:
-        _ensure_main_model()
-        mdl, tok = _mdl, _tok
-    inp = tok(board_to_prompt(board), return_tensors="pt").to(device)
-    out = mdl.generate(
-        inp.input_ids, max_new_tokens=12, do_sample=True,
-        temperature=0.3, top_k=40, top_p=0.9, repetition_penalty=1.1,
-        pad_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id,
+def get_model_move(board: chess.Board):
+    prompt  = board_to_prompt(board)
+    inputs  = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(
+        inputs.input_ids,
+        max_new_tokens=12,
+        do_sample=True,
+        temperature=0.3,
+        top_k=40,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
-    txt = tok.decode(out[0][inp.input_ids.shape[1]:], skip_special_tokens=True)
-    mv  = parse_model_move(txt, board)
-    if mv: return mv, True
+    new_tokens = outputs[0][inputs.input_ids.shape[1]:]
+    generated  = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    move       = extract_move(generated, board)
+    if move:
+        return move, True
     return random.choice(list(board.legal_moves)), False
 
-def game_over_status(board: chess.Board, label_w: str, label_b: str) -> tuple[str, bool]:
-    if board.is_checkmate():
-        winner = label_w if board.turn == chess.BLACK else label_b
-        return f"♟ Checkmate — {winner} wins!", True
-    if board.is_stalemate():              return "½ Stalemate — Draw", True
-    if board.is_insufficient_material(): return "½ Draw (material)", True
-    if board.is_seventyfive_moves():     return "½ Draw (75 moves)", True
-    if board.is_fivefold_repetition():   return "½ Draw (repetition)", True
-    if board.is_check():
-        in_check = label_w if board.turn == chess.WHITE else label_b
-        return f"⚠ Check on {in_check}!", False
-    to_move = label_w if board.turn == chess.WHITE else label_b
-    return f"{to_move} to move…", False
-
-# ══════════════════════════════════════════════════════════════════════════════
-# HALLUCINATION CARD
-# ══════════════════════════════════════════════════════════════════════════════
-
-def _ratio_color(legal: int, halluc: int) -> str:
-    if halluc == 0:                      return "#50e090"
-    if legal / max(halluc, 1) >= 3:      return "#50e090"
-    if legal / max(halluc, 1) >= 1:      return "#e0b840"
-    return "#e05050"
-
-def halluc_side_html(legal: int, halluc: int, label: str) -> str:
-    total = legal + halluc
-    ratio = f"{legal/halluc:.2f}" if halluc > 0 else "∞"
-    pct   = (halluc / total * 100) if total > 0 else 0
-    col   = _ratio_color(legal, halluc)
-    bar_w = max(2, int(pct)) if total > 0 else 0
-    return f"""
-<div style="flex:1;min-width:0;padding:0 .5rem;">
-  <div style="font-size:.68em;color:#6a5a38;letter-spacing:.06em;margin-bottom:.3rem;
-      white-space:nowrap;overflow:hidden;text-overflow:ellipsis;">{label}</div>
-  <div style="font-size:2em;font-weight:700;color:{col};line-height:1;
-      font-family:'Cinzel',serif;">{ratio}</div>
-  <div style="font-size:.75em;color:#7a6a48;margin:.2rem 0 .45rem;">
-    <span style="color:#e8d5a3">{legal}</span>✓&nbsp;
-    <span style="color:#e8d5a3">{halluc}</span>✗&nbsp;
-    <span style="color:#555">{pct:.0f}% halluc</span>
-  </div>
-  <div style="background:#1a1408;border-radius:3px;height:5px;overflow:hidden;">
-    <div style="height:100%;width:{bar_w}%;background:{col};border-radius:3px;
-        transition:width .5s ease;"></div>
-  </div>
-</div>"""
-
-def halluc_html(
-    legal_w: int, halluc_w: int, label_w: str,
-    legal_b: int, halluc_b: int, label_b: str,
-) -> str:
-    return f"""
-<div style="background:#0d0a05;border:1px solid #2e1f08;border-radius:8px;
-    padding:.75rem 1rem;margin-top:.5rem;">
-  <div style="font-size:.68em;font-family:'Cinzel',serif;letter-spacing:.09em;
-      color:#5a4a28;margin-bottom:.6rem;">HALLUCINATION RATIO</div>
-  <div style="display:flex;gap:.5rem;">
-    {halluc_side_html(legal_w, halluc_w, label_w)}
-    <div style="width:1px;background:#2e1f08;align-self:stretch;"></div>
-    {halluc_side_html(legal_b, halluc_b, label_b)}
-  </div>
-</div>"""
-
-# ══════════════════════════════════════════════════════════════════════════════
-# RENDERING
-# ══════════════════════════════════════════════════════════════════════════════
-
-_SVG_COLORS = {
-    "square light":          "#f0d9b5",
-    "square dark":           "#b58863",
+# ──────────────────────────────────────────────────────────────────────────────
+# Board rendering
+# ──────────────────────────────────────────────────────────────────────────────
+
+PIECE_COLORS = {
+    "square light":        "#f0d9b5",
+    "square dark":         "#b58863",
     "square light lastmove": "#cdd16e",
     "square dark lastmove":  "#aaa23a",
 }
 
-def render_board(board: chess.Board, last_move=None) -> str:
-    check_sq = board.king(board.turn) if board.is_check() else None
+def render_board_html(board: chess.Board, last_move=None, flipped=False, size=480):
+    check_square = None
+    if board.is_check():
+        check_square = board.king(board.turn)
+
     svg = chess.svg.board(
-        board, lastmove=last_move, check=check_sq,
-        size=BOARD_SZ, coordinates=True, colors=_SVG_COLORS,
+        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,.8);">
-  <div style="border-radius:4px;overflow:hidden;
-      box-shadow:0 8px 32px rgba(0,0,0,.6),0 0 0 3px #3d2b0e,0 0 0 5px #6b4c1e;">
-    {svg}
-  </div>
-</div>"""
-
-def fmt_history(board: chess.Board) -> str:
+    <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:#555'>No moves yet.</em>"
-    tmp, lines, mvs, i = chess.Board(), [], list(board.move_stack), 0
-    while i < len(mvs):
-        n = tmp.fullmove_number
-        w = tmp.san(mvs[i]); tmp.push(mvs[i]); i += 1
-        if i < len(mvs):
-            b = tmp.san(mvs[i]); tmp.push(mvs[i]); i += 1
+        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:#555;font-size:.78em'>{n}.</span> "
-                f"<span style='color:#e8d5a3'>{w}</span> "
-                f"<span style='color:#c4b48a'>{b}</span>")
+                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:#555;font-size:.78em'>{n}.</span> "
-                f"<span style='color:#e8d5a3'>{w}</span>")
-    inner = "<br>".join(lines[-16:])
-    return (f"<div style='font-family:\"Courier New\",monospace;"
-            f"line-height:2;font-size:.88em'>{inner}</div>")
-
-def status_html(msg: str) -> str:
-    return f"<div id='status-bar'>{msg}</div>"
-
-def log_html(lines: list[str]) -> str:
-    if not lines: return "<em style='color:#555'>—</em>"
-    items = "".join(f"<div>{l}</div>" for l in lines[-8:])
-    return (f"<div style='font-family:\"Crimson Text\",serif;"
-            f"line-height:1.9;font-size:.95em;color:#c8a96e'>{items}</div>")
-
-# ══════════════════════════════════════════════════════════════════════════════
-# MODE 1 — AI vs AI  (ChessSLM self-play)
-# ══════════════════════════════════════════════════════════════════════════════
-
-def aivsai_run(delay_val: float):
-    board = chess.Board()
-    lw, hw, lb, hb = 0, 0, 0, 0
-    label_w, label_b = "ChessSLM (W)", "ChessSLM (B)"
-    log_lines = []
-
-    yield (render_board(board),
-           status_html("Starting…"),
-           fmt_history(board),
-           log_html([]),
-           halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-    while not board.is_game_over() and len(board.move_stack) < 200:
-        is_white = board.turn == chess.WHITE
-        mv, legal = ai_move(board)
-        san = board.san(mv)
-        board.push(mv)
-
-        if is_white:
-            if legal: lw += 1
-            else:     hw += 1
-            tag = label_w
+                f"<span style='color:#8a7a5a;font-size:0.8em'>{move_num}.</span> "
+                f"<span style='color:#e8d5a3'>{white_san}</span>"
+            )
+
+    # Last 10 pairs
+    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 game_status(board: chess.Board, player_color: str):
+    if board.is_checkmate():
+        loser  = "White" if board.turn == chess.WHITE else "Black"
+        winner = "Black" if board.turn == chess.WHITE else "White"
+        if (winner == "White") == (player_color == "white"):
+            return "♟ Checkmate — You win! 🎉", "win"
         else:
-            if legal: lb += 1
-            else:     hb += 1
-            tag = label_b
-
-        flag = "" if legal else " <em style='color:#e05050'>(hallucinated)</em>"
-        log_lines.append(f"<b>{tag}</b>: {san}{flag}")
-        stxt, over = game_over_status(board, label_w, label_b)
-
-        yield (render_board(board, mv),
-               status_html(stxt),
-               fmt_history(board),
-               log_html(log_lines),
-               halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-        if over: return
-        time.sleep(max(0.05, float(delay_val)))
-
-    if not board.is_game_over():
-        yield (render_board(board),
-               status_html("½ Draw — move limit reached"),
-               fmt_history(board),
-               log_html(log_lines),
-               halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-# ══════════════════════════════════════════════════════════════════════════════
-# MODE 2 — MODEL ARENA  (ChessSLM vs any HF causal LM)
-# ══════════════════════════════════════════════════════════════════════════════
-
-def arena_load(model_id: str):
-    mid = model_id.strip()
-    if not mid:
-        return status_html("⚠ Enter a model repo ID first"), gr.update()
-    try:
-        load_custom(mid)
-        return (status_html(f"✓ <b>{mid}</b> loaded — press ▶ Play"),
-                gr.update(interactive=True))
-    except Exception as e:
-        return status_html(f"✗ Could not load <em>{mid}</em>: {e}"), gr.update()
-
-def arena_run(model_id: str, delay_val: float):
-    mid = model_id.strip()
-    if not mid:
-        yield (gr.update(), status_html("⚠ Load a model first"),
-               gr.update(), gr.update(), gr.update())
-        return
+            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):
+    """Reset the board and, if player chose Black, let the model move first."""
+    board       = chess.Board()
+    player_color = "white" if player_color_choice == "⬜ White (move first)" else "black"
+    flipped      = (player_color == "black")
+    last_move    = None
+    log_lines    = []
+
+    # If player chose Black, model plays White first
+    if player_color == "black":
+        move, legal = get_model_move(board)
+        san         = board.san(move)
+        board.push(move)
+        last_move   = move
+        log_lines.append(f"ChessSLM opens with **{san}**")
+
+    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)
+    log_html        = "<br>".join(log_lines) if log_lines else "<em style='color:#666'>Game started.</em>"
+
+    # Serialise state
+    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):
+    """Apply the player's chosen move, then let the model respond."""
+    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
+        )
+
+    # Rebuild board from 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    = []
+
+    # Apply player move
     try:
-        mdl2, tok2 = load_custom(mid)
-    except Exception as e:
-        yield (gr.update(), status_html(f"✗ {e}"),
-               gr.update(), gr.update(), gr.update())
-        return
-
-    label_w = "ChessSLM (W)"
-    label_b = f"{mid.split('/')[-1]} (B)"
-    board   = chess.Board()
-    lw, hw, lb, hb = 0, 0, 0, 0
-    log_lines = []
-
-    yield (render_board(board),
-           status_html(f"{label_w}  vs  {label_b}"),
-           fmt_history(board),
-           log_html([]),
-           halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-    while not board.is_game_over() and len(board.move_stack) < 200:
-        if board.turn == chess.WHITE:
-            mv, legal = ai_move(board)
-            tag = label_w
-            if legal: lw += 1
-            else:     hw += 1
-        else:
-            mv, legal = ai_move(board, mdl2, tok2)
-            tag = label_b
-            if legal: lb += 1
-            else:     hb += 1
-
-        san = board.san(mv)
-        board.push(mv)
-        flag = "" if legal else " <em style='color:#e05050'>(hallucinated)</em>"
-        log_lines.append(f"<b>{tag}</b>: {san}{flag}")
-        stxt, over = game_over_status(board, label_w, label_b)
-
-        yield (render_board(board, mv),
-               status_html(stxt),
-               fmt_history(board),
-               log_html(log_lines),
-               halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-        if over: return
-        time.sleep(max(0.05, float(delay_val)))
-
-    if not board.is_game_over():
-        yield (render_board(board),
-               status_html("½ Draw — move limit reached"),
-               fmt_history(board),
-               log_html(log_lines),
-               halluc_html(lw, hw, label_w, lb, hb, label_b))
-
-# ══════════════════════════════════════════════════════════════════════════════
+        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 **{move_san}**")
+    last_move = player_move
+
+    # Check game over after player move
+    status_text, status_key = game_status(board, player_color)
+    game_over = status_key in ("win", "loss", "draw")
+
+    model_san = None
+    if not game_over:
+        # Model responds
+        model_move_obj, legal = get_model_move(board)
+        model_san = board.san(model_move_obj)
+        board.push(model_move_obj)
+        last_move = model_move_obj
+        flag = "" if legal else " *(random fallback)*"
+        log_lines.append(f"ChessSLM plays **{model_san}**{flag}")
+
+        status_text, status_key = game_status(board, player_color)
+        game_over = status_key in ("win", "loss", "draw")
+
+    # Update state
+    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]
+    )
+
+    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');
 
-*, *::before, *::after { box-sizing: border-box; }
-
 body, .gradio-container {
-    background: #0a0805 !important;
+    background: #0d0d0d !important;
     color: #e8d5a3 !important;
+}
+
+.gradio-container {
+    max-width: 1100px !important;
+    margin: 0 auto !important;
     font-family: 'Crimson Text', Georgia, serif !important;
 }
-.gradio-container { max-width: 1180px !important; margin: 0 auto !important; }
 
-#site-header {
+h1, h2, h3 {
+    font-family: 'Cinzel', serif !important;
+    letter-spacing: 0.08em;
+}
+
+/* Title */
+#title-block {
     text-align: center;
-    padding: 2rem 0 1.2rem;
-    border-bottom: 1px solid #2e1f08;
-    margin-bottom: 1.4rem;
+    padding: 2rem 0 1rem;
+    border-bottom: 1px solid #3d2b0e;
+    margin-bottom: 1.5rem;
 }
 
+/* Panels */
+.panel {
+    background: #141008 !important;
+    border: 1px solid #3d2b0e !important;
+    border-radius: 8px !important;
+    padding: 1rem !important;
+}
+
+/* Status bar */
 #status-bar {
     text-align: center;
     font-family: 'Cinzel', serif;
-    font-size: 1em;
+    font-size: 1.1em;
     letter-spacing: 0.05em;
     padding: 0.6rem 1rem;
-    border-radius: 5px;
-    background: #110e07;
-    border: 1px solid #3d2b0e;
+    border-radius: 6px;
+    background: #1a1208;
+    border: 1px solid #4a3520;
     color: #f0c060;
-    min-height: 2.6rem;
-}
-
-.tabs { border: none !important; }
-.tabitem { background: transparent !important; }
-.tab-nav button {
-    font-family: 'Cinzel', serif !important;
-    font-size: 0.8em !important;
-    letter-spacing: 0.07em !important;
-    color: #8a7a5a !important;
-    border-bottom: 2px solid transparent !important;
-    background: transparent !important;
-    padding: 0.55rem 1.1rem !important;
-    transition: color .2s, border-color .2s !important;
-}
-.tab-nav button.selected {
-    color: #e8c060 !important;
-    border-bottom-color: #c49020 !important;
 }
 
+/* Buttons */
 button.primary {
-    background: linear-gradient(135deg, #7a5a10, #c4922a, #7a5a10) !important;
-    border: 1px solid #c49020 !important;
+    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.07em !important;
-    font-size: 0.85em !important;
+    letter-spacing: 0.06em !important;
+    font-size: 0.9em !important;
     border-radius: 4px !important;
-    transition: all .2s !important;
+    transition: all 0.2s ease !important;
 }
 button.primary:hover {
-    box-shadow: 0 0 20px rgba(196,144,32,.45) !important;
-    filter: brightness(1.12) !important;
+    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: #1a140a !important;
-    border: 1px solid #4a3318 !important;
+    background: #1e1810 !important;
+    border: 1px solid #5a4020 !important;
     color: #c8a96e !important;
     font-family: 'Cinzel', serif !important;
-    font-size: 0.82em !important;
+    letter-spacing: 0.04em !important;
     border-radius: 4px !important;
 }
-button.stop {
-    background: #2a0808 !important;
-    border: 1px solid #6b2020 !important;
-    color: #e07070 !important;
-}
 
-input, textarea, select {
-    background: #130f08 !important;
-    border: 1px solid #3d2b0e !important;
+/* Dropdown */
+select, .gr-dropdown select {
+    background: #1a1208 !important;
+    border: 1px solid #5a4020 !important;
     color: #e8d5a3 !important;
     font-family: 'Crimson Text', serif !important;
-    border-radius: 4px !important;
-}
-label span {
-    font-family: 'Cinzel', serif !important;
-    font-size: 0.73em !important;
-    letter-spacing: 0.07em !important;
-    text-transform: uppercase !important;
-    color: #8a6a30 !important;
+    font-size: 1em !important;
 }
 
-.panel {
-    background: #0d0a05;
-    border: 1px solid #2e1f08;
+/* Move log */
+#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;
-    max-height: 260px;
+    color: #c8a96e;
+}
+
+/* History panel */
+#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;
 }
 
-.div-line { height: 1px; background: #2e1f08; margin: 0.9rem 0; }
+/* Radio buttons */
+.gr-radio label {
+    color: #e8d5a3 !important;
+    font-family: 'Crimson Text', serif !important;
+}
 
-::-webkit-scrollbar { width: 5px; }
-::-webkit-scrollbar-track { background: #0a0805; }
-::-webkit-scrollbar-thumb { background: #4a3318; border-radius: 3px; }
+/* Labels */
+label span {
+    color: #a08050 !important;
+    font-family: 'Cinzel', serif !important;
+    font-size: 0.8em !important;
+    letter-spacing: 0.06em !important;
+    text-transform: uppercase !important;
+}
 
-input[type=range] { accent-color: #c49020 !important; }
+/* Scrollbar */
+::-webkit-scrollbar { width: 6px; }
+::-webkit-scrollbar-track { background: #0d0d0d; }
+::-webkit-scrollbar-thumb { background: #5a4020; border-radius: 3px; }
 """
 
-# ══════════════════════════════════════════════════════════════════════════════
-# LAYOUT
-# ══════════════════════════════════════════════════════════════════════════════
+# ──────────────────────────────────────────────────────────────────────────────
+# Layout
+# ──────────────────────────────────────────────────────────────────────────────
 
-_BOARD_0  = render_board(chess.Board())
-_HALLUC_0 = halluc_html(0, 0, "White", 0, 0, "Black")
+with gr.Blocks(css=CSS, title="ChessSLM — Play vs AI") as demo:
 
-with gr.Blocks(css=CSS, title="ChessSLM Arena") as demo:
+    state = gr.State({})
 
+    # ── Header ────────────────────────────────────────────────────────────────
     gr.HTML("""
-    <div id="site-header">
-      <h1 style="font-family:'Cinzel',serif;font-size:2.3em;font-weight:700;
-          color:#e8c060;text-shadow:0 0 30px rgba(232,180,60,.35);
-          margin:0 0 .3rem;letter-spacing:.14em;">♛ ChessSLM Arena</h1>
-      <p style="font-family:'Crimson Text',serif;color:#7a6a48;
-          font-size:1.05em;font-style:italic;margin:0;">
-        GPT-2 pre-trained on 100,000 chess games · watch the models battle</p>
-    </div>""")
-
+    <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;
+        ">Play against a GPT-2 model trained on 100,000 chess games</p>
+    </div>
+    """)
+
+    # ── Main layout ───────────────────────────────────────────────────────────
     with gr.Row():
 
-        # ── Left column: board + status + hallucination card ──────────────────
-        with gr.Column(scale=3, min_width=480):
-            board_out  = gr.HTML(value=_BOARD_0)
-            status_out = gr.HTML(value=status_html("Choose a mode and press ▶ Play"))
-            halluc_out = gr.HTML(value=_HALLUC_0)
-
-        # ── Right column: tabs ────────────────────────────────────────────────
-        with gr.Column(scale=2, min_width=300):
-
-            with gr.Tabs():
-
-                # ── Tab 1: AI vs AI ───────────────────────────────────────────
-                with gr.Tab("🤖  AI vs AI"):
-                    gr.HTML("<p style='font-family:Crimson Text,serif;color:#7a6a48;"
-                            "font-style:italic;font-size:.95em;margin:.3rem 0 .8rem;'>"
-                            "ChessSLM plays itself from both sides.</p>")
-                    delay1 = gr.Slider(minimum=0.1, maximum=3.0, value=0.6, step=0.1,
-                                       label="Delay between moves (s)")
-                    with gr.Row():
-                        btn1_play = gr.Button("▶  Play", variant="primary")
-                        btn1_stop = gr.Button("⏹  Stop", variant="secondary",
-                                             elem_classes=["stop"])
-
-                # ── Tab 2: Model Arena ────────────────────────────────────────
-                with gr.Tab("⚔️  Model Arena"):
-                    gr.HTML("<p style='font-family:Crimson Text,serif;color:#7a6a48;"
-                            "font-style:italic;font-size:.95em;margin:.3rem 0 .8rem;'>"
-                            "ChessSLM (White) vs any HuggingFace causal LM (Black).</p>")
-                    arena_id = gr.Textbox(
-                        placeholder="e.g. another-user/chess-gpt",
-                        label="Opponent model repo ID",
-                    )
-                    btn_load = gr.Button("⬇  Load Opponent", variant="secondary")
-                    gr.HTML("<div class='div-line'></div>")
-                    delay2 = gr.Slider(minimum=0.1, maximum=3.0, value=0.6, step=0.1,
-                                       label="Delay between moves (s)")
-                    with gr.Row():
-                        btn2_play = gr.Button("▶  Play", variant="primary",
-                                             interactive=False)
-                        btn2_stop = gr.Button("⏹  Stop", variant="secondary",
-                                             elem_classes=["stop"])
-
-            gr.HTML("<div class='div-line'></div>")
-            gr.HTML("<p style='font-family:Cinzel,serif;font-size:.73em;"
-                    "color:#8a6a30;letter-spacing:.08em;margin:0 0 .4rem;'>COMMENTARY</p>")
-            log_out = gr.HTML(
-                value="<div class='panel'><em style='color:#555'>—</em></div>")
-
-            gr.HTML("<div class='div-line'></div>")
-            gr.HTML("<p style='font-family:Cinzel,serif;font-size:.73em;"
-                    "color:#8a6a30;letter-spacing:.08em;margin:0 0 .4rem;'>MOVE LIST</p>")
-            hist_out = gr.HTML(
-                value="<div class='panel'><em style='color:#555'>No moves yet.</em></div>")
+        # Left: board
+        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>"
+            )
 
-    gr.HTML("""
-    <div style="text-align:center;margin-top:1.6rem;padding-top:.9rem;
-        border-top:1px solid #1e1508;font-family:'Crimson Text',serif;
-        font-size:.82em;color:#4a3a20;">
-      <a href="https://huggingface.co/FlameF0X/ChessSLM" target="_blank"
-          style="color:#7a5a28;text-decoration:none;">FlameF0X/ChessSLM</a>
-      &nbsp;·&nbsp; GPT-2 · 100K PGN games · top-k sampling (temp=0.3)
-    </div>""")
+        # Right: controls + history
+        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>")
 
-    # ── Event wiring ──────────────────────────────────────────────────────────
-    _outs = [board_out, status_out, hist_out, log_out, halluc_out]
+            move_dropdown = gr.Dropdown(
+                choices=[],
+                value=None,
+                label="Select move (SAN notation)",
+                interactive=False,
+            )
 
-    ev1 = btn1_play.click(fn=aivsai_run, inputs=[delay1], outputs=_outs)
-    btn1_stop.click(fn=None, cancels=[ev1])
+            move_btn = gr.Button("▶  Make Move", variant="secondary")
 
-    btn_load.click(fn=arena_load, inputs=[arena_id],
-                   outputs=[status_out, btn2_play])
+            gr.HTML("<div style='height:1px;background:#3d2b0e;margin:1rem 0;'></div>")
 
-    ev2 = btn2_play.click(fn=arena_run, inputs=[arena_id, delay2], outputs=_outs)
-    btn2_stop.click(fn=None, cancels=[ev2])
+            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>",
+            )
+
+    # ── Footer ────────────────────────────────────────────────────────────────
+    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 pre-trained on 100K PGN games
+        &nbsp;·&nbsp; Move selection uses top-k sampling (temp=0.3)
+    </div>
+    """)
+
+    # ── Wiring ────────────────────────────────────────────────────────────────
+    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__":