shared/bot/PuctSearch.ts

import { deserializeCoord, isBoardCoordInBounds, serializeCoord } from '../board.ts'
import { checkWin } from '../checkWin.ts'
import { SeededRandom, weightedIndex } from '../random.ts'
import type { Cell, GameRules } from '../types.ts'
import {
  CH_CURRENT_THREATS,
  CH_OPPONENT_THREATS,
  CH_OVERLINE_TRAPS,
  encodeFeaturePlanes,
} from './FeatureEncoder.ts'
import type {
  PlayerColor,
  PolicyTarget,
  Swap2ColorDecision,
  Swap2Phase,
  Swap2Phase1Decision,
} from '../dataset/types.ts'
import { BitBoard, bitBoardCheckWin } from './BitBoard.ts'
import type { PuctTranspositionEntry, PuctTranspositionTable } from './PuctTranspositionTable.ts'

export interface PuctOptions {
  simulations: number
  cPuct: number
  candidateRadius: number
  maxCandidates: number
  enableTranspositionTable?: boolean
  transpositionTable?: PuctTranspositionTable
  generation?: number
}

export type MaybePromise<T> = T | Promise<T>

export interface PuctContext {
  board: Record<string, Cell>
  currentPlayer: PlayerColor
  playerSymbols: [string, string]
  rules: GameRules
  boardSize: number
  phase: Swap2Phase
  moveNumber: number
}

export interface PuctSearchResult {
  selectedMove?: { x: number; y: number }
  decision?: Swap2Phase1Decision | Swap2ColorDecision
  policyTarget: PolicyTarget
  mctsValue: number
  rootActionStats: PuctRootActionStats[]
  metadata: SearchStateMetadata
}

export type SearchPhase = Swap2Phase | 'offer_place_1' | 'offer_place_2'

export type SearchAction =
  | { kind: 'place'; x: number; y: number }
  | { kind: 'special'; decision: Swap2Phase1Decision | Swap2ColorDecision }

export interface PuctRootActionStats {
  action: SearchAction
  visits: number
  totalValue: number
  totalSquaredValue: number
  meanValue: number
  variance: number
  standardError: number
}

export interface SearchState {
  board: Map<string, Cell>
  bitBoard: BitBoard
  phase: SearchPhase
  actorSeat: PlayerColor
  rootActorSeat: PlayerColor
  moveNumber: number
  rules: GameRules
  boardSize: number
  playerSymbols: [string, string]
  stoneColor?: PlayerColor
  playerToMove?: PlayerColor
  seatColors?: [PlayerColor, PlayerColor]
  terminalValue?: number
}

export interface SearchStateMetadata {
  actorSeat: PlayerColor
  actorColor?: PlayerColor
  stoneColor?: PlayerColor
  playerToMove?: PlayerColor
}

interface SearchEdge {
  action: SearchAction
  prior: number
  visits: number
  totalValue: number
  totalSquaredValue: number
  child?: SearchNode
}

interface SearchNode {
  state: SearchState
  expanded: boolean
  value: number
  edges: SearchEdge[]
  ttKey?: PuctStateKey
}

interface PuctStateKey {
  key: string
  transform: BoardTransform
}

type BoardTransform = 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7

const SWAP2_MOBILITY_WEIGHT = 350
const SWAP2_THREAT_WEIGHT = 0.75

export interface PolicyValueEvaluator {
  evaluate(state: SearchState, legalActions: SearchAction[]): MaybePromise<{ value: number; priors: number[] }>
}

export class HeuristicPolicyValueEvaluator implements PolicyValueEvaluator {
  evaluate(state: SearchState, legalActions: SearchAction[]): { value: number; priors: number[] } {
    return {
      value: clamp(evaluateState(state), -1, 1),
      priors: normalizePriors(legalActions.map(action => Math.max(0.001, actionPrior(state, action)))),
    }
  }
}

