Update app.py

app.py

@@ -1,394 +1,673 @@
 import os
-import re
-from io import BytesIO
-from typing import List, Tuple
+import random
+import math
+from dataclasses import dataclass
+from typing import List, Tuple, Dict, Optional
 
 import gradio as gr
 import torch
-import numpy as np
-from PIL import Image, ImageDraw, ImageOps
-import fitz  # PyMuPDF
+from PIL import Image, ImageDraw, ImageFont
 
 from transformers import (
-    TrOCRProcessor,
-    VisionEncoderDecoderModel,
-    BlipProcessor,
-    BlipForConditionalGeneration,
+    AutoTokenizer,
+    AutoModel,
+    AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM,
 )
-from transformers.utils import logging as hf_logging
 
-# -------------------------
-# CPU-only, quieter logs
-# -------------------------
-hf_logging.set_verbosity_error()
+# ============================================================
+# CPU setup
+# ============================================================
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-
 DEVICE = torch.device("cpu")
 torch.set_num_threads(int(os.getenv("TORCH_NUM_THREADS", "4")))
 
-TROCR_NAME = os.getenv("TROCR_MODEL", "microsoft/trocr-base-printed")
-BLIP_NAME = os.getenv("BLIP_MODEL", "Salesforce/blip-image-captioning-base")
+# ============================================================
+# 3 Transformers (minimum)
+#   1) Coach (Seq2Seq)
+#   2) Opponent (Causal LM)
+#   3) Embeddings (Encoder)
+# ============================================================
+COACH_MODEL_NAME = os.getenv("COACH_MODEL", "google/flan-t5-small")
+OPP_MODEL_NAME = os.getenv("OPP_MODEL", "distilgpt2")
+EMB_MODEL_NAME = os.getenv("EMB_MODEL", "sentence-transformers/paraphrase-MiniLM-L3-v2")
+
+coach_tok = AutoTokenizer.from_pretrained(COACH_MODEL_NAME)
+coach_model = AutoModelForSeq2SeqLM.from_pretrained(COACH_MODEL_NAME).eval().to(DEVICE)
+
+opp_tok = AutoTokenizer.from_pretrained(OPP_MODEL_NAME)
+opp_model = AutoModelForCausalLM.from_pretrained(OPP_MODEL_NAME).eval().to(DEVICE)
+
+emb_tok = AutoTokenizer.from_pretrained(EMB_MODEL_NAME)
+emb_model = AutoModel.from_pretrained(EMB_MODEL_NAME).eval().to(DEVICE)
+
+
+# ============================================================
+# Checkers engine (English draughts-like)
+# Pieces:
+#   '.' empty
+#   'w' white man (user)
+#   'W' white king
+#   'b' black man (bot)
+#   'B' black king
+#
+# Coordinates:
+#   internal: r=0..7 top->bottom, c=0..7 left->right
+#   dark squares: (r+c)%2==1
+# Move string:
+#   "b6-a5" or "c3-e5-g7" using a-h and 1-8 (1 is bottom row).
+# ============================================================
+
+def inside(r: int, c: int) -> bool:
+    return 0 <= r < 8 and 0 <= c < 8
+
+def is_dark(r: int, c: int) -> bool:
+    return (r + c) % 2 == 1
+
+def rc_to_alg(r: int, c: int) -> str:
+    # a1 bottom-left => internal (7,0)
+    file_ = chr(ord("a") + c)
+    rank = str(8 - r)
+    return f"{file_}{rank}"
+
+def alg_to_rc(s: str) -> Tuple[int, int]:
+    s = s.strip().lower()
+    c = ord(s[0]) - ord("a")
+    r = 8 - int(s[1])
+    return r, c
+
+def move_seq_to_str(seq: List[Tuple[int, int]]) -> str:
+    return "-".join(rc_to_alg(r, c) for r, c in seq)
+
+def move_str_to_seq(s: str) -> List[Tuple[int, int]]:
+    parts = [p.strip() for p in s.split("-") if p.strip()]
+    return [alg_to_rc(p) for p in parts]
+
+def piece_color(p: str) -> Optional[str]:
+    if p in ("w", "W"):
+        return "w"
+    if p in ("b", "B"):
+        return "b"
+    return None
+
+def is_king(p: str) -> bool:
+    return p in ("W", "B")
+
+
+@dataclass
+class GameState:
+    board: List[List[str]]
+    turn: str  # "w" user, "b" bot
+    history: List[str]
+    last_analysis: str
+
+
+def initial_board() -> List[List[str]]:
+    b = [["." for _ in range(8)] for _ in range(8)]
+    # Black at top rows 0-2 on dark squares
+    for r in range(0, 3):
+        for c in range(8):
+            if is_dark(r, c):
+                b[r][c] = "b"
+    # White at bottom rows 5-7 on dark squares
+    for r in range(5, 8):
+        for c in range(8):
+            if is_dark(r, c):
+                b[r][c] = "w"
+    return b
+
+def clone_board(board: List[List[str]]) -> List[List[str]]:
+    return [row[:] for row in board]
+
+def board_to_ascii(board: List[List[str]]) -> str:
+    # compact representation for prompting
+    lines = []
+    for r in range(8):
+        lines.append("".join(board[r]))
+    return "\n".join(lines)
+
+def count_material(board: List[List[str]]) -> Dict[str, float]:
+    score = {"w": 0.0, "b": 0.0}
+    for r in range(8):
+        for c in range(8):
+            p = board[r][c]
+            if p == "w":
+                score["w"] += 1.0
+            elif p == "W":
+                score["w"] += 1.6
+            elif p == "b":
+                score["b"] += 1.0
+            elif p == "B":
+                score["b"] += 1.6
+    return score
+
+def promote_if_needed(p: str, r: int) -> str:
+    if p == "w" and r == 0:
+        return "W"
+    if p == "b" and r == 7:
+        return "B"
+    return p
+
+
+# ----------------------------
+# Move generation
+# ----------------------------
+def move_dirs(p: str) -> List[Tuple[int, int]]:
+    # movement directions (step)
+    if p == "w":
+        return [(-1, -1), (-1, +1)]
+    if p == "b":
+        return [(+1, -1), (+1, +1)]
+    # kings
+    if p in ("W", "B"):
+        return [(-1, -1), (-1, +1), (+1, -1), (+1, +1)]
+    return []
+
+def capture_dirs(p: str) -> List[Tuple[int, int]]:
+    # English draughts: men capture forward only; kings both ways
+    return move_dirs(p)
+
+def gen_simple_moves(board: List[List[str]], color: str) -> List[List[Tuple[int, int]]]:
+    moves = []
+    for r in range(8):
+        for c in range(8):
+            p = board[r][c]
+            if piece_color(p) != color:
+                continue
+            for dr, dc in move_dirs(p):
+                r2, c2 = r + dr, c + dc
+                if inside(r2, c2) and board[r2][c2] == ".":
+                    moves.append([(r, c), (r2, c2)])
+    return moves
+
+def gen_captures_from(board: List[List[str]], r: int, c: int, p: str) -> List[List[Tuple[int, int]]]:
+    """
+    Returns capture sequences starting at (r,c), including start and landings.
+    If man reaches king row during capture, we stop (promotion at end of move).
+    """
+    color = piece_color(p)
+    assert color in ("w", "b")
+
+    sequences = []
+    found_any = False
+
+    for dr, dc in capture_dirs(p):
+        r_mid, c_mid = r + dr, c + dc
+        r2, c2 = r + 2 * dr, c + 2 * dc
+        if not (inside(r2, c2) and inside(r_mid, c_mid)):
+            continue
+        mid_piece = board[r_mid][c_mid]
+        if mid_piece == ".":
+            continue
+        if piece_color(mid_piece) == color:
+            continue
+        if board[r2][c2] != ".":
+            continue
 
