Add annotation support

app.py

@@ -1,16 +1,53 @@
 from typing import Optional, Tuple
 import gradio as gr
 
+from src.util.board_vis import colored_unicode_board
+from src.util.pgn_util import export_pgn, read_pgn
 from src.thinksqure_engine import ThinkSquareEngine
 
 
+def annotate_pgn(file, analysis_time_per_move: float = 0.1):
+    """Annotate a PGN file with engine analysis.
+
+    Args:
+        file: The PGN file to be annotated.
+        analysis_time_per_move: Time in seconds for engine analysis per move.
+
+    Returns:
+        The annotated PGN file.
+    """
+    game = read_pgn(file.name)
+
+    try:
+        analysis_time_per_move = float(analysis_time_per_move)
+    except ValueError:
+        raise ValueError("Analysis time must be a number.")
+
+    annotated_game = ThinkSquareEngine.annotate(
+        game, analysis_time=analysis_time_per_move
+    )
+
+    output_path = export_pgn(annotated_game)
+
+    return output_path
+
+
 def play_chess(
-    move: Optional[str] = "", fen: Optional[str] = "", draw_ascii: bool = True
+    move: Optional[str] = "",
+    fen: Optional[str] = "",
+    draw_board: bool = True,
+    render_mode: str = "ascii",
 ) -> Tuple:
     """Play a move in a chess game.
     Prerequisites:
-        - User must be asked if they want to play as white or black. If user chooses black, pass an empty string in the first move for the engine to play as white.
+        - User must be asked if they want to play as white or black.
+            - If user chooses black, pass an empty string in the first move for the engine to play as white.
+            - If the user chooses white, ask the user for a move and pass it in the first move for the engine to play as black.
+
         - User must be asked if they want a board drawn in ASCII format.
+            - If they do, pass `draw_board=True` to this function and print the board in ASCII format.
+            - If they do not, pass `draw_board=False`.
+
         - If a move is provided, it must be in long algebraic notation (e.g., "e4", "Nf3", "Bb5").
 
     To start a new game:
@@ -25,10 +62,11 @@ def play_chess(
     Args:
         move: The move to play in long algebraic notation. If None, the engine will play a move.
         fen: The FEN string representing the board state prior to the user's last move. If None, the game starts from the initial position.
-        draw_ascii: Whether to draw the board in ASCII format. Defaults to True.
+        draw_board: Whether to draw the board in ASCII/Unicode format. Defaults to True.
+        render_mode: The rendering mode for the board. Defaults to "ascii". This can be "ascii", "svg", or "unicode".
 
     Returns:
-        The best move played by the engine, the updated board state in FEN notation, and  a string representation of the board if draw_ascii is True else None.
+        The best move played by the engine, the updated board state in FEN notation, and  a string representation of the board if draw_board is True else None.
     """
 
     if move is None or move == "" or move.lower() == "none" or move.lower() == "null":