export async function runPuctSearch(
  context: PuctContext,
  options: PuctOptions,
  temperature: number,
  rng: SeededRandom,
  evaluator: PolicyValueEvaluator = new HeuristicPolicyValueEvaluator(),
): Promise<PuctSearchResult> {
  const table = options.enableTranspositionTable === false ? undefined : options.transpositionTable
  const generation = options.generation ?? 1
  const root = createNode(rootState(context), table, generation)
  await expand(root, evaluator, options, rng)

  for (let i = 0; i < options.simulations; i++) {
    await runSimulation(root, evaluator, options, rng, table, generation)
  }

  writeNodeToTranspositionTable(root, table, generation)

  const selected = selectRootEdge(root.edges, temperature, rng)
  const rootVisits = root.edges.reduce((sum, edge) => sum + edge.visits, 0)
  const rootValue = rootVisits > 0
    ? root.edges.reduce((sum, edge) => sum + edge.totalValue, 0) / rootVisits
    : root.value

  return {
    selectedMove: selected.action.kind === 'place' ? { x: selected.action.x, y: selected.action.y } : undefined,
    decision: selected.action.kind === 'special' ? selected.action.decision : undefined,
    policyTarget: policyFromEdges(root.edges),
    mctsValue: clamp(rootValue, -1, 1),
    rootActionStats: rootStatsFromEdges(root.edges),
    metadata: metadataFromState(root.state),
  }
}

function createNode(state: SearchState, table?: PuctTranspositionTable, generation = 1): SearchNode {
  const ttKey = canonicalStateKey(state)
  const entry = table?.get(ttKey.key, generation)
  const node: SearchNode = { state, expanded: false, value: state.terminalValue ?? entry?.value ?? 0, edges: [], ttKey }
  if (entry && table) hydrateNodeFromTranspositionEntry(node, entry, table.reuse)
  return node
}

function hydrateNodeFromTranspositionEntry(
  node: SearchNode,
  entry: PuctTranspositionEntry,
  reuse: 'priors' | 'stats',
): void {
  if (node.state.terminalValue !== undefined || entry.actions.length === 0 || !node.ttKey) return
  const inverse = inverseTransform(node.ttKey.transform)
  node.value = entry.value
  node.expanded = true
  node.edges = entry.actions.map(action => ({
    action: transformAction(action.action, inverse, node.state.boardSize),
    prior: action.prior,
    visits: reuse === 'stats' ? action.visits : 0,
    totalValue: reuse === 'stats' ? action.totalValue : 0,
    totalSquaredValue: reuse === 'stats' ? action.totalSquaredValue ?? 0 : 0,
  }))
}

function writeNodeToTranspositionTable(
  node: SearchNode,
  table: PuctTranspositionTable | undefined,
  generation: number,
): void {
  if (!table || !node.expanded || !node.ttKey) return
  const visits = node.edges.reduce((sum, edge) => sum + edge.visits, 0)
  const totalValue = node.edges.reduce((sum, edge) => sum + edge.totalValue, 0)
  const totalSquaredValue = node.edges.reduce((sum, edge) => sum + edge.totalSquaredValue, 0)
  table.set({
    key: node.ttKey.key,
    generation,
    value: node.value,
    visits,
    totalValue,
    totalSquaredValue,
    actions: node.edges.map(edge => ({
      action: transformAction(edge.action, node.ttKey?.transform ?? 0, node.state.boardSize),
      prior: edge.prior,
      visits: edge.visits,
      totalValue: edge.totalValue,
      totalSquaredValue: edge.totalSquaredValue,
    })),
  })
}

function canonicalStateKey(state: SearchState): PuctStateKey {
  const canonical = canonicalBitBoard(state.bitBoard)
  return {
    key: [
      canonical.key,
      state.phase,
      state.actorSeat,
      state.rootActorSeat,
      state.moveNumber,
      state.stoneColor ?? '-',
      state.playerToMove ?? '-',
      state.seatColors ? state.seatColors.join(':') : '-',
      state.rules.noOverlines ? 'no-overlines' : 'overlines',
    ].join('|'),
    transform: canonical.transform,
  }
}

export function canonicalBitBoardKey(board: BitBoard): string {
  return canonicalBitBoard(board).key
}