-DEFAULT_DPI = 200
-MAX_SIDE = int(os.getenv("MAX_SIDE", "1600"))  # soft cap for CPU speed
+        # perform capture on a cloned board
+        nb = clone_board(board)
+        nb[r][c] = "."
+        nb[r_mid][c_mid] = "."
+        nb[r2][c2] = p  # promotion deferred
 
-# -------------------------
-# Models (CPU)
-# -------------------------
-trocr_processor = TrOCRProcessor.from_pretrained(TROCR_NAME)
-trocr_model = VisionEncoderDecoderModel.from_pretrained(TROCR_NAME).eval().to(DEVICE)
+        # stop extending if this is a man that reaches king row
+        if (p == "w" and r2 == 0) or (p == "b" and r2 == 7):
+            sequences.append([(r, c), (r2, c2)])
+            found_any = True
+            continue
 
-blip_processor = BlipProcessor.from_pretrained(BLIP_NAME)
-blip_model = BlipForConditionalGeneration.from_pretrained(BLIP_NAME).eval().to(DEVICE)
+        tails = gen_captures_from(nb, r2, c2, p)
+        if tails:
+            for t in tails:
+                sequences.append([(r, c)] + t[1:])
+            found_any = True
+        else:
+            sequences.append([(r, c), (r2, c2)])
+            found_any = True
+
+    return sequences if found_any else []
+
+def gen_legal_moves(board: List[List[str]], color: str) -> List[List[Tuple[int, int]]]:
+    captures = []
+    for r in range(8):
+        for c in range(8):
+            p = board[r][c]
+            if piece_color(p) != color:
+                continue
+            caps = gen_captures_from(board, r, c, p)
+            captures.extend(caps)
+
+    # forced capture rule
+    if captures:
+        # remove duplicates (can arise via different recursion paths)
+        uniq = {}
+        for seq in captures:
+            key = tuple(seq)
+            uniq[key] = seq
+        return list(uniq.values())
+
+    return gen_simple_moves(board, color)
+
+def apply_move(board: List[List[str]], seq: List[Tuple[int, int]]) -> List[List[str]]:
+    nb = clone_board(board)
+    (r0, c0) = seq[0]
+    p = nb[r0][c0]
+    nb[r0][c0] = "."
+
+    for i in range(1, len(seq)):
+        (r1, c1) = seq[i - 1]
+        (r2, c2) = seq[i]
+        # capture if jump
+        if abs(r2 - r1) == 2 and abs(c2 - c1) == 2:
+            rm = (r1 + r2) // 2
+            cm = (c1 + c2) // 2
+            nb[rm][cm] = "."
+
+    (rf, cf) = seq[-1]
+    p2 = promote_if_needed(p, rf)
+    nb[rf][cf] = p2
+    return nb
+
+def winner(board: List[List[str]]) -> Optional[str]:
+    # winner if opponent has no pieces or no moves
+    w_cnt = 0
+    b_cnt = 0
+    for r in range(8):
+        for c in range(8):
+            if board[r][c] in ("w", "W"):
+                w_cnt += 1
+            elif board[r][c] in ("b", "B"):
+                b_cnt += 1
+    if w_cnt == 0:
+        return "b"
+    if b_cnt == 0:
+        return "w"
+    if not gen_legal_moves(board, "w"):
+        return "b"
+    if not gen_legal_moves(board, "b"):
+        return "w"
+    return None
+
+
+# ============================================================
+# Simple engine for analysis (not a transformer):
+# minimax on material + mobility, small depth for CPU.
+# ============================================================
+def eval_board(board: List[List[str]]) -> float:
+    m = count_material(board)
+    # positive => good for white
+    score = (m["w"] - m["b"])
+    # mobility bonus
+    score += 0.04 * (len(gen_legal_moves(board, "w")) - len(gen_legal_moves(board, "b")))
+    return score
+
+def minimax(board: List[List[str]], color: str, depth: int, alpha: float, beta: float) -> Tuple[float, Optional[List[Tuple[int, int]]]]:
+    win = winner(board)
+    if win == "w":
+        return 10_000.0, None
+    if win == "b":
+        return -10_000.0, None
+
+    if depth == 0:
+        return eval_board(board), None
+
+    moves = gen_legal_moves(board, color)
+    if not moves:
+        # no moves => lose
+        return (-10_000.0 if color == "w" else 10_000.0), None
+
+    best_move = None
+
+    if color == "w":
+        best = -math.inf
+        for mv in moves:
+            nb = apply_move(board, mv)
+            val, _ = minimax(nb, "b", depth - 1, alpha, beta)
+            if val > best:
+                best = val
+                best_move = mv
+            alpha = max(alpha, best)
+            if beta <= alpha:
+                break
+        return best, best_move
+    else:
+        best = math.inf
+        for mv in moves:
+            nb = apply_move(board, mv)
+            val, _ = minimax(nb, "w", depth - 1, alpha, beta)
+            if val < best:
+                best = val
+                best_move = mv
+            beta = min(beta, best)
+            if beta <= alpha:
+                break
+        return best, best_move
+
+
+# ============================================================
+# Embeddings (transformer #3) for retrieving tips
+# ============================================================
+TIPS = [
+    "Всегда проверяй обязательный бой: если есть взятие, обычный ход запрещён.",
+    "Старайся сохранять дамочную линию: не открывай край без причины.",
+    "Не меняйся, если это приводит к потере темпа и отдаёт центр.",
+    "Центр важен: контроль диагоналей увеличивает мобильность и шансы на многоходовые взятия.",
+    "Перед ходом оцени ответ соперника: что он берёт или чем отвечает на диагонали?",
+    "Если видишь возможность мультибоя, считай траекторию до конца — важно, где ты остановишься.",
+    "Дамка сильнее: иногда стоит пожертвовать шашку ради прохода в дамки.",
+    "Не оставляй одиночные шашки без поддержки — их легко поймать взятием.",
+    "Думай про 'вилку' (двойную угрозу) и про то, чтобы не подставлять шашку под обязательный бой.",
+]
+
+@torch.no_grad()
+def embed_text(text: str) -> torch.Tensor:
+    toks = emb_tok(text, return_tensors="pt", truncation=True, max_length=128, padding=True)
+    toks = {k: v.to(DEVICE) for k, v in toks.items()}
+    out = emb_model(**toks)
+    # mean pooling
+    last = out.last_hidden_state  # [B,T,H]
+    mask = toks["attention_mask"].unsqueeze(-1)  # [B,T,1]
+    pooled = (last * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+    pooled = torch.nn.functional.normalize(pooled, p=2, dim=1)
+    return pooled[0].cpu()
+
+TIP_EMBS = torch.stack([embed_text(t) for t in TIPS], dim=0)  # [N,H]
+
+def retrieve_tips(query: str, k: int = 3) -> List[str]:
+    q = embed_text(query)
+    sims = (TIP_EMBS @ q.unsqueeze(1)).squeeze(1)  # [N]
+    top = torch.topk(sims, k=min(k, len(TIPS))).indices.tolist()
+    return [TIPS[i] for i in top]
+
+
+# ============================================================
+# Coach (transformer #1): generates explanation/feedback
+# ============================================================
+@torch.no_grad()
+def coach_generate(prompt: str, max_new_tokens: int = 160) -> str:
+    inp = coach_tok(prompt, return_tensors="pt", truncation=True, max_length=512)
+    inp = {k: v.to(DEVICE) for k, v in inp.items()}
+    out = coach_model.generate(
+        **inp,
+        max_new_tokens=max_new_tokens,
+        do_sample=False,
+        num_beams=1,
+    )
+    text = coach_tok.decode(out[0], skip_special_tokens=True)
+    return text.strip()
+
+
+# ============================================================
+# Opponent (transformer #2): chooses a legal move
+# ============================================================
+@torch.no_grad()
+def opponent_choose_move(board: List[List[str]], legal_moves: List[str]) -> str:
+    # distilgpt2 is not instruction-tuned, so we keep it extremely constrained and parse output.
+    board_ascii = board_to_ascii(board)
+    moves_block = "\n".join([f"- {m}" for m in legal_moves[:40]])  # cap list
+    prompt = (
+        "You are playing checkers as Black.\n"
+        "Choose ONE move exactly from the list. Output only that move.\n"
+        f"Board:\n{board_ascii}\n"
+        f"Moves:\n{moves_block}\n"
+        "Move:"
+    )
+    inp = opp_tok(prompt, return_tensors="pt", truncation=True, max_length=512)
+    inp = {k: v.to(DEVICE) for k, v in inp.items()}
+    gen = opp_model.generate(
+        **inp,
+        max_new_tokens=24,
+        do_sample=True,
+        top_p=0.85,
+        temperature=0.7,
+        pad_token_id=opp_tok.eos_token_id,
+    )
+    out = opp_tok.decode(gen[0], skip_special_tokens=True)
+    tail = out.split("Move:")[-1].strip()
+
+    # parse: pick the first legal move that appears in the generated tail
+    for m in legal_moves:
+        if m in tail:
+            return m
+
+    # fallback: try extract token pattern like a1-b2
+    cand = re.findall(r"[a-h][1-8](?:-[a-h][1-8])+", tail.lower())
+    if cand:
+        for c in cand:
+            if c in legal_moves:
+                return c
+
+    # final fallback: random legal
+    return random.choice(legal_moves)
+
+
+# ============================================================
+# Rendering board
+# ============================================================
+def render_board(board: List[List[str]], size: int = 520) -> Image.Image:
+    pad = 20
+    cell = (size - 2 * pad) // 8
+    img = Image.new("RGB", (size, size), (245, 245, 245))
+    d = ImageDraw.Draw(img)
+
+    dark = (150, 110, 80)
+    light = (235, 220, 200)
+
+    # grid
+    for r in range(8):
+        for c in range(8):
+            x0 = pad + c * cell
+            y0 = pad + r * cell
+            x1 = x0 + cell
+            y1 = y0 + cell
+            d.rectangle([x0, y0, x1, y1], fill=(dark if is_dark(r, c) else light))
+
+    # pieces
+    for r in range(8):
+        for c in range(8):
+            p = board[r][c]
+            if p == ".":
+                continue
+            cx = pad + c * cell + cell // 2
+            cy = pad + r * cell + cell // 2
+            rad = int(cell * 0.38)
 
-# -------------------------
-# Optional: Tesseract for image boxes
-# -------------------------
-def _try_import_tesseract():
-    try:
-        import pytesseract  # type: ignore
-        _ = pytesseract.get_tesseract_version()
-        return pytesseract
-    except Exception:
-        return None
+            if p in ("w", "W"):
+                fill = (245, 245, 245)
+                outline = (30, 30, 30)
+            else:
+                fill = (40, 40, 40)
+                outline = (230, 230, 230)
 
-PYTESS = _try_import_tesseract()
+            d.ellipse([cx - rad, cy - rad, cx + rad, cy + rad], fill=fill, outline=outline, width=3)
 
-TASKS = ["OCR", "Markdown", "Locate", "Describe"]
+            if is_king(p):
+                # crown marker
+                d.ellipse([cx - rad // 2, cy - rad // 2, cx + rad // 2, cy + rad // 2], outline=(255, 215, 0), width=4)
 
+    # coordinates
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
+    except Exception:
+        font = None
 
-# -------------------------
-# Helpers
-# -------------------------
-def _to_rgb(img: Image.Image) -> Image.Image:
-    if img.mode in ("RGBA", "LA", "P"):
-        img = img.convert("RGB")
-    img = ImageOps.exif_transpose(img)
+    for c in range(8):
+        d.text((pad + c * cell + 3, pad + 8 * cell + 2), chr(ord("a") + c), fill=(30, 30, 30), font=font)
+    for r in range(8):
+        d.text((3, pad + r * cell + 3), str(8 - r), fill=(30, 30, 30), font=font)
 
-    # Keep CPU inference reasonable
-    w, h = img.size
-    m = max(w, h)
-    if m > MAX_SIDE:
-        scale = MAX_SIDE / float(m)
-        img = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
     return img
 
 
-def _tokenize(s: str) -> List[str]:
-    return re.findall(r"[A-Za-zА-Яа-я0-9]+", (s or "").lower())
-
-
-def trocr_ocr(img: Image.Image) -> str:
-    img = _to_rgb(img)
-    inputs = trocr_processor(images=img, return_tensors="pt")
-    pixel_values = inputs.pixel_values.to(DEVICE)
-    with torch.no_grad():
-        ids = trocr_model.generate(pixel_values, max_new_tokens=256)
-    text = trocr_processor.batch_decode(ids, skip_special_tokens=True)[0]
-    return (text or "").strip()
-
-
-def blip_describe(img: Image.Image) -> str:
-    img = _to_rgb(img)
-    inputs = blip_processor(images=img, return_tensors="pt").to(DEVICE)
-    with torch.no_grad():
-        out = blip_model.generate(**inputs, max_new_tokens=80)
-    return blip_processor.decode(out[0], skip_special_tokens=True).strip()
-
-
-def render_pdf_page(path: str, page_num: int, dpi: int = DEFAULT_DPI):
-    doc = fitz.open(path)
-    page_idx = max(0, min(int(page_num) - 1, len(doc) - 1))
-    page = doc.load_page(page_idx)
-    zoom = dpi / 72.0
-    pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
-    img = Image.open(BytesIO(pix.tobytes("png")))
-    return doc, page, _to_rgb(img), zoom
-
-
-def pdf_has_text(page: fitz.Page) -> bool:
-    return bool(page.get_text("words"))
+# ============================================================
+# Game logic wrapper
+# ============================================================
+def new_game() -> GameState:
+    return GameState(
+        board=initial_board(),
+        turn="w",
+        history=[],
+        last_analysis="",
+    )
+
+def legal_moves_str(board: List[List[str]], color: str) -> List[str]:
+    moves = gen_legal_moves(board, color)
+    ms = [move_seq_to_str(mv) for mv in moves]
+    # stable ordering: captures first (longer sequences first), then lexicographic
+    ms.sort(key=lambda s: (-s.count("-"), s))
+    return ms
+
+def analyze_user_move(board_before: List[List[str]], user_move_str: str) -> str:
+    # engine "best move" as baseline (not a transformer)
+    depth = int(os.getenv("ANALYSIS_DEPTH", "3"))
+    best_val, best_mv = minimax(board_before, "w", depth=depth, alpha=-math.inf, beta=math.inf)
+    best_str = move_seq_to_str(best_mv) if best_mv else "(none)"
+
+    tips = retrieve_tips("шашки: как улучшить ход и не подставиться", k=3)
+
+    prompt = (
+        "Ты тренер по шашкам. Коротко и по делу.\n"
+        f"Ход игрока: {user_move_str}\n"
+        f"Рекомендованный ход (по анализу): {best_str}\n"
+        "Дай объяснение: почему рекомендованный лучше, и какая ошибка/риск в ходе игрока.\n"
+        "Добавь 2-3 практических совета.\n"
+        "Подсказки:\n"
+        + "\n".join(f"- {t}" for t in tips)
+    )
+    return coach_generate(prompt, max_new_tokens=180)
+
+
+def step_user_and_bot(state: GameState, user_move: str) -> Tuple[GameState, str]:
+    if winner(state.board) is not None:
+        return state, "Game already finished."
+
+    if state.turn != "w":
+        return state, "Not your turn."
+
+    leg = legal_moves_str(state.board, "w")
+    if user_move not in leg:
+        return state, "Invalid move (not in legal list)."
+
+    board_before = clone_board(state.board)
+    seq = move_str_to_seq(user_move)
+    state.board = apply_move(state.board, seq)
+    state.history.append(f"White: {user_move}")
+    state.turn = "b"
+
+    # analysis (coach transformer)
+    state.last_analysis = analyze_user_move(board_before, user_move)
+
+    win = winner(state.board)
+    if win is not None:
+        state.history.append("Result: " + ("White wins" if win == "w" else "Black wins"))
+        return state, ("White wins." if win == "w" else "Black wins.")
+
+    # bot move
+    bot_leg = legal_moves_str(state.board, "b")
+    if not bot_leg:
+        state.history.append("Result: White wins")
+        return state, "White wins."
+
+    bot_move = opponent_choose_move(state.board, bot_leg)
+    bot_seq = move_str_to_seq(bot_move)
+    state.board = apply_move(state.board, bot_seq)
+    state.history.append(f"Black: {bot_move}")
+    state.turn = "w"
+
+    win = winner(state.board)
+    if win is not None:
+        state.history.append("Result: " + ("White wins" if win == "w" else "Black wins"))
+        return state, ("White wins." if win == "w" else "Black wins.")
+
+    return state, f"Bot played: {bot_move}"
+
+
+# ============================================================
+# Coach chat (transformer #1 + embeddings #3)
+# ============================================================
+def coach_chat(state: GameState, message: str, chat_hist: List[Tuple[str, str]]):
+    msg = (message or "").strip()
+    if not msg:
+        return chat_hist, ""
+
+    # Retrieve tips relevant to the question
+    tips = retrieve_tips(msg, k=3)
+
+    # Provide board context
+    context = board_to_ascii(state.board)
+    last = state.history[-6:] if state.history else []
+
+    prompt = (
+        "Ты тренер по шашкам. Отвечай кратко, но конкретно.\n"
+        f"Вопрос игрока: {msg}\n"
+        "Контекст партии (последние ходы):\n"
+        + ("\n".join(last) if last else "(нет)")
+        + "\n"
+        "Доска (ASCII):\n"
+        + context
+        + "\n"
+        "Полезные подсказки:\n"
+        + "\n".join(f"- {t}" for t in tips)
+        + "\n"
+        "Ответ:"
+    )
+
+    answer = coach_generate(prompt, max_new_tokens=180)
+    chat_hist = chat_hist + [(msg, answer)]
+    return chat_hist, ""
+
+
+# ============================================================
+# UI
+# ============================================================
+theme = gr.themes.Monochrome(font=[gr.themes.GoogleFont("Inter"), "system-ui"])
+
+with gr.Blocks(theme=theme, title="Checkers Coach (CPU, 3 Transformers)") as demo:
+    state = gr.State(new_game())
 
+    with gr.Row():
+        with gr.Column(scale=1, min_width=360):
+            board_img = gr.Image(label="Board", type="pil", height=520)
+            status = gr.Textbox(label="Status", value="", interactive=False)
 
-def pdf_extract_text(page: fitz.Page) -> str:
-    return (page.get_text("text") or "").strip()
+            move_dd = gr.Dropdown(label="Your move (White)", choices=[], value=None)
+            play_btn = gr.Button("Play move", variant="primary")
+            new_btn = gr.Button("New game")
 
+            analysis = gr.Textbox(label="Coach analysis", lines=10, interactive=False)
 
-def pdf_to_markdown_simple(page: fitz.Page) -> str:
-    data = page.get_text("dict")
-    spans = []
-    for b in data.get("blocks", []):
-        for ln in b.get("lines", []):
-            for sp in ln.get("spans", []):
-                t = (sp.get("text") or "").strip()
-                if t:
-                    spans.append(float(sp.get("size", 0.0)))
-    if not spans:
-        return ""
+        with gr.Column(scale=1, min_width=360):
+            hist = gr.Markdown("")
+            gr.Markdown("### Coach chat")
+            chat = gr.Chatbot(height=360)
+            msg = gr.Textbox(label="Message", placeholder="Ask about strategy, mistakes, next plan…")
+            send = gr.Button("Send")
 
-    med = float(np.median(spans))
-    h1_thr = med * 1.60
-    h2_thr = med * 1.35
+    def refresh_ui(gs: GameState):
+        img = render_board(gs.board)
+        leg = legal_moves_str(gs.board, "w") if winner(gs.board) is None else []
+        h = "### History\n" + ("\n".join([f"- {x}" for x in gs.history]) if gs.history else "- (empty)")
+        return img, ("" if gs.turn == "w" else "Bot thinking / waiting…"), gr.update(choices=leg, value=(leg[0] if leg else None)), gs.last_analysis, h
 
-    out_lines: List[str] = []
-    for b in data.get("blocks", []):
-        if b.get("type") != 0:
-            continue
-        for ln in b.get("lines", []):
-            parts = []
-            sizes = []
-            for sp in ln.get("spans", []):
-                t = (sp.get("text") or "").strip()
-                if t:
-                    parts.append(t)
-                    sizes.append(float(sp.get("size", 0.0)))
-            if not parts:
-                continue
-            line = " ".join(parts).strip()
-            sz = max(sizes) if sizes else med
-            if sz >= h1_thr:
-                out_lines.append("# " + line)
-            elif sz >= h2_thr:
-                out_lines.append("## " + line)
-            else:
-                out_lines.append(line)
-        out_lines.append("")
-    return "\n".join(out_lines).strip()
-
-
-def draw_rects(img: Image.Image, rects_px: List[Tuple[int, int, int, int]]) -> Image.Image:
-    out = img.copy()
-    draw = ImageDraw.Draw(out)
-    overlay = Image.new("RGBA", out.size, (0, 0, 0, 0))
-    draw2 = ImageDraw.Draw(overlay)
-    for (x0, y0, x1, y1) in rects_px:
-        draw.rectangle([x0, y0, x1, y1], outline=(0, 160, 255), width=3)
-        draw2.rectangle([x0, y0, x1, y1], fill=(0, 160, 255, 60))
-    out.paste(overlay, (0, 0), overlay)
-    return out
-
-
-def locate_in_pdf_words(page: fitz.Page, query: str) -> List[Tuple[float, float, float, float]]:
-    q = _tokenize(query)
-    if not q:
-        return []
-    words = page.get_text("words")
-    if not words:
-        return []
-
-    w_tokens = []
-    for w in words:
-        toks = _tokenize(w[4])
-        w_tokens.append(toks[0] if toks else "")
-
-    rects = []
-    n, m = len(w_tokens), len(q)
-    for i in range(0, n - m + 1):
-        if w_tokens[i:i + m] == q:
-            xs0 = [float(words[j][0]) for j in range(i, i + m)]
-            ys0 = [float(words[j][1]) for j in range(i, i + m)]
-            xs1 = [float(words[j][2]) for j in range(i, i + m)]
-            ys1 = [float(words[j][3]) for j in range(i, i + m)]
-            rects.append((min(xs0), min(ys0), max(xs1), max(ys1)))
-    return rects
-
-
-def locate_in_image_tesseract(img: Image.Image, query: str):
-    if PYTESS is None:
-        return [], "Tesseract not available."
-    q = _tokenize(query)
-    if not q:
-        return [], "Empty query."
-
-    img = _to_rgb(img)
-    data = PYTESS.image_to_data(img, output_type=PYTESS.Output.DICT)
-
-    texts = data.get("text", [])
-    left = data.get("left", [])
-    top = data.get("top", [])
-    width = data.get("width", [])
-    height = data.get("height", [])
-    conf = data.get("conf", [])
-
-    tokens = []
-    boxes = []
-    for i, t in enumerate(texts):
-        t = (t or "").strip()
-        if not t:
-            continue
-        toks = _tokenize(t)
-        if not toks:
-            continue
-        try:
-            c = float(conf[i])
-            if c < 0:
-                continue
-        except Exception:
-            pass
-        tokens.append(toks[0])
-        boxes.append((int(left[i]), int(top[i]), int(left[i] + width[i]), int(top[i] + height[i])))
-
-    rects_px = []
-    n, m = len(tokens), len(q)
-    for i in range(0, n - m + 1):
-        if tokens[i:i + m] == q:
-            xs0 = [boxes[j][0] for j in range(i, i + m)]
-            ys0 = [boxes[j][1] for j in range(i, i + m)]
-            xs1 = [boxes[j][2] for j in range(i, i + m)]
-            ys1 = [boxes[j][3] for j in range(i, i + m)]
-            rects_px.append((min(xs0), min(ys0), max(xs1), max(ys1)))
-
-    return rects_px, ("Found." if rects_px else "Not found.")
-
-
-def as_text_block(s: str) -> str:
-    s = (s or "").strip()
-    return s if s else ""
-
-
-# -------------------------
-# Core processing
-# -------------------------
-def process(file_path: str, task: str, page_num: int, query: str):
-    if not file_path:
-        return "Upload a file.", "", None, None
-
-    ext = os.path.splitext(file_path)[1].lower()
-
-    # PDF
-    if ext == ".pdf":
-        doc, page, page_img, zoom = render_pdf_page(file_path, page_num, dpi=DEFAULT_DPI)
-        try:
-            preview = page_img
-
-            if task == "Describe":
-                cap = blip_describe(page_img)
-                return cap, cap, None, preview
-
-            if task == "OCR":
-                txt = pdf_extract_text(page) if pdf_has_text(page) else trocr_ocr(page_img)
-                return txt, txt, None, preview
-
-            if task == "Markdown":
-                if pdf_has_text(page):
-                    md = pdf_to_markdown_simple(page)
-                    if not md:
-                        md = pdf_extract_text(page)
-                else:
-                    md = trocr_ocr(page_img)
-                return md, md, None, preview
-
-            if task == "Locate":
-                if not (query or "").strip():
-                    return "Enter query.", "", preview, preview
-
-                # selectable-text PDF: precise boxes
-                rects_pdf = locate_in_pdf_words(page, query)
-                if rects_pdf:
-                    rects_px = [(int(x0 * zoom), int(y0 * zoom), int(x1 * zoom), int(y1 * zoom)) for x0, y0, x1, y1 in rects_pdf]
-                    boxed = draw_rects(page_img, rects_px)
-                    return "Found.", "", boxed, preview
-
-                # fallback: render + tesseract
-                rects_px, msg = locate_in_image_tesseract(page_img, query)
-                boxed = draw_rects(page_img, rects_px) if rects_px else page_img
-                return msg, "", boxed, preview
-
-            return "Unknown task.", "", None, preview
-        finally:
-            doc.close()
-
-    # Image
-    img = _to_rgb(Image.open(file_path))
-    preview = img
-
-    if task == "Describe":
-        cap = blip_describe(img)
-        return cap, cap, None, preview
-
-    if task == "OCR":
-        txt = trocr_ocr(img)
-        return txt, txt, None, preview
-
-    if task == "Markdown":
-        md = trocr_ocr(img)
-        return md, md, None, preview
-
-    if task == "Locate":
-        if not (query or "").strip():
-            return "Enter query.", "", img, preview
-        rects_px, msg = locate_in_image_tesseract(img, query)
-        boxed = draw_rects(img, rects_px) if rects_px else img
-        return msg, "", boxed, preview
-
-    return "Unknown task.", "", None, preview
-
-
-# -------------------------
-# UI wiring
-# -------------------------
-def update_page_ui(file_path: str):
-    if not file_path:
-        return gr.update(visible=False), None
-
-    ext = os.path.splitext(file_path)[1].lower()
-    if ext != ".pdf":
-        return gr.update(visible=False), _to_rgb(Image.open(file_path))
-
-    doc = fitz.open(file_path)
-    pages = max(1, len(doc))
-    doc.close()
-
-    _, _, img, _ = render_pdf_page(file_path, 1, dpi=DEFAULT_DPI)
-    return gr.update(visible=True, minimum=1, maximum=pages, value=1), img
-
-
-def update_preview(file_path: str, page_num: int):
-    if not file_path:
-        return None
-    ext = os.path.splitext(file_path)[1].lower()
-    if ext != ".pdf":
-        return _to_rgb(Image.open(file_path))
-    _, _, img, _ = render_pdf_page(file_path, int(page_num), dpi=DEFAULT_DPI)
-    return img
+    def on_new():
+        gs = new_game()
+        return (gs, ) + refresh_ui(gs) + ([], "")
 
+    def on_play(gs: GameState, mv: str):
+        gs, st = step_user_and_bot(gs, mv or "")
+        img, _, dd, an, h = refresh_ui(gs)
+        return gs, img, st, dd, an, h
 
-def toggle_query(task: str):
-    return gr.update(visible=(task == "Locate"))
+    def on_send(gs: GameState, m: str, ch: List[Tuple[str, str]]):
+        ch, cleared = coach_chat(gs, m, ch or [])
+        return ch, cleared
 
+    demo.load(lambda gs: refresh_ui(gs), inputs=[state], outputs=[board_img, status, move_dd, analysis, hist])
 
-# -------------------------
-# Minimal UI style
-# -------------------------
-theme = gr.themes.Monochrome(
-    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"]
-)
+    new_btn.click(on_new, inputs=[], outputs=[state, board_img, status, move_dd, analysis, hist, chat, msg])
+    play_btn.click(on_play, inputs=[state, move_dd], outputs=[state, board_img, status, move_dd, analysis, hist])
 
-with gr.Blocks(theme=theme, title="Doc Tool (CPU)") as demo:
-    with gr.Row():
-        with gr.Column(scale=1, min_width=320):
-            file_in = gr.File(label="File", file_types=["image", ".pdf"], type="filepath")
-            page = gr.Slider(label="Page", minimum=1, maximum=1, value=1, step=1, visible=False)
-            task = gr.Dropdown(label="Task", choices=TASKS, value="OCR")
-            query = gr.Textbox(label="Query", placeholder="Text to locate", visible=False)
-            run_btn = gr.Button("Run", variant="primary")
-
-        with gr.Column(scale=2):
-            with gr.Row():
-                preview = gr.Image(label="Preview", type="pil", height=320)
-                boxes = gr.Image(label="Boxes", type="pil", height=320)
-            out = gr.Textbox(label="Output", lines=10)
-
-    file_in.change(update_page_ui, inputs=[file_in], outputs=[page, preview])
-    page.change(update_preview, inputs=[file_in, page], outputs=[preview])
-    task.change(toggle_query, inputs=[task], outputs=[query])
-
-    def on_run(fp, t, p, q):
-        text, _, boxed, prev = process(fp, t, int(p), q or "")
-        # keep preview stable; boxes only when relevant
-        return prev, boxed, as_text_block(text)
-
-    run_btn.click(on_run, inputs=[file_in, task, page, query], outputs=[preview, boxes, out])
+    send.click(on_send, inputs=[state, msg, chat], outputs=[chat, msg])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+    demo.queue(max_size=32).launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)