@@ -48,32 +86,63 @@ def play_chess(
 
     bestmove_san = ThinkSquareEngine.get_best_move(fen)
     fen_after_move = ThinkSquareEngine.get_fen_after_move(bestmove_san, fen)
-
-    board_str = (
-        ThinkSquareEngine.get_ascii_board(fen_after_move) if draw_ascii else None
-    )
+    if draw_board:
+        if render_mode == "ascii":
+            board_str = ThinkSquareEngine.render_board_ascii(fen_after_move)
+        elif render_mode == "svg":
+            board_str = ThinkSquareEngine.render_board_svg(fen_after_move)
+        elif render_mode == "unicode":
+            board_str = colored_unicode_board(fen_after_move)
+        else:
+            raise ValueError(
+                "Invalid render mode. Choose 'ascii', 'svg', or 'unicode'."
+            )
+
+    else:
+        board_str = None
 
     return bestmove_san, fen_after_move, board_str
 
 
-demo = gr.Interface(
-    fn=play_chess,
-    inputs=[
-        gr.Textbox(
-            label="Your Move (in standard algebraic notation - SAN) (Optional)",
-            placeholder="e.g. e4, Nf3, Bb5, etc. Optional - ",
-            value=None,
-        ),
-        gr.Textbox(
+with gr.Blocks(title="ThinkSquare") as app:
+
+    with gr.Tab("Play Chess"):
+        gr.Markdown("### Play Chess with an engine")
+        move_input = gr.Textbox(
+            label="Your Move (SAN)", placeholder="e4, Nf3...", value=None
+        )
+        fen_input = gr.Textbox(
             label="FEN String (optional)",
-            placeholder="Optional - Leave blank to start from initial position",
+            placeholder="Leave blank to start from initial position",
             value=None,
-        ),
-        gr.Checkbox(label="Draw ASCII Board", value=True),
-    ],
-    outputs=("text", "text", "text"),
-    title="Play Chess with an engine",
-    description="This interface allows you to play chess against the Stockfish engine. You can provide a move in standard algebraic notation and an optional FEN string representing the board state. The engine will respond with its best move, the updated FEN string, and an ASCII representation of the board if requested.",
-)
-
-demo.launch(mcp_server=True)
+        )
+        draw_board_checkbox = gr.Checkbox(label="Draw Board", value=True)
+        play_btn = gr.Button("Submit Move")
+
+        best_move_output = gr.Textbox(label="Best Move by Engine (SAN)")
+        updated_fen_output = gr.Textbox(label="Updated FEN")
+        board_output = gr.Textbox(label="Board (ASCII)", visible=True)
+
+        play_btn.click(
+            fn=play_chess,
+            inputs=[move_input, fen_input, draw_board_checkbox],
+            outputs=[best_move_output, updated_fen_output, board_output],
+        )
+
+    with gr.Tab("PGN Annotation"):
+        gr.Markdown("### PGN Analyzer")
+        pgn_file = gr.File(label="Upload PGN", file_types=[".pgn"])
+        analysis_time = gr.Textbox(
+            label="Analysis Time per Move (seconds)", value="0.1"
+        )
+        analyze_btn = gr.Button("Annotate PGN")
+
+        annotated_pgn_file = gr.File(label="Download Annotated PGN")
+
+        analyze_btn.click(
+            fn=annotate_pgn,
+            inputs=[pgn_file, analysis_time],
+            outputs=annotated_pgn_file,
+        )
+
+app.launch(mcp_server=True)

src/thinksqure_engine.py

@@ -2,6 +2,10 @@ from pathlib import Path
 from typing import Optional
 import chess
 import chess.engine
+import chess.svg
+import chess.pgn
+
+from src.util.pgn_util import add_variation
 
 
 class ThinkSquareEngine:
@@ -21,6 +25,108 @@ class ThinkSquareEngine:
 
         return bestmove_san
 
+    @staticmethod
+    def get_engine_analysis(board, analysis_time=0.1):
+        with chess.engine.SimpleEngine.popen_uci(ThinkSquareEngine._ENGINE) as engine:
+            pre_info = engine.analyse(board, chess.engine.Limit(time=analysis_time))
+        return pre_info
+
+    @staticmethod
+    def _perform_post_analysis_and_add_comment(
+        analysis_time,
+        board,
+        played_node,
+        pre_eval,
+        engine_best_move_san,
+        pv,
+    ):
+        post_info = ThinkSquareEngine.get_engine_analysis(board, analysis_time)
+        post_eval = post_info["score"].white().score(mate_score=100000)
+
+        # Evaluation drop
+        eval_drop = (
+            (pre_eval - post_eval)
+            if pre_eval is not None and post_eval is not None
+            else 0
+        )
+
+        # Classification
+        if eval_drop > 200:
+            label = "Blunder"
+        elif eval_drop > 100:
+            label = "Mistake"
+        elif eval_drop > 50:
+            label = "Inaccuracy"
+        elif eval_drop < -150:
+            label = "Super Brilliant"
+        elif eval_drop < -60:
+            label = "Brilliant"
+        elif abs(eval_drop) <= 30:
+            label = None
+        else:
+            label = "Good"
+
+        if label is not None:
+            comment = f"{label}. "
+            if eval_drop > 0 and engine_best_move_san is not None:
+                comment += f"Better was {engine_best_move_san} "
+                played_node.comment = comment
+                if pv is not None:
+                    add_variation(played_node.parent, pv)
+            else:
+                played_node.comment = comment
+
+    @staticmethod
+    def annotate(game, analysis_time: float = 0.1):
+
+        if not isinstance(game, chess.pgn.Game):
+            raise ValueError("Input must be a chess.pgn.Game object")
+
+        if not game.variations:
+            raise ValueError("Game must have at least one variation")
+
+        if analysis_time <= 0:
+            raise ValueError("Analysis time must be greater than 0")
+
+        node = game
+
+        while node.variations:
+            board = node.board()
+            played_node = node.variation(0)
+            played_move = played_node.move
+
+            # Get engine's best move BEFORE the actual move
+            pre_info = ThinkSquareEngine.get_engine_analysis(board, analysis_time)
+
+            pre_eval = pre_info["score"].white().score(mate_score=100000)
+
+            # Best move suggestion
+            engine_best_move = pre_info.get("pv", [None])[0]
+            engine_best_move_san = (
+                board.san(engine_best_move) if engine_best_move else None
+            )
+
+            # Get principal variation (PV)
+            pv = pre_info.get("pv", [])
+
+            # Make the played move and get new evaluation
+            played_move_san = board.san(played_move) if played_move else None
+            board.push(played_move)
+
+            if played_move_san != engine_best_move_san:
+                ThinkSquareEngine._perform_post_analysis_and_add_comment(
+                    analysis_time,
+                    board,
+                    played_node,
+                    pre_eval,
+                    engine_best_move_san,
+                    pv,
+                )
+
+            node = played_node
+
+        return game
+
     @staticmethod
     def is_valid_move(
         move_san: str,
@@ -58,7 +164,7 @@ class ThinkSquareEngine:
             return None
 
     @staticmethod
-    def get_ascii_board(fen: Optional[str] = None) -> str:
+    def render_board_ascii(fen: Optional[str] = None) -> str:
         if fen is None:
             fen = chess.STARTING_FEN
 
@@ -66,3 +172,22 @@ class ThinkSquareEngine:
         ascii_representation = str(board)
 
         return ascii_representation
+
+    @staticmethod
+    def render_board_svg(fen: Optional[str] = None):
+        if fen is None:
+            fen = chess.STARTING_FEN
+
+        board = chess.Board(fen)
+        svg = chess.svg.board(board=board)
+        return svg
+
+    @staticmethod
+    def render_board_unicode(fen: Optional[str] = None) -> str:
+        if fen is None:
+            fen = chess.STARTING_FEN
+
+        board = chess.Board(fen)
+        unicode_representation = board.unicode(invert_color=False)
+
+        return unicode_representation