function canonicalBitBoard(board: BitBoard): { key: string; transform: BoardTransform } {
  let bestKey = ''
  let bestTransform: BoardTransform = 0
  for (let transform = 0 as BoardTransform; transform < 8; transform = (transform + 1) as BoardTransform) {
    const transformed = BitBoard.empty(board.boardSize)
    for (let y = 0; y < board.boardSize; y++) {
      for (let x = 0; x < board.boardSize; x++) {
        const player = board.playerAt(x, y)
        if (player === null) continue
        const coord = transformCoord({ x, y }, transform, board.boardSize)
        transformed.place(coord.x, coord.y, player)
      }
    }
    const key = transformed.key()
    if (bestKey === '' || key < bestKey) {
      bestKey = key
      bestTransform = transform
    }
  }
  return { key: bestKey, transform: bestTransform }
}

function transformAction(action: SearchAction, transform: BoardTransform, boardSize: number): SearchAction {
  if (action.kind === 'special') return action
  const coord = transformCoord(action, transform, boardSize)
  return { kind: 'place', x: coord.x, y: coord.y }
}

function transformCoord(coord: { x: number; y: number }, transform: BoardTransform, boardSize: number): { x: number; y: number } {
  const max = boardSize - 1
  if (transform === 0) return coord
  if (transform === 1) return { x: max - coord.y, y: coord.x }
  if (transform === 2) return { x: max - coord.x, y: max - coord.y }
  if (transform === 3) return { x: coord.y, y: max - coord.x }
  if (transform === 4) return { x: max - coord.x, y: coord.y }
  if (transform === 5) return { x: coord.x, y: max - coord.y }
  if (transform === 6) return { x: coord.y, y: coord.x }
  return { x: max - coord.y, y: max - coord.x }
}

function inverseTransform(transform: BoardTransform): BoardTransform {
  if (transform === 1) return 3
  if (transform === 3) return 1
  return transform
}

function rootState(context: PuctContext): SearchState {
  const board = new Map(Object.entries(context.board))
  const bitBoard = BitBoard.fromBoard(board, context.boardSize)
  if (context.phase === 'play') {
    return {
      board,
      bitBoard,
      phase: 'play',
      actorSeat: context.currentPlayer,
      rootActorSeat: context.currentPlayer,
      playerToMove: context.currentPlayer,
      seatColors: [0, 1],
      moveNumber: context.moveNumber,
      rules: context.rules,
      boardSize: context.boardSize,
      playerSymbols: context.playerSymbols,
    }
  }

  if (context.phase === 'swap2_phase1') {
    return baseOpeningState(board, context, 'swap2_phase1', 1)
  }

  if (context.phase === 'swap2_phase2') {
    return baseOpeningState(board, context, 'swap2_phase2', 0)
  }

  const actorSeat: PlayerColor = context.moveNumber <= 3 ? 0 : 1
  return {
    ...baseOpeningState(board, context, 'swap2_phase0', actorSeat),
    stoneColor: context.currentPlayer,
  }
}

function baseOpeningState(
  board: Map<string, Cell>,
  context: PuctContext,
  phase: SearchPhase,
  actorSeat: PlayerColor,
): SearchState {
  return {
    board,
    bitBoard: BitBoard.fromBoard(board, context.boardSize),
    phase,
    actorSeat,
    rootActorSeat: actorSeat,
    moveNumber: context.moveNumber,
    rules: context.rules,
    boardSize: context.boardSize,
    playerSymbols: context.playerSymbols,
  }
}

async function runSimulation(
  root: SearchNode,
  evaluator: PolicyValueEvaluator,
  options: PuctOptions,
  rng: SeededRandom,
  table?: PuctTranspositionTable,
  generation = 1,
): Promise<void> {
  let node = root
  const path: SearchEdge[] = []
  const pathNodes: SearchNode[] = [root]

  while (node.expanded && node.edges.length > 0 && node.state.terminalValue === undefined) {
    const edge = selectPuctEdge(node, options, rng)
    if (!edge.child) edge.child = createNode(nextState(node.state, edge.action), table, generation)
    path.push(edge)
    node = edge.child
    pathNodes.push(node)
  }

  const value = node.state.terminalValue ?? await expand(node, evaluator, options, rng)
  for (const edge of path) {
    edge.visits++
    edge.totalValue += value
    edge.totalSquaredValue += value * value
  }
  for (const pathNode of pathNodes) writeNodeToTranspositionTable(pathNode, table, generation)
}

