예제 소스

recorder.py

#!/usr/bin/env python3
"""모노몹 장기 대국을 화면 캡처로 읽어 SQLite에 기보로 저장 (기록 전용, 탭/추천 없음).

  python recorder.py                 지금 화면부터 기록 시작 (대국 중간이면 현재 국면부터)
  python recorder.py --start-ply 18  중간부터 기록할 때 이미 둔 수 (표시용)
  python recorder.py list            저장된 대국 목록
  python recorder.py show ID         대국 수순 출력

저장 위치: data/monomob_kibo.sqlite
좌표/수 표기는 엔진과 같음: "ffrr" = from.file from.rank to.file to.rank (rank 0 = 한 궁성 쪽)
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path

import cv2
import numpy as np

import bot
import rules
from bot import CHO, EMPTY, HAN, PIECE_KO, SIDE_KO, View

ROOT = Path(__file__).resolve().parent
DB_PATH = ROOT / "data" / "monomob_kibo.sqlite"
FULL_TEMPLATES = ROOT / "templates_full"


# ---------------------------------------------------------------- DB

def db() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.executescript(
        """
        CREATE TABLE IF NOT EXISTS games (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          source TEXT NOT NULL DEFAULT 'monomob',
          started_at TEXT NOT NULL,
          ended_at TEXT,
          from_start INTEGER NOT NULL,
          start_ply INTEGER NOT NULL DEFAULT 0,
          start_board TEXT NOT NULL,
          flipped INTEGER NOT NULL,
          my_side TEXT,
          move_count INTEGER NOT NULL DEFAULT 0,
          end_reason TEXT,
          note TEXT
        );
        CREATE TABLE IF NOT EXISTS moves (
          game_id INTEGER NOT NULL,
          ply INTEGER NOT NULL,
          side TEXT NOT NULL,
          move TEXT NOT NULL,
          piece TEXT,
          captured TEXT,
          board_after TEXT NOT NULL,
          created_at TEXT NOT NULL,
          PRIMARY KEY (game_id, ply)
        );
        """
    )
    return conn


def now() -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S")


# ---------------------------------------------------------------- 판 전체 인식

_tpl: dict[tuple[str, str], list[np.ndarray]] | None = None


def full_templates() -> dict[tuple[str, str], list[np.ndarray]]:
    global _tpl
    if _tpl is None:
        _tpl = defaultdict(list)
        for path in FULL_TEMPLATES.glob("*.png"):
            side, kind, _ = path.stem.split("_")
            _tpl[(CHO if side == "cho" else HAN, kind)].append(cv2.imread(str(path), cv2.IMREAD_GRAYSCALE))
        if not _tpl:
            sys.exit("templates_full/ 이 비어 있습니다")
    return _tpl


def read_board(img: np.ndarray) -> tuple[View, list[str]] | None:
    """화면 → (방향, 엔진 좌표 기물 배열 'cR' 등)."""
    screen_occ = bot.classify_screen(img)
    tpl = full_templates()
    screen_board = [""] * 90
    for i, o in enumerate(screen_occ):
        if o == EMPTY:
            continue
        patch = bot.patch_gray(img, i % 9, i // 9)
        kind = max((bot.match_score(patch, v), k) for (s, k), v in tpl.items() if s == o)[1]
        screen_board[i] = o + kind
    kings = {p[0]: i // 9 for i, p in enumerate(screen_board) if p and p[1] == "K"}
    if CHO not in kings or HAN not in kings:
        return None
    view = View(flipped=kings[CHO] <= 2)
    board = [""] * 90
    for rank in range(10):
        for file in range(9):
            col, row = view.to_screen(file, rank)
            board[rank * 9 + file] = screen_board[row * 9 + col]
    return view, board


def occ_of(board: list[str]) -> list[str]:
    return [p[0] if p else EMPTY for p in board]


def sq(i: int) -> str:
    return f"{i % 9}{i // 9}"


def label(i: int) -> str:
    return f"{i % 9 + 1}{i // 9 + 1}"


# ---------------------------------------------------------------- 기록 루프

class Recorder:
    def __init__(self, conn: sqlite3.Connection, my_side: str | None, source: str = "monomob"):
        self.conn = conn
        self.my_side = my_side
        self.source = source
        self.game_id: int | None = None
        self.view: View | None = None
        self.board: list[str] = []
        self.ply = 0
        self.last_side: str | None = None

    def open_game(self, view: View, board: list[str], from_start: bool, start_ply: int) -> None:
        self.close_game("new_game")
        self.view, self.board = view, board
        self.ply = 0
        self.last_side = None
        cur = self.conn.execute(
            "INSERT INTO games (source, started_at, from_start, start_ply, start_board, flipped, my_side)"
            " VALUES (?,?,?,?,?,?,?)",
            (self.source, now(), int(from_start), start_ply, json.dumps(board), int(view.flipped), self.my_side),
        )
        self.conn.commit()
        self.game_id = cur.lastrowid
        where = "시작 국면" if from_start else f"중간 국면 ({start_ply}수 이후)"
        print(f"[대국 #{self.game_id}] 기록 시작: {where}, 화면 위쪽={'초' if view.flipped else '한'}")

    def close_game(self, reason: str) -> None:
        if self.game_id is None:
            return
        self.conn.execute(
            "UPDATE games SET ended_at=?, move_count=?, end_reason=? WHERE id=?",
            (now(), self.ply, reason, self.game_id),
        )
        self.conn.commit()
        print(f"[대국 #{self.game_id}] 기록 종료 ({reason}, {self.ply}수)")
        self.game_id = None

    def add(self, side: str, move: str, piece: str, captured: str) -> None:
        self.ply += 1
        self.conn.execute(
            "INSERT INTO moves (game_id, ply, side, move, piece, captured, board_after, created_at) VALUES (?,?,?,?,?,?,?,?)",
            (self.game_id, self.ply, side, move, piece, captured, json.dumps(self.board), now()),
        )
        self.conn.execute("UPDATE games SET move_count=? WHERE id=?", (self.ply, self.game_id))
        self.conn.commit()
        self.last_side = side

    def apply_target(self, target: list[str]) -> bool:
        """현재 판 → target(화면에서 읽은 판)을 설명하는 수순(최대 MAX_GAP수)을 찾아 기록."""
        sides = [bot.opposite_side(self.last_side)] if self.last_side else [CHO, HAN]
        if self.last_side:
            sides.append(self.last_side)  # 상대가 한수 쉼을 한 경우
        best: list[tuple[int, int]] | None = None
        best_first = CHO
        for first in sides:
            best = explain(self.board, target, first)
            if best:
                best_first = first
                break
        if not best:
            return False
        if self.last_side and best_first == self.last_side:
            other = bot.opposite_side(best_first)
            self.add(other, "pass", "", "")
            print(f"{self.ply:3d}. {SIDE_KO[other]} 한수 쉼")
        for f, t in best:
            piece, captured = self.board[f], self.board[t]
            self.board[t], self.board[f] = piece, ""
            self.add(piece[0], sq(f) + sq(t), piece, captured)
            cap = f" x{PIECE_KO[captured[1]]}" if captured else ""
            print(f"{self.ply:3d}. {SIDE_KO[piece[0]]} {PIECE_KO[piece[1]]} {label(f)}→{label(t)}{cap}")
        return True


MAX_GAP = 4


def explain(board: list[str], target: list[str], first: str) -> list[tuple[int, int]] | None:
    """board에서 first부터 번갈아 둬서 target이 되는 최단 수순. 움직임은 바뀐 칸 안에서만 탐색."""
    if board == target:
        return []
    if sum(a != b for a, b in zip(board, target)) > MAX_GAP * 2:
        return None

    def dfs(b: list[str], side: str, depth: int) -> list[tuple[int, int]] | None:
        mism = [i for i in range(90) if b[i] != target[i]]
        if not mism:
            return []
        if depth == 0 or len(mism) > depth * 2:
            return None
        # 움직이는 말은 아직 목표와 다른 칸에서 출발해야 함
        for f in mism:
            if not b[f] or b[f][0] != side:
                continue
            for t in rules.targets(b, f):
                if b[t] and b[t][1] == "K":
                    continue
                nb = b[:]
                nb[t], nb[f] = nb[f], ""
                rest = dfs(nb, bot.opposite_side(side), depth - 1)
                if rest is not None:
                    return [(f, t), *rest]
        return None

    for depth in range(1, MAX_GAP + 1):
        seq = dfs(board, first, depth)
        if seq is not None:
            return seq
    return None


def cmd_record(args: argparse.Namespace) -> None:
    conn = db()
    rec = Recorder(conn, args.me)
    start_occ = bot.start_occupancy()
    last_obs: list[str] | None = None
    polls = 0
    fails = 0
    print("기록 대기 중... (Ctrl+C로 종료)")
    try:
        while True:
            polls += 1
            try:
                if polls % 6 == 0 and bot.is_ad_activity(bot.focused_activity()):
                    rec.close_game("ad")
                    time.sleep(2)
                    continue
                img = bot.screencap()
            except (subprocess.SubprocessError, OSError, ValueError) as e:
                print(f"  ! adb 오류 ({type(e).__name__}) - 3초 후 재시도")
                time.sleep(3)
                continue

            if rec.game_id is None:
                read = read_board(img)
                if read:
                    view, board = read
                    from_start = occ_of(board) == start_occ
                    rec.open_game(view, board, from_start, 0 if from_start else args.start_ply)
                    args.start_ply = 0
                time.sleep(1)
                continue

            assert rec.view is not None
            obs = rec.view.canonical(bot.classify_screen(img))
            occ = occ_of(rec.board)
            if obs != occ and obs == last_obs:
                read = read_board(img)
                if read and read[0].flipped == rec.view.flipped:
                    target = read[1]
                    if occ_of(target) == start_occ:
                        rec.open_game(read[0], target, True, 0)
                        fails = 0
                    elif rec.apply_target(target):
                        fails = 0
                    else:
                        fails += 1
                        if fails >= 4:
                            print("  ! 수순 복원 실패 - 현재 국면부터 다시 기록")
                            rec.close_game("resync")
                            rec.open_game(read[0], target, False, 0)
                            fails = 0
                else:
                    fails += 1
                    if fails >= 8:
                        rec.close_game("screen_changed")
                        fails = 0
            last_obs = obs
            time.sleep(0.2)
    except KeyboardInterrupt:
        rec.close_game("stopped")
        print("중단")


def cmd_list(_args: argparse.Namespace) -> None:
    conn = db()
    rows = conn.execute(
        "SELECT id, started_at, from_start, start_ply, move_count, end_reason, my_side FROM games ORDER BY id DESC"
    ).fetchall()
    for gid, started, fs, sp, mc, reason, me in rows:
        start = "처음부터" if fs else f"{sp}수 이후부터"
        print(f"#{gid}  {started}  {start}  {mc}수  종료={reason or '-'}  나={me or '-'}")


def cmd_show(args: argparse.Namespace) -> None:
    conn = db()
    for ply, side, move, piece, captured in conn.execute(
        "SELECT ply, side, move, piece, captured FROM moves WHERE game_id=? ORDER BY ply", (args.id,)
    ):
        if move == "pass":
            print(f"{ply:3d}. {SIDE_KO[side]} 한수 쉼")
            continue
        f = int(move[1]) * 9 + int(move[0])
        t = int(move[3]) * 9 + int(move[2])
        cap = f" x{PIECE_KO[captured[1]]}" if captured else ""
        print(f"{ply:3d}. {SIDE_KO[side]} {PIECE_KO[piece[1]]} {label(f)}→{label(t)}{cap}   [{move}]")


def main() -> None:
    sys.stdout.reconfigure(line_buffering=True)  # type: ignore[union-attr]
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--start-ply", type=int, default=0)
    ap.add_argument("--me", choices=["cho", "han"])
    sub = ap.add_subparsers(dest="cmd")
    sub.add_parser("list").set_defaults(func=cmd_list)
    p = sub.add_parser("show")
    p.add_argument("id", type=int)
    p.set_defaults(func=cmd_show)
    args = ap.parse_args()
    if args.cmd is None:
        cmd_record(args)
    else:
        args.func(args)


if __name__ == "__main__":
    main()