chore: clean up lib/analysis

Author: kevinjosethomas

src/hooks/useAnalysisController/useAnalysisController.ts

@@ -20,12 +20,13 @@ import { useMoveRecommendations } from './useMoveRecommendations'
 import { MaiaEngineContext } from 'src/contexts/MaiaEngineContext'
 import { generateColorSanMapping, calculateBlunderMeter } from './utils'
 import { StockfishEngineContext } from 'src/contexts/StockfishEngineContext'
+import { storeGameAnalysisCache } from 'src/api/analysis'
 import {
+  extractPlayerMistakes,
+  isBestMove,
   collectEngineAnalysisData,
   generateAnalysisCacheKey,
-} from 'src/lib/analysisStorage'
-import { storeGameAnalysisCache } from 'src/api/analysis'
-import { extractPlayerMistakes, isBestMove } from 'src/lib/analysis'
+} from 'src/lib/analysis'
 import { LearnFromMistakesState, MistakePosition } from 'src/types/analysis'
 import { LEARN_FROM_MISTAKES_DEPTH } from 'src/constants/analysis'
 

src/lib/analysis.ts

@@ -1,10 +1,13 @@
+import { Chess } from 'chess.ts'
 import { cpToWinrate } from './stockfish'
 import {
   GameTree,
   GameNode,
   RawMove,
+  MistakePosition,
   MoveValueMapping,
   StockfishEvaluation,
+  CachedEngineAnalysisEntry,
 } from 'src/types'
 
 export function convertBackendEvalToStockfishEval(
@@ -108,3 +111,233 @@ export function insertBackendStockfishEvalToGameTree(
     currentNode = currentNode?.mainChild
   }
 }