async function expand(
  node: SearchNode,
  evaluator: PolicyValueEvaluator,
  options: PuctOptions,
  rng: SeededRandom,
): Promise<number> {
  if (node.expanded) return node.value
  if (node.state.terminalValue !== undefined) {
    node.value = node.state.terminalValue
    node.expanded = true
    return node.value
  }

  const actions = legalActions(node.state, options)
  const evaluation = await evaluator.evaluate(node.state, actions)
  const priors = normalizePriors(evaluation.priors.length === actions.length ? evaluation.priors : [])
  node.edges = actions.map((action, index) => ({
    action,
    prior: priors[index] ?? 1 / Math.max(1, actions.length),
    visits: 0,
    totalValue: 0,
    totalSquaredValue: 0,
  }))
  if (node.edges.length > 1) {
    const jitter = rng.next() * 0.000001
    node.edges[0].prior += jitter
    node.edges = normalizeEdgePriors(node.edges)
  }
  node.value = clamp(evaluation.value, -1, 1)
  node.expanded = true
  return node.value
}

function selectPuctEdge(node: SearchNode, options: PuctOptions, rng: SeededRandom): SearchEdge {
  const parentVisits = Math.max(1, node.edges.reduce((sum, edge) => sum + edge.visits, 0))
  let bestScore = -Infinity
  let best: SearchEdge[] = []

  for (const edge of node.edges) {
    const q = edge.visits > 0 ? edge.totalValue / edge.visits : 0
    const actorQ = isRootActorTurn(node.state) ? q : -q
    const u = options.cPuct * edge.prior * Math.sqrt(parentVisits) / (1 + edge.visits)
    const score = actorQ + u
    if (score > bestScore) {
      bestScore = score
      best = [edge]
    } else if (score === bestScore) {
      best.push(edge)
    }
  }

  return rng.pick(best)
}

function selectRootEdge(edges: SearchEdge[], temperature: number, rng: SeededRandom): SearchEdge {
  if (edges.length === 0) {
    throw new Error('PUCT search has no legal root actions')
  }
  if (temperature <= 0.05) {
    return edges.reduce((best, edge) => edge.visits > best.visits ? edge : best, edges[0])
  }
  return edges[weightedIndex(edges.map(edge => Math.pow(Math.max(edge.visits, 1), 1 / temperature)), rng)]
}

function legalActions(state: SearchState, options: PuctOptions): SearchAction[] {
  if (state.phase === 'swap2_phase1') {
    return [
      { kind: 'special', decision: 'offer' },
      { kind: 'special', decision: 'player1' },
      { kind: 'special', decision: 'player2' },
    ]
  }
  if (state.phase === 'swap2_phase2') {
    return [
      { kind: 'special', decision: 'player1' },
      { kind: 'special', decision: 'player2' },
    ]
  }
  return generateCandidates(state.bitBoard, options.candidateRadius)
    .map(move => ({ kind: 'place' as const, x: move.x, y: move.y }))
    .sort((a, b) => centerPrior(b, state.boardSize) - centerPrior(a, state.boardSize) || a.y - b.y || a.x - b.x)
    .slice(0, options.maxCandidates)
}

function nextState(state: SearchState, action: SearchAction): SearchState {
  if (action.kind === 'special') return nextSpecialState(state, action.decision)

  const stoneColor = placementColor(state)
  const board = withMove(state.board, action, stoneColor, state.playerSymbols)
  const bitBoard = state.bitBoard.withMove(action.x, action.y, stoneColor)
  const terminal = terminalValueForPlacement(state, bitBoard, action, stoneColor)
  const next = nextPlacementState(state, board, bitBoard)
  return { ...next, terminalValue: terminal ?? next.terminalValue }
}

