02_ccc_generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from os.path import *\n",
    "import importlib\n",
    "import random\n",
    "\n",
    "import yaml\n",
    "import jsonlines\n",
    "import pickle\n",
    "import wandb\n",
    "import logging\n",
    "import tqdm\n",
    "import chess\n",
    "import chess.svg\n",
    "from cairosvg import svg2png\n",
    "\n",
    "from sklearn import svm\n",
    "import numpy as np\n",
    "import tensorflow as tf \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "sys.path.append(f\"{os.getcwd()}/lczeroTraining/tf/\")\n",
    "\n",
    "from stockfish import Stockfish\n",
    "from tfprocess import TFProcess\n",
    "from lcztools import LeelaBoard as leelaBoard\n",
    "\n",
    "probing_svm = __import__(\"01_probing_svm\", fromlist=\"*\")\n",
    "\n",
    "%matplotlib inline    \n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "mode = \"train\"          # [\"train\", \"svm_train\", \"svm_eval\"]\n",
    "wandb_project = \"lc0_concept\"\n",
    "name = \"negrand_large\"\n",
    "target_lrs = [1e-3, 1e-4, 1e-5, 1e-6]\n",
    "# target_lrs = [1e-3, 1e-4, 1e-5, 1e-6]\n",
    "num_epochs = 100\n",
    "\n",
    "lichess_eval_path = \"data/lichess_db_eval.jsonl\"\n",
    "sts_path = \"data/STS/STS1-STS15_LAN_v3.epd\"\n",
    "lichess_puzzle_path = 'data/lichess_db_puzzle.csv'\n",
    "\n",
    "data_size = 200000\n",
    "data_ratio = 0.05\n",
    "test_split = 0.1\n",
    "\n",
    "concept_extraction_model = \"linear_svm\"     # [\"linear_svm\", \"svm\", \"mlp\"]\n",
    "concept_extraction_version = \"v4.6\"     # \n",
    "\n",
    "cbm_data_size = 10000\n",
    "batch_size = 4\n",
    "# cbm_label = \"stockfish\"       # [\"lc0\", \"stockfish\"]\n",
    "cbm_label = \"lc0\"       # [\"lc0\", \"stockfish\"]\n",
    "\n",
    "target_layers = [39]\n",
    "\n",
    "stockfish_concepts = [\n",
    "    \"Material_t_mid\", \"Imbalance_t_mid\", \"Pawns_t_mid\", \n",
    "    \"Knights_w_mid\", \"Knights_b_mid\",\n",
    "    \"Bishop_w_mid\", \"Bishop_b_mid\",\n",
    "    \"Rooks_w_mid\", \"Rooks_b_mid\",\n",
    "    \"Queens_w_mid\", \"Queens_b_mid\",\n",
    "    \"Mobility_w_mid\", \"Mobility_b_mid\",\n",
    "    \"Kingsafety_w_mid\", \"Kingsafety_b_mid\",\n",
    "    \"Threats_w_mid\", \"Threats_b_mid\",\n",
    "    \"Space_w_mid\", \"Space_b_mid\",\n",
    "    \"Passedpawns_w_mid\", \"Passedpawns_b_mid\",\n",
    "]\n",
    "\n",
    "sts_concepts = [\n",
    "    \"Undermining\",\n",
    "    \"Open_Files_and_Diagonals\",\n",
    "    \"Knight_Outposts\",\n",
    "    \"Square_Vacancy\",\n",
    "    \"Bishop_vs_Knight\",\n",
    "    \"Re-Capturing\",\n",
    "    \"Offer_of_Simplification\",\n",
    "    \"Advancement_of_f/g/h_Pawns\",\n",
    "    \"Advancement_of_a/b/c_Pawns\",\n",
    "    \"Simplification\",\n",
    "    \"Activity_of_the_King\",\n",
    "    \"Center_Control\",\n",
    "    \"Pawn_Play_in_the_Center\",\n",
    "    \"Queens_and_Rooks_to_the_7th_rank\",\n",
    "    \"Avoid_Pointless_Exchange\",\n",
    "]\n",
    "\n",
    "puzzle_concepts = [\n",
    "    \"fork\",\n",
    "    \"pin\",\n",
    "    \"mate\",\n",
    "    # \"defensiveMove\",\n",
    "    \"hangingPiece\",\n",
    "    \"sacrifice\",\n",
    "    \"attraction\",\n",
    "    \"deflection\",\n",
    "    \"skewer\",\n",
    "    \"discoveredAttack\",\n",
    "    \"capturingDefender\",\n",
    "    \"exposedKing\",\n",
    "    # \"zugzwang\",\n",
    "]\n",
    "\n",
    "target_concepts = stockfish_concepts\n",
    "\n",
    "skip_concept_list = [\"Re-Capturing\", \"Square_Vacancy\", \"Queens_and_Rooks_to_the_7th_rank\", \"Avoid_Pointless_Exchange\"]\n",
    "\n",
    "cfg = {\n",
    "    \"name\": name,\n",
    "    \"target_concepts\": target_concepts,\n",
    "    \"target_lrs\": target_lrs,\n",
    "    \"target_layers\": target_layers,\n",
    "    \"num_epochs\": num_epochs,\n",
    "}\n",
    "\n",
    "result_dir = f\"results/{name}\"\n",
    "stockfish_engine = Stockfish(path=probing_svm.stockfish_8_path)\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "saved_filename = f\"{result_dir}/cbm_0.0001_l39.pkl\"\n",
    "target_layers = [39]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "clfs = {target_layer: {} for target_layer in target_layers}\n",
    "\n",
    "tfproc = probing_svm.load_tf_net(probing_svm.tf_cfg_path, probing_svm.tf_ckp_path)\n",
    "fen_data = probing_svm.get_fen_data(data_size=data_size)\n",
    "for target_layer in target_layers:\n",
    "    for target_concept in target_concepts:\n",
    "        key = f\"size_{data_size}_{data_ratio}_concept_{target_concept}_layer_{target_layer}\".replace(\"/\", \"_\")\n",
    "        cache_path = f\"cache/{concept_extraction_model}_{concept_extraction_version}_{key}.pkl\"\n",
    "        if os.path.exists(cache_path):\n",
    "            with open(cache_path, \"rb\") as f:\n",
    "                clf = pickle.load(f)\n",
    "            print(f\"loaded from {cache_path}\")\n",
    "        else:\n",
    "            train_set_fen, train_set_tag, test_set_fen, test_set_tag = get_data_for_concept(fen_data, target_concept)\n",
    "            clf, train_acc, train_pr, train_recall, test_acc, test_pr, test_recall, = concept_linear_probing(tfproc, target_layer, train_set_fen, train_set_tag, test_set_fen, test_set_tag)\n",
    "            print(f\"{key} : {(train_acc, train_pr, train_recall, test_acc, test_pr, test_recall)}\")\n",
    "            with open(cache_path, \"wb\") as f:\n",
    "                pickle.dump(clf, f)\n",
    "        clfs[target_layer][key] = clf\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def plot_board(fen, move_uci, is_san=False):\n",
    "    if is_san:\n",
    "        move_uci = san_to_uci(fen, move_uci)\n",
    "\n",
    "    \n",
    "    board = chess.Board(fen)\n",
    "    svg =  chess.svg.board(\n",
    "        board,\n",
    "        arrows=[chess.svg.Arrow(\n",
    "            chess.square(ord(move_uci[0]) - ord('a'), int(move_uci[1])-1), \n",
    "            chess.square(ord(move_uci[2]) - ord('a'), int(move_uci[3])-1), \n",
    "            color=\"#0000cccc\",\n",
    "        )],\n",
    "        size=500,\n",
    "    )  \n",
    "    display(\"image/svg+xml\", svg)\n",
    "\n",
    "    return svg\n",
    "\n",
    "def san_to_uci(fen, san):\n",
    "    board = chess.Board(fen)\n",
    "    uci_text = board.copy().push_san(san).uci()\n",
    "    return uci_text\n",
    "\n",
    "\n",
    "def lc0_pred_wdl(tfproc, fen):\n",
    "    input_tf = probing_svm.fen_to_tf_input(fen)\n",
    "    policy, wdl, move_left = tfproc.model.predict(input_tf)\n",
    "    return wdl\n",
    "\n",
    "softmax = lambda x: np.exp(x - np.max(x)) / sum(np.exp(x - np.max(x)))\n",
    "\n",
    "def eval_importance_basic(target_fen, target_move, target_layer=target_layers[0], return_img=False):\n",
    "    stockfish_engine.set_fen_position(target_fen)\n",
    "    best_move = san_to_uci(target_fen, target_move)\n",
    "\n",
    "    moves = [best_move]\n",
    "    replies = []\n",
    "    future_reply_fen = []\n",
    "    future_fen = []\n",
    "    for move in moves:\n",
    "        stockfish_engine.set_fen_position(target_fen)\n",
    "        stockfish_engine.make_moves_from_current_position([move])\n",
    "        next_fen = stockfish_engine.get_fen_position()\n",
    "        future_reply_fen.append(next_fen)\n",
    "\n",
    "        best_reply = stockfish_engine.get_best_move()\n",
    "        stockfish_engine.make_moves_from_current_position([best_reply])\n",
    "        replies.append(best_reply)\n",
    "        next_fen = stockfish_engine.get_fen_position()\n",
    "        future_fen.append(next_fen)\n",
    "    \n",
    "\n",
    "    lc0_activations = probing_svm.activation_extraction(tfproc, target_layer, [target_fen] + future_reply_fen + future_fen)\n",
    "    lc0_wdl = lc0_pred_wdl(tfproc, target_fen)\n",
    "    lc0_wdl = softmax(lc0_wdl[0])\n",
    "    concept_vectors = []\n",
    "    for idx, (key, clf) in tqdm.tqdm(enumerate(clfs[target_layer].items()), desc=\"concepts\"):\n",
    "        if idx >= len(stockfish_concepts):\n",
    "            break\n",
    "        train_concept = clf.decision_function(lc0_activations)\n",
    "        concept_vectors.append(train_concept)\n",
    "    all_concepts_np = np.stack(concept_vectors).T\n",
    "    all_concepts_np = np.clip(all_concepts_np, -10, 10)\n",
    "    train_concepts_np = all_concepts_np[0]\n",
    "\n",
    "    print(f\"input: {target_fen}, move: {best_move}, reply: {replies[0]}\")\n",
    "    stockfish_engine.set_fen_position(target_fen)\n",
    "    # board_svg = plot_board(target_fen, best_move)\n",
    "    concept_scores = [(cname, c.item()) for c, cname in zip(train_concepts_np, target_concepts[:len(stockfish_concepts)])]\n",
    "    print(f\"concept score: {concept_scores}\")\n",
    "    \n",
    "    concept_importance = all_concepts_np[2] - all_concepts_np[0]\n",
    "    \n",
    "    importance = sorted([(p.item(), name) for name, p in zip(target_concepts[:len(stockfish_concepts)], concept_importance)], reverse=True)\n",
    "    print(f\"importance: {importance}\")\n",
    "\n",
    "    if return_img:\n",
    "        return concept_scores, importance, board_svg\n",
    "    return concept_scores, importance\n",
    "\n",
    "\n",
    "target_position = \"5rk1/pp4p1/7p/3pQ3/1n1P1P2/P7/1P6/2K5 b - - 0 33\"\n",
    "target_move = \"Nd3+\" \n",
    "concept_score, concept_importance = eval_importance_basic(target_position, target_move)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "engine_cache = {}\n",
    "\n",
    "def get_engine_evalutaion(fen, target_move):\n",
    "    \n",
    "    key = f\"{fen} {target_move}\"\n",
    "    if key in engine_cache.keys():\n",
    "        return engine_cache[key]\n",
    "\n",
    "    depth = 20\n",
    "\n",
    "    stockfish_engine.set_fen_position(fen)\n",
    "    best_moves = stockfish_engine.get_top_moves(2)\n",
    "    stockfish_engine.set_depth(depth + 1)\n",
    "    eval1 = stockfish_engine.get_evaluation()\n",
    "\n",
    "    stockfish_engine.set_depth(depth + 1)\n",
    "    target_move_san = target_move.split(\" \")[0]\n",
    "    board = chess.Board(fen)\n",
    "    uci_text = board.copy().push_san(target_move_san).uci()\n",
    "    stockfish_engine.make_moves_from_current_position([uci_text])\n",
    "    eval2 = stockfish_engine.get_evaluation()\n",
    "    move_score = f\"Mate in {eval2['value'] * -1}\" if eval2['type'] == \"mate\" else f\"{eval2['value'] * -1}cp\"\n",
    "    if move_score == \"Mate in 0\": \n",
    "        move_score = \"Checkmate!\"\n",
    "    if \"Mate\" not in move_score:\n",
    "        best_reply = stockfish_engine.get_best_move()\n",
    "        move_score = f\"{move_score}, expected reply - {best_reply}\"\n",
    "\n",
    "    best_move1 = board.copy().san(board.copy().push_uci(best_moves[0]['Move']))\n",
    "    if best_moves[0]['Mate'] is not None:\n",
    "        best_score1 = f\"Mate in {best_moves[0]['Mate'] - 1}\"  \n",
    "        if best_score1 == \"Mate in 0\": \n",
    "            best_score1 = \"Checkmate!\"\n",
    "    else:\n",
    "        if \"cp\" in move_score:\n",
    "            diff = best_moves[0]['Centipawn'] - eval2['value'] * -1\n",
    "            if diff < 20:\n",
    "                best_score1 = f\"similar to actual move\"\n",
    "            elif diff < 50:\n",
    "                best_score1 = f\"better than actual move over 20cp\"\n",
    "            elif diff < 100:\n",
    "                best_score1 = f\"better than actual move over 50cp\"\n",
    "            else:\n",
    "                best_score1 = f\"better than actual move over 100cp, worth a pawn\"\n",
    "        else:\n",
    "            best_score1 = f\"{best_moves[0]['Centipawn']}cp\"\n",
    "\n",
    "    if len(best_moves) == 1:\n",
    "        eval_str = f\"evaluation: only legal move; move - {target_move_san} {move_score}\"\n",
    "        return eval_str\n",
    "\n",
    "    best_move2 = board.copy().san(board.copy().push_uci(best_moves[1]['Move']))\n",
    "    if best_moves[1]['Mate'] is not None:\n",
    "        best_score2 = f\"Mate in {best_moves[1]['Mate'] - 1}\" \n",
    "        if best_score2 == \"Mate in 0\": \n",
    "            best_score2 = \"Checkmate!\"\n",
    "    else:\n",
    "        if \"cp\" in move_score:\n",
    "            diff = best_moves[1]['Centipawn'] - eval2['value'] * -1\n",
    "            if diff < 20:\n",
    "                best_score2 = f\"similar to actual move\"\n",
    "            elif diff < 50:\n",
    "                best_score2 = f\"better than actual move over 20cp\"\n",
    "            elif diff < 100:\n",
    "                best_score2 = f\"better than actual move over 50cp\"\n",
    "            else:\n",
    "                best_score2 = f\"better than actual move over 100cp, worth a pawn\"\n",
    "        else:\n",
    "            best_score2 = f\"{best_moves[1]['Centipawn']}cp\"\n",
    "\n",
    "    eval_str = f\"evaluation: actual move - {target_move_san} {move_score}, best move - {best_move1} {best_score1}, second best move - {best_move2} {best_score2}\"\n",
    "\n",
    "    engine_cache[key] = eval_str\n",
    "    return eval_str\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_table = {\n",
    "    \"Material_t_mid\": \"Material\", \n",
    "    \"Imbalance_t_mid\": \"Imbalance\", \n",
    "    \"Pawns_t_mid\": \"Pawns\", \n",
    "    \"Knights_w_mid\": \"White Knights\",\n",
    "    \"Knights_b_mid\": \"Black Knights\",\n",
    "    \"Bishop_w_mid\": \"White Bishop\",\n",
    "    \"Bishop_b_mid\": \"Black Bishop\",\n",
    "    \"Rooks_w_mid\": \"White Rooks\",\n",
    "    \"Rooks_b_mid\": \"Black Rooks\",\n",
    "    \"Queens_w_mid\": \"White Queens\",\n",
    "    \"Queens_b_mid\": \"Black Queens\",\n",
    "    \"Mobility_w_mid\": \"White Mobility\",\n",
    "    \"Mobility_b_mid\": \"Black Mobility\",\n",
    "    \"Kingsafety_w_mid\": \"White Kingsafety\",\n",
    "    \"Kingsafety_b_mid\": \"Black Kingsafety\",\n",
    "    \"Threats_w_mid\": \"White Threats\",\n",
    "    \"Threats_b_mid\": \"Black Threats\",\n",
    "    \"Space_w_mid\": \"White Space\",\n",
    "    \"Space_b_mid\": \"Black Space\",\n",
    "    \"Passedpawns_w_mid\": \"White Passedpawns\",\n",
    "    \"Passedpawns_b_mid\": \"Black Passedpawns\",\n",
    "    }\n",
    "\n",
    "def convert_concept(text):\n",
    "    return convert_table.get(text, text)\n",
    "\n",
    "target_position = \"8/5ppk/1p6/8/1K1P4/4R1P1/2q4P/8 w - - 0 0\"\n",
    "target_move = \"d5\"        # Nd3+  =  b4d3  \n",
    "\n",
    "concept_score, concept_importance = eval_importance_basic(target_position, target_move)\n",
    "print(concept_score)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from 61 samples\n",
    "\n",
    "target_fens = ['8/1bp4p/p3P1p1/3n4/3kN1P1/5P1P/P1P2K2/8 b - - 0 0', 'r1b1k2r/ppppb1qp/2n2n2/6p1/2PP1p2/P1NB1N1P/1P2R1P1/R1BQ2K1 b - - 0 0', 'r1b1k1q1/ppppb2r/2P2n2/5N1p/2P3p1/P1NB3P/1P2Q1P1/R1B3K1 w - - 0 0', '1r3rk1/3bqNbp/pp1p2p1/2pB4/P2PP1n1/2P1B3/1P1Q2PP/R4RK1 w - - 0 0', '1r2qb2/3b1Nkp/pp1p2pn/2pB4/P2PP3/2P5/1P1Q2PP/R5K1 w - - 0 0', 'r1bqkbnr/ppp2ppp/8/3pn3/8/4P3/PPPN1PPP/R1BQKBNR w - - 0 0', 'r1bqk2r/ppp2ppp/8/3nb3/8/4PN2/PP3PPP/R1BQKB1R b - - 0 0', 'rn1q1rk1/1b2bppp/p2ppn2/1p6/3NPP2/1BN1B3/PPP3PP/R2Q1RK1 w - - 0 0', 'r2r2k1/1b1nbppp/pqn1p3/1p2P3/5B2/1BN2N2/PPP3PP/R2Q1R1K w - - 0 0', 'r2r2k1/1b2bppp/pqn1p3/1pn1P3/8/1BN1BN2/PPP1Q1PP/R4R1K b - - 0 0', 'r4rk1/1bq1bppp/p1n1p3/1p2P3/8/1PN1BN2/1PP2QPP/3R1R1K b - - 0 0', 'rnbqkbnr/pppppppp/8/8/P7/8/1PPPPPPP/RNBQKBNR b - - 0 0', 'r1bqkbnr/pppp1ppp/2n5/4p3/P1P5/8/1P1PPPPP/RNBQKBNR w - - 0 0', 'r1bqk2r/pppp1ppp/2n2n2/4p3/P1P1P3/3P4/1P1N1PPP/R2QKBNR b - - 0 0', 'r1bq1rk1/ppp2ppp/2np3P/4p3/P1P1P1n1/3P4/1P1N1PP1/R2QKBNR b - - 0 0', 'r2q1rk1/ppp2ppp/3p3n/4pB2/P1PnP3/3P4/1P1N1PP1/R2QK1NR b - - 0 0', 'r2q1rk1/ppp2pp1/5P1p/3pp2Q/P1P1N3/3P4/1P3PP1/n2K2NR b - - 0 0', 'rnbqkb1r/ppp1n2p/4p1p1/3p1p1Q/3P4/4P3/PPPN1PPP/RNB1KB1R w - - 0 0', '1k6/p7/2N5/2R5/6bP/2KB4/r7/8 b - - 0 0', 'k7/8/2N5/p4R2/7P/2K5/r7/8 w - - 0 0', '8/8/8/8/5K2/3Q4/8/6k1 w - - 0 0', 'r1b2rk1/3p1ppp/1pnb4/pP4q1/8/2P1PBP1/3P3P/RNBQK2R b - - 0 0', '2Q5/5ppk/1p3b2/p7/3Pq3/1KP1n1P1/7P/1R2R3 b - - 0 0', '8/5ppk/1p6/8/1K1P4/4R1P1/2q4P/8 w - - 0 0', 'r1bqk2r/ppp2ppp/2n1pn2/3p4/3P4/5NP1/PPPQPPP1/RN2KB1R w - - 0 0', 'rnbqkbnr/pppp1ppp/4p3/8/4P3/8/PPPP1PPP/RNBQKBNR w - - 0 0', 'r5k1/pbpp1ppp/1p6/3P4/2B3n1/2N2N2/PPP2PPP/4R1K1 w - - 0 0', 'r2qkb1r/1p1n1pp1/p1p1p2p/3pPn2/1P1P4/P1NQ1N2/2P2PPP/R1B2RK1 w - - 0 0', 'b2q1rk1/3pb1pp/n1n1p3/p3Np2/2PpPB2/P2Q1P2/1P4PP/4RRK1 b - - 0 0', 'b2q1rk1/3pb1pp/2n1p3/p1n1Np2/2PpPB2/P2Q1P2/1P4PP/4RRK1 w - - 0 0', 'rnb1kb1r/pppqnp1p/4p1p1/3pP3/3P1NQ1/2N1B3/PPP2PPP/R3KB1R b - - 0 0', 'r1bqk2r/pp1pbppp/2n2n2/2p1p1N1/2B1P3/5Q2/PPPP1PPP/RNB1K2R w - - 0 0', 'r1b1k2r/pp1pbNpp/2n2n2/q1p1p3/2B1P3/5Q2/PPPP1PPP/RNB1K2R w - - 0 0', 'r3k2N/pp2b1pp/8/q1pPpb2/8/8/PPnP1PPP/RNBQ1RK1 w - - 0 0', 'r3k2N/pp2b1pp/8/q1pPpb1Q/8/8/PPnP1PPP/RNB2RK1 b - - 0 0', 'r2qkb1r/p2nnpp1/4p2p/1p1pPb2/2pP4/1PP1BN2/P2NBPPP/R2Q1RK1 w - - 0 0', 'rnb1qrk1/ppp1p1bp/3p1np1/5p2/2PP4/2N2NP1/PPQ1PPBP/R1B2RK1 b - - 0 0', 'r1b2rk1/ppn1pqbp/2p2np1/2Pp1p2/3P3N/PPN3P1/1BQ1PPBP/R4RK1 b - - 0 0', 'r1b1k2r/ppp2ppp/4pn2/4P3/2B5/4P3/PP1N1PPP/R3K2R b - - 0 0', 'r4rk1/pbpn1ppp/1p2p3/4P3/2B2P2/4P3/PP1N2PP/2R2RK1 w - - 0 0', 'r4rk1/pbp2ppp/1p2p3/1Bn1P3/5PP1/4P3/PP1N3P/2R2RK1 b - - 0 0', 'r4rk1/p1p3pp/bp2pp2/1Bn1P3/1P3PP1/4P3/P2N3P/2R2RK1 w - - 0 0', 'r5k1/2R4p/pp3p2/1b3P2/1P6/4P3/P2N3P/6K1 w - - 0 0', '1r4k1/7p/p1RN1p2/1p3P2/PP6/3bP3/7P/6K1 w - - 0 0', '6k1/7p/3N1p2/P7/8/3bP3/5K1P/8 w - - 0 0', '6k1/7p/3N1p2/P7/8/3bPK2/7P/8 b - - 0 0', '2N5/P7/2b2p2/3k3P/8/4K3/8/8 b - - 0 0', '1kn3r1/7q/2p1p3/1p1pPp2/pP1P3P/P1P2N1Q/8/2K4R b - - 0 0', '6rr/pbpknRpN/1p2p3/3pP3/2nP4/P1PB2P1/2P4P/R1B3K1 w - - 0 0', 'rnbqkb1r/pppppppp/5n2/8/3P4/8/PPP1PPPP/RNBQKBNR w - - 0 0', 'rnbqkb1r/pppppppp/5n2/8/3P4/4P3/PPP2PPP/RNBQKBNR b - - 0 0', 'r2qr1k1/pb1n1ppp/2nbp3/1p1pN3/2pPPP2/2P5/PPBN2PP/R1B1QRK1 w - - 0 0', 'r2qr1k1/p2nb1pp/2b1p3/1p1pP3/2pP4/2P2NQ1/PPB3PP/R1B2RK1 b - - 0 0', 'rn2q2N/pp4pp/2kbp3/2p5/4pPB1/8/PPPP2PP/RNB1K2R w - - 0 0', 'r3q2N/pp1n2pp/2kb4/2p2B2/4p3/8/PPPP2PP/RNB1K2R w - - 0 0', 'r3q2N/pp1n2pp/2kb4/2p2B2/4p3/8/PPPP2PP/RNB1KR2 b - - 0 0', '4r3/ppb1q1p1/2k2n2/2p3B1/8/3B4/PPP1N1P1/R3KR2 b - - 0 0', '8/pB3pkp/P7/8/6nP/4P1P1/3r1PK1/2R5 w - - 0 0', '6rr/pp1kb3/4p3/2npP3/3B2Pp/P4N2/1PP2K2/R5R1 w - - 0 0', 'r1b2rk1/1p2nppp/2n1p3/1B1p4/Pb1B4/4PN2/2q2PPP/1N1R2K1 w - - 0 0']\n",
    "target_moves = ['27... Ne7', '17... h5', '23. Nh6', '22. Nd8+', '27. Qxh6+', '5. Ngf3', '9... Qd6', '11. e5', '16. Qe2', '17... Qc7', '20... Nb4', '1... e5', '3. e4', '6... d6', '9... Nxh6', '12... Nhxf5', '17... dxe4', '6. Qe2', '41... Ka8', '44. Rxa5+', '65. Kf3', '14... Na7', '36... Qc2+', '42. d5', '7. Qg5', '2. Nf3', '16. h3', '11. Ne2', '18... Nc5', '19. Qd1', '8... Nf5', '6. Nxf7', '7. Nxh8', '12. Qh5+', '12... g6', '11. a4', '8... c6', '13... Ng4', '10... Nd7', '14. Bb5', '15... f6', '17. Bc6', '24. Ne4', '28. Rxa6', '35. Kf3', '35... Kf8', '48... Ba8', '31... Nb6', '19. Bg5', '2. e3', '2... c5', '13. Ndf3', '17... Nf8', '12. f5', '14. Rf1', '14... Nf6', '21... Be5', '40. Rf1', '26. b4', '17. Re1']\n",
    "ref_list = ['Stops the pawn and hopes for an exchange ( bishop for a knight ) and maybe a pawn ...\\n', 'Time wise the game was relatively quick , the most time spent thinking was on this move . White can now no longer afford any passive move anymore , Black has a pawn phalanx marching straight toward the king , backed by a queen and a rook . White on the other hand has a rook targeting the pinned bishop , with good central pawns . My initial thought was that Black has the better attack and I should defend , but this being a kings gambit game and because I was a pawn down I decided to throw caution in to the wind .\\n', 'Will trade back my rook I lost for a minor piece earlier . The computer shows that this move leads to an about equal game , Nb5 would have given white a much better advantage and keeps up the pressure .\\n', 'need to remove his rook off the f file and temporarily put his b rook out of action . no idea how it will go\\n', 'queen moves up . if he moves to g8 , the windmill will strike again ( 28 . Ng5 followed by Qxh7 mate\\n', 'white to force trade or back up knight\\n', 'black moves queen up to backup bishop and possible check and attack at whites front row .\\n', \"The pawn evades the attack and , in turn , attacks Black 's N. The main alternative would have been 11 . Qf3 but this ties the Q to the e-pawn 's defence and pins the pawn against the Q , reducing its scope for involvement in attack .\\n\", 'The Q avoids the threatened discovered attack and bolsters the defence of the e-pawn .\\n', 'Black relieves the pin and renews the attack on the e-pawn\\n', 'So instead he attacks the vulnerable c-pawn , while also revealing the power of the Bb7 . It would have been dangerous to take on e5 as White would have been able to pin the N with 21 . Bf4 and pile up pressure with 22 . Rfe1 .\\n', \"I start my favorite opening which is the king 's pawn opening .\\n\", \"Finally , he returned the king 's pawn opening .\\n\", 'I free my light-square bishop .\\n', 'I take the free pawn .\\n', 'I take it with my knight .\\n', 'I take his knight .\\n', 'Moving my queen out of danger\\n', 'Moving out of check while still protecting the pawn\\n', 'I take the pawn , checking the king and deciding to trade my rook for his . Worst comes to worst , he takes my rook and I take his , and then I can concentrate on promoting that pawn .\\n', 'Moving the king closer to provide support to mate .\\n', 'Forgetting the pin I mistakenly move to attack the pawn .\\n', 'The first check forces his King to move to a3 .\\n', 'This is a questionable move , allowing me to fork his King , Rook , and d5 pawn by playing ... Qd2+ .\\n', 'Threatening a poisoned pawn .\\n', '2 . Nf3 - I thought of the usual d4 but decided to develop a piece , namely my king knight .\\n', '16. h3 - this move creates some air ( luft ) for my King , Henry and drives back the black knight on g4 .\\n', 'knowing its threat has been parried ... white seeks other attacking options ... remove the f5 knight ...\\n', 'Forcing the queen to stay back , threatening e4 and give the oportunity to advance my pawn to d4 .\\n', '? ! I expected Qd2 who looked more natural , protecting the bishop . Computer analysis : Mistake : +4.26 ? 19 . Qd1 fxe4 20 . Nxc6 Bxc6 21. fxe4 Qb6 22 . Qd2 Nxe4 23 . Qd3 Qxb2 24 . Rxe4 Bxe4 25 . Qxe4 Bxa3 26 . Bc7 Rxf1+ 27 . Kxf1 Qc1+ 28 . Kf2 Best : +3.07 ? 19 . Qd2 fxe4 20. fxe4 Nxe5 21 . Bxe5 Rxf1+ 22 . Rxf1 Nxe4 23 . Qd3 Qb6 24 . Rf4 Qxb2 25 . Rxe4 Bh4 26. g3 Bxe4 27 . Qxe4\\n', '8 . ... ... Nf5 - good move - puts knight on a more active square and black has the option of moving his bishop back to b4 or to put it on g7 . On balance , despite my previous note I think g7 is the better choice . Black now starts a pawn roll , hoping that white will castle kingside with 9. h5 .\\n', \"I am now attacking flyfish283 's queen and rook and is being protected by my bishop . flyfish283 obviously moves his queen out of the way allowing me to take the rook .\\n\", 'I take the rook .\\n', 'I put the king in check .\\n', \"A good move , blocking the escape of my knight , attacking my queen and defending his bishop at the same time . It was n't unexpected .\\n\", 'trying to exchange the weak a-pawn .\\n', \"copying white 's positional play style , blocking the long diagonal from the white bishop and protects the important d5 square\\n\", \"Black 's counter attack ; preparing e5 . ... Nd7 here would have reduced options and let black defend d4 . 14. h3 is not really a threat as we will see because the knight has an excellent sacrifying option\\n\", 'Knight moves to protect itself from e5 pawn and now threatens e5 .\\n', 'Attacking the knight and the c pawn\\n', \"Black does n't capitalise on the loss of tempo and chooses to attack the advanced pawn . Computer sees this as a mistake .\\n\", \"I decide to threaten his rook rather than close the c file with my bishop . Computer does n't like this , it removes my advantage completely to zero .\\n\", 'Moving to attack the f6 pawn .\\n', 'I decided to take the free pawn . Computer says a5 led to a great advantage for me . 28. a5 Ra8 29. e4 Kg7 30 . Kf2 h5 31 . Ke3 Bc4 32 . Nb7 Kf7 33 . Nc5 Ke7 34 . Rc7+ + ( 4.54++ )\\n', 'I moved the king to allow the pawn to advance with some protection . I was hoping that by overloading the bishop I would be able to move the a-pawn . Unfortunately my knight is wasted here . The computer would have preferred Ne8 forcing the king to protect the pawn Kf7 , then I would do Nc7 which puts the knight in a position to protect the pawn as it moves through a6 and eventually a8 , allowing for at least a sacrifice of bishop for pawn .\\n', 'Black decides to start getting king involved .\\n', 'Blunder : black lets me fork his bishop and king . Computer Mate in 12 : 49. h6 f5 50. h7 f4+ 51 . Kxf4 Kc4 52. h8=Q Kd3 53 . Nb6 Bg2 54 . Qd8+ Ke2 55. a8=Q Bxa8 56 . Nxa8 Kf2 57 . Qd2+ Kg1 58 . Kf3 Kf1 59 . Qf2 #\\n', \"What 's this my Knight is going opposite way ! his 3 pieces and a pawn vs 2 pieces , black has gone insane !\\n\", \"! Of course . The rook ca n't defend both the e7 knight and g7 pawn at once . If 18 ... Rhg8 had been played , Black had been able to play 19 ... Rae8 now .\\n\", 'I play the passive looking e3 , opening the diagonal for the Bishop and preparing f4 .\\n', \"This was n't according to plan . White hoped for e6 or d5 . But the Stonewall plan against cd is ed ! opening the diagonal for the Bishop .\\n\", 'White defends e5 with a piece .\\n', \"A haste retreat , to an unfavorable square . White 's next moves reduce Black 's defenses , and game time . The as yet unmoved Bishop on c1 will be exchanged for the developed Bishop on e7 , and the White Knight assumes a threatening post .\\n\", \"The N could n't yet be saved , but the pawn threatens to complicate things . I 'm happy to swap it off . The Knight can wait .\\n\", '? White could have won a pawn and threatened to release his Knight with 14 . Bxh7 . Furthermore it would have been better to castle than just play 14 . Rf1 , getting the King out of the centre and still getting the Rook to f1 .\\n', 'Now the N defends both h7 and e4 , freeing the Queen and Rook for action .\\n', 'Over-defending the Knight and attacking b2\\n', '? Worse than losing a pawn .\\n', \"I advance a pawn to my opponent 's knight and unless he wants to retreat he will check my king into a better position ... .\\n\", 'Re1 ? ? Hangs the rook ! !\\n']\n",
    "gac_list = ['The knight retreats to the rescue .\\n', '... and I attack his pawn with my bishop ...\\n', 'I attack his queen with my knight .\\n', 'I check him with my knight .\\n', 'I take the knight with my queen .\\n', 'I develop my knight .\\n', '... and I attack his queen with my bishop .\\n', 'I decide to push my pawn to e5 , threatening to fork the knight and the knight .\\n', 'I attack the pawn again .\\n', 'Black moves his queen out of the way to attack the pawn .\\n', \"Black 's knight is immune to defend the pawn and the knight on f3 .\\n\", 'I open with my pawn to e5 .\\n', 'I started with my usual opening .\\n', '... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...\\n', 'I take the pawn .\\n', '... and I capture his bishop with my knight ...\\n', 'I take the knight .\\n', 'I decided to push the queen .\\n', '... and I attack his king ...\\n', 'I take the pawn with my rook .\\n', \"White 's king has to retreat .\\n\", 'Black moves his knight to a good square .\\n', 'I attack his king again .\\n', 'I push up my pawn .\\n', 'I attack his knight\\n', \"I decided to start the King 's Pawn Opening .\\n\", 'I decide to push my knight to the knight .\\n', 'White has no way to protect the pawn with my knight .\\n', \"Black 's queen has no way to retreat .\\n\", 'I attack his pawn with my queen .\\n', 'Black moves his knight to e7 and attacking the pawn on e5 .\\n', 'I take the pawn with my knight .\\n', 'I take the rook .\\n', 'I attack the king again .\\n', 'Black moves his knight to a passive square .\\n', 'I push up my pawn .\\n', '... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...\\n', 'Black moves his knight to g4 and attacking the pawn on g4 .\\n', 'Black moves his knight to a better square .\\n', 'I attack his knight\\n', '... and I attack his pawn with my bishop .\\n', 'I attack his rook with my bishop .\\n', 'White has no way to take the pawn with his knight .\\n', '... and I capture with my rook ...\\n', \"White has no solution to prevent the King 's Gambit .\\n\", 'Black gets his king out of the way and I am going to get his king into the corner .\\n', '... and he attacks my bishop with my bishop ...\\n', \"Black 's knight is going to a better square .\\n\", 'I attack his knight\\n', 'I open up my bishop and queen .\\n', 'I push up my pawn to c5 .\\n', 'The knight retreats to the bishop .\\n', 'The knight comes into the attack .\\n', 'I decide to push my pawn forward to the king .\\n', 'I move my rook to threaten the queen\\n', 'The knight comes out to the corner .\\n', 'I decided to bring my bishop into play and attacking my pawn on e4 .\\n', 'I move my knight to a better square .\\n', 'I push up my pawn to attack the knight .\\n', 'I attack his queen .\\n']\n",
    "\n",
    "target_moves_strip = list(map(lambda x: x.split()[1], target_moves))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import re\n",
    "def get_all_attacks(fen, after_move=None):\n",
    "    board = chess.Board(fen)\n",
    "    if after_move != None:\n",
    "        board.push_san(after_move)\n",
    "    attacks = []\n",
    "    for move in board.pseudo_legal_moves:\n",
    "        if board.piece_at(move.to_square) != None:\n",
    "            attacks.append(board.lan(move))\n",
    "    board.turn = chess.BLACK if board.turn == chess.WHITE else chess.WHITE\n",
    "    for move in board.pseudo_legal_moves:\n",
    "        if board.piece_at(move.to_square) != None:\n",
    "            attacks.append(board.lan(move))\n",
    "    return attacks\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raise RuntimeError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"########## Your API Key ##########\"\n",
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "response = client.chat.completions.create(\n",
    "  model = \"gpt-4o\",\n",
    "  messages = [\n",
    "    {\"role\": \"system\", \"content\": \"This is a chess commentary assistant consists of board, move, evaluation by stockfish, most important concept, and comment based on concept.\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: r4rk1/p3ppbp/Pp1q1np1/3PpbB1/2B5/2N2P2/1PPQ2PP/3RR1K1 b - - 0 18\"\n",
    "        \"actual move: 18... Qc5+\"\n",
    "        \"evaluation: actual move - Qc5+ 361cp, best move - Qc5+ 342cp, second best move - Bxc2 -6cp\"\n",
    "        \"concept: fork\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: great move, fork on king and bishop .\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: 2rqrbk1/pp3ppp/8/3p1N2/3NnnQ1/2P4P/PP3PP1/R4RK1 w - - 0 23\"\n",
    "        \"actual move: 23. Nh6+\"\n",
    "        \"evaluation: actual move - Nh6+ 468cp, best move - Nh6+ 511cp, second best move - Qxf4 -25cp\"\n",
    "        \"concept: pin, fork, queen\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: nice move, plan to fork king and queen Nh6+ and Nf7+, and black cannot capture knight because of pin .\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: r4rk1/1bq1bppp/p1n1p3/1p2P3/8/1PN1BN2/1PP2QPP/3R1R1K b - - 0 0\"\n",
    "        \"actual move: Nb4\"\n",
    "        \"evaluation: actual move - Nb4 41cp, best move - Rae8 51cp, second best move - Nb4 36cp\"\n",
    "        \"concept: Re-Capturing, Queen\"\n",
    "        )},\n",
    "  ]\n",
    ")\n",
    "\n",
    "print(response)\n",
    "print(response.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### vanilla gpt\n",
    "\n",
    "base_responses_4o = []\n",
    "\n",
    "# for idx in range(50):\n",
    "for idx in range(len(target_fens)):\n",
    "    target_position = target_fens[idx]\n",
    "    target_move = target_moves_strip[idx]\n",
    "    target_move_full = target_moves[idx]\n",
    "\n",
    "\n",
    "    prompt = [\n",
    "    {\"role\": \"system\", \"content\": \"This is a chess commentary assistant consists of board, move, and comment .\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: r4rk1/p3ppbp/Pp1q1np1/3PpbB1/2B5/2N2P2/1PPQ2PP/3RR1K1 b - - 0 18\\n\"\n",
    "        \"actual move: 18... Qc5+\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: great move, fork on king and bishop .\\n\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: 2rqrbk1/pp3ppp/8/3p1N2/3NnnQ1/2P4P/PP3PP1/R4RK1 w - - 0 23\\n\"\n",
    "        \"actual move: 23. Nh6+\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: nice move, plan to fork king and queen Nh6+ and Nf7+, and black cannot capture knight because of pin .\\n\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        f\"board: {target_position}\\n\"\n",
    "        f\"actual move: {target_move_full}\\n\"\n",
    "        )},\n",
    "    ]\n",
    "    print(prompt)\n",
    "    response = client.chat.completions.create(model=\"gpt-4o\", messages=prompt, temperature=0.01)\n",
    "    base_responses_4o.append((idx, response.choices[0].message.content))\n",
    "    # print(response.choices[0].message.content)\n",
    "\n",
    "print(base_responses_4o)\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gpt + concept\n",
    "\n",
    "responses_4o = []\n",
    "\n",
    "for idx in range(len(target_fens)):\n",
    "    target_position = target_fens[idx]\n",
    "    target_move = target_moves_strip[idx]\n",
    "    target_move_full = target_moves[idx]\n",
    "    if \"O-O\" in target_move:\n",
    "        tmp = target_position.split(\"- -\")\n",
    "        target_position = tmp[0] + \"KQkq -\" + tmp[1]\n",
    "\n",
    "    concept_score, concept_importance = eval_importance_basic(target_position, target_move)\n",
    "    concept = f\"concept candidates: {convert_concept(concept_importance[0][1])}, {convert_concept(concept_importance[1][1])}, {convert_concept(concept_importance[2][1])}, {convert_concept(concept_importance[3][1])}, {convert_concept(concept_importance[4][1])}\"\n",
    "    print(concept)\n",
    "\n",
    "    engine_eval = get_engine_evalutaion(target_position, target_move)\n",
    "    attacks = get_all_attacks(target_position, target_move)\n",
    "\n",
    "    prompt = [\n",
    "    {\"role\": \"system\", \"content\": (\n",
    "            \"This is a chess commentary assistant consists of board, move, evaluation by stockfish, important concept candidates and comment.\\n\"\n",
    "            \"You are required to generate correct, clear, concise and self-contained comment for the move, using concept candidates as a hint.\\n\"\n",
    "            \"Evaluating the concept candidates may improve understanding of the position.\\n\"\n",
    "        )},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: r4rk1/p3ppbp/Pp1q1np1/3PpbB1/2B5/2N2P2/1PPQ2PP/3RR1K1 b - - 0 18\\n\"\n",
    "        \"actual move: 18... Qc5+\\n\"\n",
    "        \"evaluation: actual move - Qc5+ 361cp, best move - Qc5+ 342cp, second best move - Bxc2 -6cp\\n\"\n",
    "        \"concept candidates: fork, White king, White bishop\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": (\n",
    "        \"concept evaluation: fork - 18... Qc5+ is fork, attacking king and undefended bishop. White king - White king is attacked by Qc5+. White bishop - White bishop is attacked by queen. Only defense is Qd4, but will be captured by e5 pawn. \\n\"\n",
    "        \"comment: great move, fork on king and bishop .\\n\"\n",
    "        )},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: 2rqrbk1/pp3ppp/8/3p1N2/3NnnQ1/2P4P/PP3PP1/R4RK1 w - - 0 23\\n\"\n",
    "        \"actual move: 23. Nh6+\\n\"\n",
    "        \"evaluation: actual move - Nh6+ 468cp, best move - Nh6+ 511cp, second best move - Qxf4 -25cp\\n\"\n",
    "        \"concept candidates: pin, Black queen, White rook\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": (\n",
    "        \"concept evaluation: pin - g7 pawn is pinned and cannot capture the knight. Black queen - after Kh8, Nf7+ forks black king and queen. White rook - It does not engage in attack. It is irrelevent concept. \\n\"\n",
    "        \"comment: nice move, plan to fork king and queen Nh6+ and Nf7+, and black cannot capture knight because of pin .\\n\"\n",
    "        )},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        f\"board: {target_position}\\n\"\n",
    "        f\"actual move: {target_move_full}\\n\"\n",
    "        f\"{engine_eval}\\n\"\n",
    "        # f\"attacks: {attacks}\\n\"\n",
    "        f\"{concept}\\n\"\n",
    "        )},\n",
    "    ]\n",
    "    response = client.chat.completions.create(model=\"gpt-4o\", messages=prompt, temperature=0.01)\n",
    "    print((concept, response.choices[0].message.content))\n",
    "    responses_4o.append((idx, concept, response.choices[0].message.content))\n",
    "\n",
    "print(responses_4o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### gpt + engine\n",
    "\n",
    "engine_responses_4o = []\n",
    "\n",
    "for idx in range(len(target_fens)):\n",
    "    target_position = target_fens[idx]\n",
    "    target_move = target_moves_strip[idx]\n",
    "    target_move_full = target_moves[idx]\n",
    "    if \"O-O\" in target_move:\n",
    "        tmp = target_position.split(\"- -\")\n",
    "        target_position = tmp[0] + \"KQkq -\" + tmp[1]\n",
    "\n",
    "    engine_eval = get_engine_evalutaion(target_position, target_move)\n",
    "\n",
    "    prompt = [\n",
    "    {\"role\": \"system\", \"content\": (\n",
    "        \"This is a chess commentary assistant consists of board, move, evaluation by stockfish, and comment based on engine evaluation.\\n\"\n",
    "        \"Engine evaluation score in centi-pawn is advantage of the player over the oppenent.\\n\"\n",
    "        )},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: r4rk1/p3ppbp/Pp1q1np1/3PpbB1/2B5/2N2P2/1PPQ2PP/3RR1K1 b - - 0 18\\n\"\n",
    "        \"actual move: 18... Qc5+\\n\"\n",
    "        \"evaluation: actual move - Qc5+ 361cp, best move - Qc5+ 342cp, second best move - Bxc2 -6cp\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: great move, fork on king and bishop .\\n\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        \"board: 2rqrbk1/pp3ppp/8/3p1N2/3NnnQ1/2P4P/PP3PP1/R4RK1 w - - 0 23\\n\"\n",
    "        \"actual move: 23. Nh6+\\n\"\n",
    "        \"evaluation: actual move - Nh6+ 468cp, best move - Nh6+ 511cp, second best move - Qxf4 -25cp\\n\"\n",
    "        )},\n",
    "    {\"role\": \"assistant\", \"content\": \"comment: nice move, plan to fork king and queen Nh6+ and Nf7+, and black cannot capture knight because of pin .\\n\"},\n",
    "    {\"role\": \"user\", \"content\": (\n",
    "        f\"board: {target_position}\\n\"\n",
    "        f\"actual move: {target_move_full}\\n\"\n",
    "        f\"{engine_eval}\\n\"\n",
    "        )},\n",
    "    ]\n",
    "    print(prompt)\n",
    "    response = client.chat.completions.create(model=\"gpt-4o\", messages=prompt, temperature=0.01)\n",
    "    engine_responses_4o.append((idx, response.choices[0].message.content))\n",
    "\n",
    "print(engine_responses_4o)\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_file = \"results/pipeline_v9.log\"\n",
    "with open(log_file, \"w\") as f: \n",
    "    for idx in range(len(target_fens)):\n",
    "        target_position = target_fens[idx]\n",
    "        target_move = target_moves_strip[idx]\n",
    "        target_move_full = target_moves[idx]\n",
    "\n",
    "        f.write(\n",
    "            str((\n",
    "                target_position, target_move_full, \n",
    "                \"vanilla \" + base_responses_4o[idx][1], \n",
    "                \"engine \" + engine_responses_4o[idx][1], \n",
    "                responses_4o[idx][1], responses_4o[idx][2],\n",
    "            )) + \"\\n\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### load from log file\n",
    "\n",
    "import ast \n",
    "\n",
    "log_file = \"results/pipeline_v9.log\"\n",
    "base_responses_4o = []\n",
    "engine_responses_4o = []\n",
    "responses_4o = []\n",
    "responses_4o_wrong_concept = []\n",
    "\n",
    "with open(log_file, \"r\") as f:\n",
    "    for idx, line in enumerate(f):\n",
    "        data_tuple = ast.literal_eval(line.strip())\n",
    "\n",
    "        target_position, target_move_full, vanilla_response, engine_response, response1, response2, wrong_concept_response1, wrong_concept_response2 = data_tuple\n",
    "\n",
    "        base_responses_4o.append((idx, vanilla_response.replace(\"vanilla \", \"\"))) \n",
    "        engine_responses_4o.append((idx, engine_response.replace(\"engine \", \"\"))) \n",
    "        responses_4o.append((idx, response1, response2))  \n",
    "        responses_4o_wrong_concept.append((idx, wrong_concept_response1, wrong_concept_response2))  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 50\n",
    "\n",
    "svg = plot_board(target_fens[idx], target_moves_strip[idx], is_san=True)\n",
    "print(target_fens[idx])\n",
    "print(target_moves[idx])\n",
    "print(ref_list[idx].replace(\"\\n\", \" \"))\n",
    "print(gac_list[idx].replace(\"\\n\", \" \"))\n",
    "print(base_responses_4o[idx][1].replace(\"comment: \", \"\").replace(\"\\n\", \" \"))\n",
    "print(engine_responses_4o[idx][1].replace(\"comment: \", \"\").replace(\"\\n\", \" \"))\n",
    "print(responses_4o[idx][1])\n",
    "print(responses_4o[idx][2])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}