+
+export const collectEngineAnalysisData = (
+  gameTree: GameTree,
+): CachedEngineAnalysisEntry[] => {
+  const positions: CachedEngineAnalysisEntry[] = []
+  const mainLine = gameTree.getMainLine()
+
+  mainLine.forEach((node, index) => {
+    if (!node.analysis.maia && !node.analysis.stockfish) {
+      return
+    }
+
+    const position: CachedEngineAnalysisEntry = {
+      ply: index,
+      fen: node.fen,
+    }
+
+    if (node.analysis.maia) {
+      position.maia = node.analysis.maia
+    }
+
+    if (node.analysis.stockfish) {
+      position.stockfish = {
+        depth: node.analysis.stockfish.depth,
+        cp_vec: node.analysis.stockfish.cp_vec,
+      }
+    }
+
+    positions.push(position)
+  })
+
+  return positions
+}
+
+const reconstructCachedStockfishAnalysis = (
+  cpVec: { [move: string]: number },
+  depth: number,
+  fen: string,
+) => {
+  const board = new Chess(fen)
+  const isBlackTurn = board.turn() === 'b'
+
+  let bestCp = isBlackTurn ? Infinity : -Infinity
+  let bestMove = ''
+
+  for (const move in cpVec) {
+    const cp = cpVec[move]
+    if (isBlackTurn) {
+      if (cp < bestCp) {
+        bestCp = cp
+        bestMove = move
+      }
+    } else {
+      if (cp > bestCp) {
+        bestCp = cp
+        bestMove = move
+      }
+    }
+  }
+
+  const cp_relative_vec: { [move: string]: number } = {}
+  for (const move in cpVec) {
+    const cp = cpVec[move]
+    cp_relative_vec[move] = isBlackTurn ? bestCp - cp : cp - bestCp
+  }
+
+  const winrate_vec: { [move: string]: number } = {}
+  for (const move in cpVec) {
+    const cp = cpVec[move]
+    const winrate = cpToWinrate(cp * (isBlackTurn ? -1 : 1), false)
+    winrate_vec[move] = winrate
+  }
+
+  let bestWinrate = -Infinity
+  for (const move in winrate_vec) {
+    const wr = winrate_vec[move]
+    if (wr > bestWinrate) {
+      bestWinrate = wr
+    }
+  }
+
+  const winrate_loss_vec: { [move: string]: number } = {}
+  for (const move in winrate_vec) {
+    winrate_loss_vec[move] = winrate_vec[move] - bestWinrate
+  }
+
+  const sortedEntries = Object.entries(winrate_vec).sort(
+    ([, a], [, b]) => b - a,
+  )
+
+  const sortedWinrateVec = Object.fromEntries(sortedEntries)
+  const sortedWinrateLossVec = Object.fromEntries(
+    sortedEntries.map(([move]) => [move, winrate_loss_vec[move]]),
+  )
+
+  return {
+    sent: true,
+    depth,
+    model_move: bestMove,
+    model_optimal_cp: bestCp,
+    cp_vec: cpVec,
+    cp_relative_vec,
+    winrate_vec: sortedWinrateVec,
+    winrate_loss_vec: sortedWinrateLossVec,
+  }
+}
+
+export const applyEngineAnalysisData = (
+  gameTree: GameTree,
+  analysisData: CachedEngineAnalysisEntry[],
+): void => {
+  const mainLine = gameTree.getMainLine()
+
+  analysisData.forEach((positionData) => {
+    const { ply, maia, stockfish } = positionData
+
+    if (ply >= 0 && ply < mainLine.length) {
+      const node = mainLine[ply]
+
+      if (node.fen === positionData.fen) {
+        if (maia) {
+          node.addMaiaAnalysis(maia)
+        }
+
+        if (stockfish) {
+          const stockfishEval = reconstructCachedStockfishAnalysis(
+            stockfish.cp_vec,
+            stockfish.depth,
+            node.fen,
+          )
+
+          if (
+            !node.analysis.stockfish ||
+            node.analysis.stockfish.depth < stockfish.depth
+          ) {
+            node.addStockfishAnalysis(stockfishEval)
+          }
+        }
+      }
+    }
+  })
+}
+
+export const generateAnalysisCacheKey = (
+  analysisData: CachedEngineAnalysisEntry[],
+): string => {
+  const keyData = analysisData.map((pos) => ({
+    ply: pos.ply,
+    fen: pos.fen,
+    hasStockfish: !!pos.stockfish,
+    stockfishDepth: pos.stockfish?.depth || 0,
+    hasMaia: !!pos.maia,
+    maiaModels: pos.maia ? Object.keys(pos.maia).sort() : [],
+  }))
+
+  return JSON.stringify(keyData)
+}
+
+export function extractPlayerMistakes(
+  gameTree: GameTree,
+  playerColor: 'white' | 'black',
+): MistakePosition[] {
+  const mainLine = gameTree.getMainLine()
+  const mistakes: MistakePosition[] = []
+
+  for (let i = 1; i < mainLine.length; i++) {
+    const node = mainLine[i]
+    const isPlayerMove = node.turn === (playerColor === 'white' ? 'b' : 'w')
+
+    if (
+      isPlayerMove &&
+      (node.blunder || node.inaccuracy) &&
+      node.move &&
+      node.san
+    ) {
+      const parentNode = node.parent
+      if (!parentNode) continue
+
+      const stockfishEval = parentNode.analysis.stockfish
+      if (!stockfishEval || !stockfishEval.model_move) continue
+
+      const chess = new Chess(parentNode.fen)
+      const bestMoveResult = chess.move(stockfishEval.model_move, {
+        sloppy: true,
+      })
+      if (!bestMoveResult) continue
+
+      mistakes.push({
+        nodeId: `move-${i}`, // Simple ID based on position in main line
+        moveIndex: i, // Index of the mistake node in the main line
+        fen: parentNode.fen, // Position before the mistake
+        playedMove: node.move,
+        san: node.san,
+        type: node.blunder ? 'blunder' : 'inaccuracy',
+        bestMove: stockfishEval.model_move,
+        bestMoveSan: bestMoveResult.san,
+        playerColor,
+      })
+    }
+  }
+
+  return mistakes
+}
+
+export function getBestMoveForPosition(node: GameNode): {
+  move: string
+  san: string
+} | null {
+  const stockfishEval = node.analysis.stockfish
+  if (!stockfishEval || !stockfishEval.model_move) {
+    return null
+  }
+
+  const chess = new Chess(node.fen)
+  const moveResult = chess.move(stockfishEval.model_move, { sloppy: true })
+
+  if (!moveResult) {
+    return null
+  }
+
+  return {
+    move: stockfishEval.model_move,
+    san: moveResult.san,
+  }
+}
+
+export function isBestMove(node: GameNode, moveUci: string): boolean {
+  const bestMove = getBestMoveForPosition(node)
+  return bestMove ? bestMove.move === moveUci : false
+}