function nextSpecialState(
  state: SearchState,
  decision: Swap2Phase1Decision | Swap2ColorDecision,
): SearchState {
  if (state.phase === 'swap2_phase1') {
    if (decision === 'offer') {
      return {
        ...state,
        phase: 'offer_place_1',
        actorSeat: 1,
        stoneColor: 1,
      }
    }
    return {
      ...state,
      phase: 'play',
      actorSeat: 1,
      playerToMove: 1,
      seatColors: decision === 'player1' ? [1, 0] : [0, 1],
    }
  }

  return {
    ...state,
    phase: 'play',
    actorSeat: 1,
    playerToMove: 1,
    seatColors: decision === 'player1' ? [0, 1] : [1, 0],
  }
}

function nextPlacementState(state: SearchState, board: Map<string, Cell>, bitBoard: BitBoard): SearchState {
  if (state.phase === 'play') {
    const nextPlayer = otherPlayer(state.playerToMove ?? 0)
    return {
      ...state,
      board,
      bitBoard,
      actorSeat: nextPlayer,
      playerToMove: nextPlayer,
      moveNumber: state.moveNumber + 1,
    }
  }

  if (state.phase === 'offer_place_1') {
    return {
      ...state,
      board,
      bitBoard,
      phase: 'offer_place_2',
      actorSeat: 1,
      stoneColor: 0,
      moveNumber: state.moveNumber + 1,
    }
  }

  if (state.phase === 'offer_place_2') {
    return {
      ...state,
      board,
      bitBoard,
      phase: 'swap2_phase2',
      actorSeat: 0,
      stoneColor: undefined,
      moveNumber: state.moveNumber + 1,
    }
  }

  if (state.phase === 'swap2_phase0' && state.moveNumber <= 2) {
    return {
      ...state,
      board,
      bitBoard,
      actorSeat: 0,
      stoneColor: state.moveNumber === 1 ? 1 : 0,
      moveNumber: state.moveNumber + 1,
    }
  }

  if (state.phase === 'swap2_phase0' && state.moveNumber === 3) {
    return {
      ...state,
      board,
      bitBoard,
      phase: 'swap2_phase1',
      actorSeat: 1,
      stoneColor: undefined,
      moveNumber: state.moveNumber + 1,
    }
  }

  if (state.phase === 'swap2_phase0' && state.moveNumber === 4) {
    return {
      ...state,
      board,
      bitBoard,
      phase: 'offer_place_2',
      actorSeat: 1,
      stoneColor: 0,
      moveNumber: state.moveNumber + 1,
    }
  }

  return {
    ...state,
    board,
    bitBoard,
    phase: 'swap2_phase2',
    actorSeat: 0,
    stoneColor: undefined,
    moveNumber: state.moveNumber + 1,
  }
}

function placementColor(state: SearchState): PlayerColor {
  if (state.phase === 'play') return state.playerToMove ?? state.actorSeat
  return state.stoneColor ?? state.actorSeat
}

function terminalValueForPlacement(
  state: SearchState,
  bitBoard: BitBoard,
  move: { x: number; y: number },
  stoneColor: PlayerColor,
): number | null {
  if (bitBoardCheckWin(bitBoard, move.x, move.y, stoneColor, state.rules.noOverlines)) {
    return valueForWinner(state, stoneColor)
  }
  return bitBoard.occupiedCount >= state.boardSize * state.boardSize ? 0 : null
}

function valueForWinner(state: SearchState, winner: PlayerColor): number {
  const rootColor = rootActorColor(state)
  if (rootColor === undefined) return winner === state.rootActorSeat ? 1 : -1
  return winner === rootColor ? 1 : -1
}

function evaluateState(state: SearchState): number {
  const rootColor = rootActorColor(state) ?? fallbackRootColor(state)
  const features = encodeFeaturePlanes({
    board: state.board,
    boardSize: state.boardSize,
    rules: state.rules,
    currentPlayer: rootColor,
    phase: state.phase,
  })
  const cells = state.boardSize * state.boardSize
  let ownThreats = 0
  let opponentThreats = 0
  let overlineTraps = 0
  for (let i = 0; i < cells; i++) {
    ownThreats += features[CH_CURRENT_THREATS * cells + i]
    opponentThreats += features[CH_OPPONENT_THREATS * cells + i]
    overlineTraps += features[CH_OVERLINE_TRAPS * cells + i]
  }
  const lineScore = linePotential(state.board, rootColor, state) - linePotential(state.board, otherPlayer(rootColor), state) * 1.05
  const tempoScore = earlyTempoScore(state, rootColor)
  return Math.tanh((lineScore + tempoScore + ownThreats * 35_000 - opponentThreats * 45_000 - overlineTraps * 20_000) / 100_000)
}

function rootActorColor(state: SearchState): PlayerColor | undefined {
  return state.seatColors?.[state.rootActorSeat]
}

export function valueSignForRootActor(state: SearchState, valueActor: PlayerColor): 1 | -1 {
  const rootColor = rootActorColor(state) ?? fallbackRootColor(state)
  return valueActor === rootColor ? 1 : -1
}

function isRootActorTurn(state: SearchState): boolean {
  const rootColor = rootActorColor(state)
  if (rootColor !== undefined && state.playerToMove !== undefined) return state.playerToMove === rootColor
  return state.actorSeat === state.rootActorSeat
}

function metadataFromState(state: SearchState): SearchStateMetadata {
  return {
    actorSeat: state.actorSeat,
    ...(rootActorColor(state) !== undefined ? { actorColor: rootActorColor(state) } : {}),
    ...(state.stoneColor !== undefined ? { stoneColor: state.stoneColor } : {}),
    ...(state.playerToMove !== undefined ? { playerToMove: state.playerToMove } : {}),
  }
}

function fallbackRootColor(state: SearchState): PlayerColor {
  if (state.stoneColor !== undefined) return state.stoneColor
  return state.rootActorSeat
}

function earlyTempoScore(state: SearchState, rootColor: PlayerColor): number {
  if (state.phase !== 'play' || state.playerToMove === undefined || state.moveNumber > 6) return 0
  const gain = bestImmediateMoveGain(state, state.playerToMove) * 2.5
  return state.playerToMove === rootColor ? gain : -gain
}

function actionPrior(state: SearchState, action: SearchAction): number {
  if (action.kind === 'special') {
    const scores = specialActionScores(state)
    if (action.decision === 'offer') return scores.offer ?? 1
    if (action.decision === 'player1') return scores.player1 ?? 1
    return scores.player2 ?? 1
  }

  const color = placementColor(state)
  const next = withMove(state.board, action, color, state.playerSymbols)
  const ownWin = checkWin(next, action.x, action.y, String(color), state.rules.noOverlines).won
  if (ownWin) return 1_000_000
  const opponent = otherPlayer(color)
  const blockBoard = withMove(state.board, action, opponent, state.playerSymbols)
  const blocksWin = checkWin(blockBoard, action.x, action.y, String(opponent), state.rules.noOverlines).won
  if (blocksWin) return 900_000
  return Math.max(1, linePotential(next, color, state) - linePotential(state.board, color, state) + centerPrior(action, state.boardSize))
}

function specialActionScores(state: SearchState): Partial<Record<Swap2Phase1Decision, number>> {
  if (state.phase === 'swap2_phase1') {
    const player1 = swap2DecisionUtility(state, 'player1')
    const player2 = swap2DecisionUtility(state, 'player2')
    return positiveSpecialScores({
      player1,
      player2,
      offer: Math.min(player1, player2) + Math.abs(player1 - player2) * 0.35,
    })
  }
  return positiveSpecialScores({
    player1: swap2DecisionUtility(state, 'player1'),
    player2: swap2DecisionUtility(state, 'player2'),
  })
}

function swap2DecisionUtility(
  state: SearchState,
  decision: Swap2Phase1Decision | Swap2ColorDecision,
): number {
  const next = nextSpecialState(state, decision)
  const actorColor = rootActorColor(next) ?? fallbackRootColor(next)
  const opponent = otherPlayer(actorColor)
  const actorLineScore = linePotential(state.board, actorColor, state)
  const opponentLineScore = linePotential(state.board, opponent, state)
  const nextMover = next.playerToMove
  const nextMoveGain = nextMover === undefined ? 0 : bestImmediateMoveGain(state, nextMover)
  const tempoScore = nextMover === actorColor ? nextMoveGain * 2.5 : -nextMoveGain * 2.5
  const threatScore = maxThreatDanger(state, actorColor) - maxThreatDanger(state, opponent) * 1.1
  const mobilityScore = mobilityPotential(state, actorColor) - mobilityPotential(state, opponent) * 1.05
  return (
    actorLineScore
    - opponentLineScore * 1.05
    + tempoScore
    + threatScore * SWAP2_THREAT_WEIGHT
    + mobilityScore * SWAP2_MOBILITY_WEIGHT
  )
}

function bestImmediateMoveGain(state: SearchState, color: PlayerColor): number {
  let best = 0
  for (const move of generateCandidates(state.bitBoard, 2).slice(0, 24)) {
    const next = withMove(state.board, move, color, state.playerSymbols)
    if (checkWin(next, move.x, move.y, String(color), state.rules.noOverlines).won) return 1_000_000
    best = Math.max(best, linePotential(next, color, state) - linePotential(state.board, color, state) + centerPrior(move, state.boardSize))
  }
  return best
}

function maxThreatDanger(state: SearchState, color: PlayerColor): number {
  let danger = 0
  for (const move of generateCandidates(state.bitBoard, 2).slice(0, 24)) {
    const next = withMove(state.board, move, color, state.playerSymbols)
    if (checkWin(next, move.x, move.y, String(color), state.rules.noOverlines).won) return 1_000_000
    danger = Math.max(danger, linePotential(next, color, state) - linePotential(state.board, color, state))
  }
  return danger
}

function mobilityPotential(state: SearchState, color: PlayerColor): number {
  let mobility = 0
  for (const move of generateCandidates(state.bitBoard, 2).slice(0, 24)) {
    const next = withMove(state.board, move, color, state.playerSymbols)
    const gain = linePotential(next, color, state) - linePotential(state.board, color, state)
    if (gain >= 20_000) mobility += 6
    else if (gain >= 1_000) mobility += 2
    else if (gain > 0) mobility += 0.5
  }
  return mobility
}

function positiveSpecialScores<T extends string>(scores: Record<T, number>): Record<T, number> {
  const values = Object.values(scores) as number[]
  const min = Math.min(...values)
  const shifted = {} as Record<T, number>
  for (const [key, value] of Object.entries(scores) as [T, number][]) {
    shifted[key] = Math.max(1, value - min + 1)
  }
  return shifted
}

function centerPrior(move: { x: number; y: number }, boardSize: number): number {
  const center = (boardSize - 1) / 2
  return Math.max(1, boardSize - Math.abs(move.x - center) - Math.abs(move.y - center))
}

function policyFromEdges(edges: SearchEdge[]): PolicyTarget {
  const placements = edges.filter(edge => edge.action.kind === 'place')
  if (placements.length > 0) {
    const total = placements.reduce((sum, edge) => sum + edge.visits, 0)
    return {
      placements: placements.map(edge => {
        const action = edge.action as Extract<SearchAction, { kind: 'place' }>
        return {
          x: action.x,
          y: action.y,
          probability: total > 0 ? edge.visits / total : 1 / placements.length,
        }
      }),
    }
  }

  const total = edges.reduce((sum, edge) => sum + edge.visits, 0)
  const specialActions: NonNullable<PolicyTarget['specialActions']> = {}
  for (const edge of edges) {
    if (edge.action.kind !== 'special') continue
    const probability = total > 0 ? edge.visits / total : 1 / edges.length
    if (edge.action.decision === 'offer') specialActions.offerTwo = probability
    if (edge.action.decision === 'player1') specialActions.choosePlayer1 = probability
    if (edge.action.decision === 'player2') specialActions.choosePlayer2 = probability
  }
  return { placements: [], specialActions }
}

function rootStatsFromEdges(edges: SearchEdge[]): PuctRootActionStats[] {
  return edges.map(edge => {
    const meanValue = edge.visits > 0 ? edge.totalValue / edge.visits : 0
    const variance = edge.visits > 0
      ? Math.max(0, edge.totalSquaredValue / edge.visits - meanValue * meanValue)
      : 0
    return {
      action: edge.action,
      visits: edge.visits,
      totalValue: edge.totalValue,
      totalSquaredValue: edge.totalSquaredValue,
      meanValue,
      variance,
      standardError: edge.visits > 0 ? Math.sqrt(variance) / Math.sqrt(edge.visits) : 0,
    }
  })
}

function generateCandidates(board: BitBoard, radius: number): { x: number; y: number }[] {
  if (board.occupiedCount === 0) {
    const boardSize = board.boardSize
    const center = Math.floor(boardSize / 2)
    return [{ x: center, y: center }]
  }

  const candidates = new Map<string, { x: number; y: number }>()
  for (let y = 0; y < board.boardSize; y++) {
    for (let x = 0; x < board.boardSize; x++) {
      if (!board.has(x, y)) continue
      for (let dy = -radius; dy <= radius; dy++) {
        for (let dx = -radius; dx <= radius; dx++) {
          const nx = x + dx
          const ny = y + dy
          const nKey = serializeCoord(nx, ny)
          if (!isBoardCoordInBounds(nx, ny, board.boardSize) || board.has(nx, ny)) continue
          candidates.set(nKey, { x: nx, y: ny })
        }
      }
    }
  }

  if (candidates.size === 0) {
    for (let y = 0; y < board.boardSize; y++) {
      for (let x = 0; x < board.boardSize; x++) {
        const key = serializeCoord(x, y)
        if (!board.has(x, y)) candidates.set(key, { x, y })
      }
    }
  }

  return [...candidates.values()].sort((a, b) => a.y - b.y || a.x - b.x)
}

function withMove(
  board: Map<string, Cell>,
  move: { x: number; y: number },
  player: PlayerColor,
  symbols: [string, string],
): Map<string, Cell> {
  const next = new Map(board)
  next.set(serializeCoord(move.x, move.y), {
    playerId: String(player),
    symbol: symbols[player],
    timestamp: 0,
  })
  return next
}

function linePotential(board: Map<string, Cell>, player: PlayerColor, state: Pick<SearchState, 'boardSize' | 'rules'>): number {
  const directions: readonly [number, number][] = [[1, 0], [0, 1], [1, 1], [1, -1]]
  let score = 0
  const playerId = String(player)
  const center = (state.boardSize - 1) / 2

  for (const [key, cell] of board) {
    if (cell.playerId !== playerId) continue
    const { x, y } = deserializeCoord(key)
    score += Math.max(0, state.boardSize - Math.abs(x - center) - Math.abs(y - center)) * 4
    for (const [dx, dy] of directions) {
      let count = 1
      let openEnds = 0
      let cx = x + dx
      let cy = y + dy
      while (board.get(serializeCoord(cx, cy))?.playerId === playerId) {
        count++
        cx += dx
        cy += dy
      }
      if (isBoardCoordInBounds(cx, cy, state.boardSize) && !board.has(serializeCoord(cx, cy))) openEnds++
      cx = x - dx
      cy = y - dy
      while (board.get(serializeCoord(cx, cy))?.playerId === playerId) {
        count++
        cx -= dx
        cy -= dy
      }
      if (isBoardCoordInBounds(cx, cy, state.boardSize) && !board.has(serializeCoord(cx, cy))) openEnds++
      if (state.rules.noOverlines && count > 5) continue
      score += Math.pow(10, Math.min(count, 5)) * (openEnds + 1)
    }
  }

  return score
}

function normalizePriors(values: number[]): number[] {
  if (values.length === 0) return []
  const clean = values.map(value => Number.isFinite(value) && value > 0 ? value : 0)
  const total = clean.reduce((sum, value) => sum + value, 0)
  if (total <= 0) return values.map(() => 1 / values.length)
  return clean.map(value => value / total)
}

function normalizeEdgePriors(edges: SearchEdge[]): SearchEdge[] {
  const priors = normalizePriors(edges.map(edge => edge.prior))
  return edges.map((edge, index) => ({ ...edge, prior: priors[index] }))
}

function otherPlayer(player: PlayerColor): PlayerColor {
  return player === 0 ? 1 : 0
}

function clamp(value: number, min: number, max: number): number {
  return Math.max(min, Math.min(max, value))
}