Blame for shared/bot/BotDefinitions.ts

import { BOARD_SIZE, deserializeCoord, isBoardCoordInBounds, serializeCoord } from '../board.ts'
import { checkWin } from '../checkWin.ts'
import type { Cell, GameRules } from '../types.ts'
import { SeededRandom, weightedIndex } from '../random.ts'
import type {
  MctsTarget,
  PlayerColor,
  PolicyTarget,
  Swap2ColorDecision,
  Swap2Phase,
  Swap2Phase1Decision,
} from '../dataset/types.ts'
import { MinimaxAlphaBetaBot } from './MinimaxAlphaBetaBot.ts'
import { runPuctSearch } from './PuctSearch.ts'
import type { PolicyValueEvaluator, PuctRootActionStats } from './PuctSearch.ts'
import type { PuctTranspositionTable } from './PuctTranspositionTable.ts'

export interface BotMove {
  x: number
  y: number
}

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

export interface BotMoveDecision {
  move: BotMove
  mctsTarget?: MctsTarget
}

export interface DatasetBot {
  id: string
  params: Record<string, number | string | boolean>
  selectMove(context: DatasetBotContext, rng: SeededRandom): Promise<BotMoveDecision>
  chooseSwap2Phase1(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2Phase1Decision; mctsTarget?: MctsTarget }>
  chooseSwap2Phase2(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2ColorDecision; mctsTarget?: MctsTarget }>
}

export interface DatasetBotOptions {
  puctTranspositionTable?: PuctTranspositionTable
  puctGeneration?: number
  modelEvaluator?: PolicyValueEvaluator
}

interface MctsOptions {
  simulations: number
  rolloutDepth: number
  exploration: number
  candidateRadius: number
  maxCandidates?: number
  earlyTemperature: number
  lateTemperature: number
}

const PLAYER_SYMBOLS: [string, string] = ['X', 'O']
const BOT_IDS = [
  'random',
  'greedy',
  'minimax_weak',
  'minimax_balanced',
  'minimax_strong',
  'mcts_weak',
  'mcts_balanced',
  'mcts_strong',
  'mcts_model_guided',
] as const
const THREAT_WIN = 1_000_000
const THREAT_OPEN_FOUR = 120_000
const THREAT_CLOSED_FOUR = 80_000
const THREAT_BROKEN_FOUR = 70_000
const THREAT_OPEN_THREE = 20_000
const SWAP2_OFFER_VALUE_MARGIN = 0.08
const SWAP2_OFFER_UNCERTAIN_MARGIN = 0.16
const SWAP2_OFFER_STANDARD_ERROR = 0.1
const SWAP2_OFFER_MIN_VISITS = 4
const SWAP2_UTILITY_SCALE = 10_000
const SWAP2_MOBILITY_WEIGHT = 350
const SWAP2_THREAT_WEIGHT = 0.75
const SWAP2_NEUTRAL_OPENING_LOOKAHEAD = 32
const SWAP2_NEUTRAL_OPENING_POLICY_SIZE = 8
const SELF_PLAY_BOT_IDS = [
  'minimax_balanced',
  'minimax_strong',
  'mcts_balanced',
  'mcts_strong',
] as const

export type DatasetBotId = typeof BOT_IDS[number]

export function getDatasetBotIds(): DatasetBotId[] {
  return [...BOT_IDS]
}

export function getSelfPlayBotIds(): DatasetBotId[] {
  return [...SELF_PLAY_BOT_IDS]
}

export function createDatasetBot(id: DatasetBotId, options: DatasetBotOptions = {}): DatasetBot {
  if (id === 'random') return new RandomBot()
  if (id === 'greedy') return new GreedyBot()
  if (id === 'minimax_weak') return new MinimaxDatasetBot(id, { maxDepth: 1, timeBudgetMs: 80, candidateRadius: 1, panicPly: 1 })
  if (id === 'minimax_balanced') return new MinimaxDatasetBot(id, { maxDepth: 2, timeBudgetMs: 180, candidateRadius: 2, panicPly: 1 })
  if (id === 'minimax_strong') return new MinimaxDatasetBot(id, { maxDepth: 3, timeBudgetMs: 320, candidateRadius: 2, panicPly: 2 })
  if (id === 'mcts_weak') return new MctsDatasetBot(id, { simulations: 48, rolloutDepth: 12, exploration: 1.4, candidateRadius: 1, earlyTemperature: 1, lateTemperature: 0.15 }, options)
  if (id === 'mcts_balanced') return new MctsDatasetBot(id, { simulations: 96, rolloutDepth: 18, exploration: 1.2, candidateRadius: 2, earlyTemperature: 0.75, lateTemperature: 0.08 }, options)
  if (id === 'mcts_strong') return new MctsDatasetBot(id, { simulations: 192, rolloutDepth: 24, exploration: 1.05, candidateRadius: 2, earlyTemperature: 0.6, lateTemperature: 0.03 }, options)
  return new MctsDatasetBot(id, { simulations: 96, rolloutDepth: 0, exploration: 1.2, candidateRadius: 2, maxCandidates: 24, earlyTemperature: 0.75, lateTemperature: 0.08 }, options)
}

function minimaxAttitude(id: DatasetBotId): string {
  if (id === 'minimax_weak') return 'shallow_tactical'
  if (id === 'minimax_strong') return 'deep_defensive'
  return 'balanced_tactical'
}

function mctsAttitude(id: DatasetBotId): string {
  if (id === 'mcts_weak') return 'exploratory_rollout'
  if (id === 'mcts_strong') return 'low_noise_search'
  if (id === 'mcts_model_guided') return 'model_guided_search'
  return 'balanced_search'
}

class RandomBot implements DatasetBot {
  readonly id = 'random'
  readonly params = { strategy: 'random', attitude: 'exploratory', noise: 1, aggressionWeight: 0.4, defensiveBias: 0.35 }

  async selectMove(context: DatasetBotContext, rng: SeededRandom): Promise<BotMoveDecision> {
    return { move: rng.pick(generateCandidates(toMap(context.board), context.boardSize, 2)) }
  }

  async chooseSwap2Phase1(_context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2Phase1Decision }> {
    return { decision: rng.pick(['offer', 'player1', 'player2'] as const) }
  }

  async chooseSwap2Phase2(_context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2ColorDecision }> {
    return { decision: rng.pick(['player1', 'player2'] as const) }
  }
}

class GreedyBot implements DatasetBot {
  readonly id = 'greedy'
  readonly params = { strategy: 'greedy', attitude: 'aggressive', noise: 0.35, aggressionWeight: 0.8, defensiveBias: 0.45 }

  async selectMove(context: DatasetBotContext, rng: SeededRandom): Promise<BotMoveDecision> {
    const board = toMap(context.board)
    const candidates = generateCandidates(board, context.boardSize, 2)
    let best = candidates[0]
    let bestScore = -Infinity
    for (const move of candidates) {
      const score = scoreMove(board, move, context.currentPlayer, context)
      if (score > bestScore || (score === bestScore && rng.next() < 0.35)) {
        best = move
        bestScore = score
      }
    }
    return { move: best }
  }

  async chooseSwap2Phase1(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2Phase1Decision }> {
    return { decision: choosePhase1Heuristic(toMap(context.board), context, rng, 0.12) }
  }

  async chooseSwap2Phase2(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2ColorDecision }> {
    return { decision: choosePhase2Heuristic(toMap(context.board), context, rng, 0.1) }
  }
}

class MinimaxDatasetBot implements DatasetBot {
  readonly id: DatasetBotId
  readonly params: Record<string, number | string | boolean>
  private readonly bot: MinimaxAlphaBetaBot

  constructor(id: DatasetBotId, options: { maxDepth: number; timeBudgetMs: number; candidateRadius: number; panicPly?: number }) {
    this.id = id
    this.params = { strategy: 'minimax', attitude: minimaxAttitude(id), aggressionWeight: 0.65, defensiveBias: 0.85, ...options }
    this.bot = new MinimaxAlphaBetaBot({ ...options, enableThreatSpaceSearch: false })
  }

  async selectMove(context: DatasetBotContext, rng: SeededRandom): Promise<BotMoveDecision> {
    if (context.phase === 'swap2_phase0') {
      const board = toMap(context.board)
      const candidates = generateCandidates(board, context.boardSize, 2)
      const ranked = rankNeutralOpeningMovesNoLookahead(board, candidates, context)
      const selected = chooseNeutralOpeningMove(ranked, rng)
      return { move: selected.move }
    }
    return {
      move: this.bot.selectMove({
        board: context.board,
        currentPlayer: context.currentPlayer,
        playerSymbols: context.playerSymbols,
        rules: context.rules,
        boardSize: context.boardSize,
      }),
    }
  }

  async chooseSwap2Phase1(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2Phase1Decision }> {
    return { decision: choosePhase1Heuristic(toMap(context.board), context, rng, 0.06) }
  }

  async chooseSwap2Phase2(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2ColorDecision }> {
    return { decision: choosePhase2Heuristic(toMap(context.board), context, rng, 0.04) }
  }
}

class MctsDatasetBot implements DatasetBot {
  readonly id: DatasetBotId
  readonly params: Record<string, number | string | boolean>
  private readonly options: MctsOptions
  private readonly botOptions: DatasetBotOptions

  constructor(id: DatasetBotId, options: MctsOptions, botOptions: DatasetBotOptions) {
    this.id = id
    this.options = options
    this.botOptions = botOptions
    this.params = { strategy: 'mcts', attitude: mctsAttitude(id), aggressionWeight: 0.6, defensiveBias: 0.75, ...options }
  }

  async selectMove(context: DatasetBotContext, rng: SeededRandom): Promise<BotMoveDecision> {
    const board = toMap(context.board)
    const candidates = generateCandidates(board, context.boardSize, this.options.candidateRadius)
    const temperature = this.temperatureForMove(context.moveNumber)
    if (context.phase === 'swap2_phase0') {
      return neutralSwap2OpeningDecision(this.id, this.params, context, board, candidates, temperature, rng)
    }

    const ownWins = immediateWinningMoves(board, candidates, context.currentPlayer, context)
    if (ownWins.length > 0) {
      return tacticalMctsDecision(this.id, this.params, context, ownWins[0], temperature, 1)
    }

    const opponentWins = immediateWinningMoves(board, candidates, otherPlayer(context.currentPlayer), context)
    if (opponentWins.length > 0) {
      return tacticalMctsDecision(this.id, this.params, context, opponentWins[0], temperature, opponentWins.length === 1 ? 0 : -1)
    }

    const defense = findDefensiveThreatMove(board, candidates, context.currentPlayer, context, rng)
    if (defense) {
      return tacticalMctsDecision(this.id, this.params, context, defense, temperature, 0)
    }

    const search = await runPuctSearch(context, this.puctOptions(), temperature, rng, this.botOptions.modelEvaluator)
    const move = search.selectedMove ?? candidates[0]

    return {
      move,
      mctsTarget: {
        phase: context.phase,
        moveNumber: context.moveNumber,
        actorBotId: this.id,
        actorBotParams: this.params,
        player: context.currentPlayer,
        ...search.metadata,
        selectedMove: move,
        policyTarget: search.policyTarget,
        mctsValue: search.mctsValue,
        temperature,
      },
    }
  }

  async chooseSwap2Phase1(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2Phase1Decision; mctsTarget: MctsTarget }> {
    const temperature = this.temperatureForMove(context.moveNumber)
    const board = toMap(context.board)
    const uncertaintySearch = await runPuctSearch(context, this.puctOptions(), Math.min(temperature, 0.35), rng, this.botOptions.modelEvaluator)
    const decision = choosePhase1MctsDecision(board, context, uncertaintySearch.rootActionStats)
    return {
      decision,
      mctsTarget: {
        phase: 'swap2_phase1',
        actorBotId: this.id,
        actorBotParams: this.params,
        player: context.currentPlayer,
        ...uncertaintySearch.metadata,
        decision,
        policyTarget: uncertaintySearch.policyTarget,
        mctsValue: uncertaintySearch.mctsValue,
        temperature,
      },
    }
  }

  async chooseSwap2Phase2(context: DatasetBotContext, rng: SeededRandom): Promise<{ decision: Swap2ColorDecision; mctsTarget: MctsTarget }> {
    const temperature = this.temperatureForMove(context.moveNumber)
    const search = await runPuctSearch(context, this.puctOptions(), Math.min(temperature, 0.25), rng, this.botOptions.modelEvaluator)
    const decision = search.decision as Swap2ColorDecision | undefined ?? choosePhase2MctsDecision(toMap(context.board), context, rng)
    return {
      decision,
      mctsTarget: {
        phase: 'swap2_phase2',
        actorBotId: this.id,
        actorBotParams: this.params,
        player: context.currentPlayer,
        ...search.metadata,
        decision,
        policyTarget: search.policyTarget,
        mctsValue: search.mctsValue,
        temperature,
      },
    }
  }

  private temperatureForMove(moveNumber: number): number {
    return moveNumber < 24 ? this.options.earlyTemperature : this.options.lateTemperature
  }

  private puctOptions(): {
    simulations: number
    cPuct: number
    candidateRadius: number
    maxCandidates: number
    transpositionTable?: PuctTranspositionTable
    generation: number
  } {
    return {
      simulations: this.options.simulations,
      cPuct: this.options.exploration,
      candidateRadius: this.options.candidateRadius,
      maxCandidates: this.options.maxCandidates ?? 24,
      transpositionTable: this.botOptions.puctTranspositionTable,
      generation: this.botOptions.puctGeneration ?? 1,
    }
  }
}

export interface Swap2UncertaintyOfferOptions {
  valueMargin: number
  uncertainMargin: number
  standardError: number
  minVisits: number
}

export const DEFAULT_SWAP2_UNCERTAINTY_OFFER: Swap2UncertaintyOfferOptions = {
  valueMargin: SWAP2_OFFER_VALUE_MARGIN,
  uncertainMargin: SWAP2_OFFER_UNCERTAIN_MARGIN,
  standardError: SWAP2_OFFER_STANDARD_ERROR,
  minVisits: SWAP2_OFFER_MIN_VISITS,
}

export function shouldOfferSwap2FromUncertainty(
  stats: PuctRootActionStats[],
  options: Swap2UncertaintyOfferOptions = DEFAULT_SWAP2_UNCERTAINTY_OFFER,
): boolean {
  const player1 = stats.find(stat => stat.action.kind === 'special' && stat.action.decision === 'player1')
  const player2 = stats.find(stat => stat.action.kind === 'special' && stat.action.decision === 'player2')
  if (!player1 || !player2) return false
  if (player1.visits < options.minVisits || player2.visits < options.minVisits) return false

  const margin = Math.abs(player1.meanValue - player2.meanValue)
  if (margin <= options.valueMargin) return true

  const uncertainty = Math.max(player1.standardError, player2.standardError)
  return margin <= options.uncertainMargin && uncertainty >= options.standardError
}

function neutralSwap2OpeningDecision(
  botId: string,
  botParams: Record<string, number | string | boolean>,
  context: DatasetBotContext,
  board: Map<string, Cell>,
  candidates: BotMove[],
  temperature: number,
  rng: SeededRandom,
): BotMoveDecision {
  const ranked = rankNeutralOpeningMoves(board, candidates, context)
  const selected = chooseNeutralOpeningMove(ranked, rng)
  return {
    move: selected.move,
    mctsTarget: {
      phase: context.phase,
      moveNumber: context.moveNumber,
      actorBotId: botId,
      actorBotParams: { ...botParams, openingObjective: 'swap2_neutral' },
      player: context.currentPlayer,
      ...targetMetadataFromContext(context),
      selectedMove: selected.move,
      policyTarget: neutralOpeningPolicy(ranked, context.boardSize),
      mctsValue: 0,
      temperature,
    },
  }
}

export function rankNeutralOpeningMovesNoLookahead(
  board: Map<string, Cell>,
  candidates: BotMove[],
  context: DatasetBotContext,
): { move: BotMove; imbalance: number }[] {
  if (board.size === 0) {
    const center = Math.floor(context.boardSize / 2)
    return [{ move: { x: center, y: center }, imbalance: 0 }]
  }

  return candidates
    .map(move => {
      const nextBoard = withMove(board, move, context.currentPlayer)
      const utilities = rawSwap2ColorUtilities(nextBoard, context)
      const imbalance = Math.abs(normalizeSwap2Utility(utilities.player1 - utilities.player2))
      return { move, imbalance }
    })
    .sort((a, b) => a.imbalance - b.imbalance || centerDistance(a.move, context.boardSize) - centerDistance(b.move, context.boardSize) || a.move.y - b.move.y || a.move.x - b.move.x)
}

export function rankNeutralOpeningMoves(
  board: Map<string, Cell>,
  candidates: BotMove[],
  context: DatasetBotContext,
): { move: BotMove; imbalance: number }[] {
  if (board.size === 0) {
    const center = Math.floor(context.boardSize / 2)
    return [{ move: { x: center, y: center }, imbalance: 0 }]
  }

  return candidates
    .map(move => ({
      move,
      imbalance: neutralOpeningImbalance(withMove(board, move, context.currentPlayer), context),
    }))
    .sort((a, b) => a.imbalance - b.imbalance || centerDistance(a.move, context.boardSize) - centerDistance(b.move, context.boardSize) || a.move.y - b.move.y || a.move.x - b.move.x)
}

function neutralOpeningImbalance(board: Map<string, Cell>, context: DatasetBotContext): number {
  if (board.size >= 3) {
    const utilities = rawSwap2ColorUtilities(board, context)
    return Math.abs(normalizeSwap2Utility(utilities.player1 - utilities.player2))
  }

  const nextColor: PlayerColor = board.size === 1 ? 1 : 0
  const nextCandidates = generateCandidates(board, context.boardSize, Math.max(2, context.boardSize > 15 ? 2 : 1))
    .slice(0, SWAP2_NEUTRAL_OPENING_LOOKAHEAD)
  if (nextCandidates.length === 0) return 1
  return Math.min(...nextCandidates.map(move => neutralOpeningImbalance(withMove(board, move, nextColor), context)))
}

export function chooseNeutralOpeningMove(
  ranked: { move: BotMove; imbalance: number }[],
  rng: SeededRandom,
): { move: BotMove; imbalance: number } {
  const pool = ranked.slice(0, Math.min(SWAP2_NEUTRAL_OPENING_POLICY_SIZE, ranked.length))
  const min = Math.min(...pool.map(entry => entry.imbalance))
  const weights = pool.map(entry => Math.exp(-(entry.imbalance - min) / 0.05))
  return pool[weightedIndex(weights, rng)] ?? ranked[0]
}

function neutralOpeningPolicy(ranked: { move: BotMove; imbalance: number }[], boardSize: number): PolicyTarget {
  const pool = ranked.slice(0, Math.min(SWAP2_NEUTRAL_OPENING_POLICY_SIZE, ranked.length))
  const min = Math.min(...pool.map(entry => entry.imbalance))
  const weights = pool.map(entry => Math.exp(-(entry.imbalance - min) / 0.05))
  const total = weights.reduce((sum, weight) => sum + weight, 0)
  return {
    placements: pool.map((entry, index) => ({
      x: entry.move.x,
      y: entry.move.y,
      probability: total > 0 ? weights[index] / total : 1 / Math.max(1, pool.length),
    })).filter(placement => isBoardCoordInBounds(placement.x, placement.y, boardSize)),
  }
}

function centerDistance(move: BotMove, boardSize: number): number {
  const center = (boardSize - 1) / 2
  return Math.abs(move.x - center) + Math.abs(move.y - center)
}

function toMap(board: Record<string, Cell>): Map<string, Cell> {
  return new Map(Object.entries(board))
}

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

export function generateCandidates(board: Map<string, Cell>, boardSize: number, radius: number): BotMove[] {
  if (board.size === 0) {
    const center = Math.floor(boardSize / 2)
    return [{ x: center, y: center }]
  }

  const candidates = new Map<string, BotMove>()
  for (const key of board.keys()) {
    const { x, y } = deserializeCoord(key)
    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, boardSize) || board.has(nKey)) continue
        candidates.set(nKey, { x: nx, y: ny })
      }
    }
  }

  if (candidates.size === 0) {
    for (let y = 0; y < boardSize; y++) {
      for (let x = 0; x < boardSize; x++) {
        const key = serializeCoord(x, y)
        if (!board.has(key)) 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: BotMove, player: PlayerColor): Map<string, Cell> {
  const next = new Map(board)
  next.set(serializeCoord(move.x, move.y), {
    playerId: String(player),
    symbol: PLAYER_SYMBOLS[player],
    timestamp: 0,
  })
  return next
}

function immediateWinningMoves(
  board: Map<string, Cell>,
  candidates: BotMove[],
  player: PlayerColor,
  context: DatasetBotContext,
): BotMove[] {
  return candidates.filter(move => {
    const next = withMove(board, move, player)
    return checkWin(next, move.x, move.y, String(player), context.rules.noOverlines).won
  })
}

function tacticalMctsDecision(
  botId: DatasetBotId,
  botParams: Record<string, number | string | boolean>,
  context: DatasetBotContext,
  move: BotMove,
  temperature: number,
  value: number,
): BotMoveDecision {
  return {
    move,
    mctsTarget: {
      phase: context.phase,
      moveNumber: context.moveNumber,
      actorBotId: botId,
      actorBotParams: botParams,
      player: context.currentPlayer,
      ...targetMetadataFromContext(context),
      selectedMove: move,
      policyTarget: { placements: [{ ...move, probability: 1 }] },
      mctsValue: value,
      temperature,
    },
  }
}

function targetMetadataFromContext(context: DatasetBotContext): Pick<MctsTarget, 'actorSeat' | 'actorColor' | 'stoneColor' | 'playerToMove'> {
  if (context.phase === 'play') {
    return {
      actorSeat: context.currentPlayer,
      actorColor: context.currentPlayer,
      stoneColor: context.currentPlayer,
      playerToMove: context.currentPlayer,
    }
  }
  if (context.phase === 'swap2_phase1') return { actorSeat: 1 }
  if (context.phase === 'swap2_phase2') return { actorSeat: 0 }
  return {
    actorSeat: context.moveNumber <= 3 ? 0 : 1,
    stoneColor: context.currentPlayer,
  }
}

function findDefensiveThreatMove(
  board: Map<string, Cell>,
  candidates: BotMove[],
  player: PlayerColor,
  context: DatasetBotContext,
  rng: SeededRandom,
): BotMove | null {
  const opponent = otherPlayer(player)
  const currentDanger = maxThreatDanger(board, opponent, context)
  if (currentDanger < THREAT_OPEN_THREE) return null

  let best: BotMove | null = null
  let bestDanger = Infinity
  let bestOwnScore = -Infinity
  for (const move of candidates) {
    const next = withMove(board, move, player)
    const danger = maxThreatDanger(next, opponent, context)
    const ownScore = scoreMove(next, move, player, context)
    if (
      danger < bestDanger
      || (danger === bestDanger && ownScore > bestOwnScore)
      || (danger === bestDanger && ownScore === bestOwnScore && rng.next() < 0.25)
    ) {
      best = move
      bestDanger = danger
      bestOwnScore = ownScore
    }
  }

  return best && bestDanger < currentDanger ? best : null
}

export function maxThreatDanger(board: Map<string, Cell>, player: PlayerColor, context: DatasetBotContext): number {
  let danger = 0
  for (const move of generateCandidates(board, context.boardSize, 2)) {
    danger = Math.max(danger, threatDangerForMove(board, move, player, context))
    if (danger >= THREAT_WIN) return danger
  }
  return danger
}

function threatDangerForMove(
  board: Map<string, Cell>,
  move: BotMove,
  player: PlayerColor,
  context: DatasetBotContext,
): number {
  const next = withMove(board, move, player)
  if (checkWin(next, move.x, move.y, String(player), context.rules.noOverlines).won) return THREAT_WIN

  let danger = 0
  for (const [dx, dy] of [[1, 0], [0, 1], [1, 1], [1, -1]] as const) {
    const line = threatLineStats(board, move, player, dx, dy, context.boardSize)
    if (line.count === 4 && line.openEnds === 2) danger = Math.max(danger, THREAT_OPEN_FOUR)
    else if (line.count === 4 && line.openEnds === 1) danger = Math.max(danger, THREAT_CLOSED_FOUR)
    else if (isBrokenFourThreat(board, move, player, dx, dy, context.boardSize)) danger = Math.max(danger, THREAT_BROKEN_FOUR)
    else if (line.count === 3 && line.openEnds === 2) danger = Math.max(danger, THREAT_OPEN_THREE)
  }
  return danger
}

function threatLineStats(
  board: Map<string, Cell>,
  move: BotMove,
  player: PlayerColor,
  dx: number,
  dy: number,
  boardSize: number,
): { count: number; openEnds: number } {
  const playerId = String(player)
  let count = 1
  let openEnds = 0
  let cx = move.x + dx
  let cy = move.y + dy
  while (board.get(serializeCoord(cx, cy))?.playerId === playerId) {
    count++
    cx += dx
    cy += dy
  }
  if (isBoardCoordInBounds(cx, cy, boardSize) && !board.has(serializeCoord(cx, cy))) openEnds++
  cx = move.x - dx
  cy = move.y - dy
  while (board.get(serializeCoord(cx, cy))?.playerId === playerId) {
    count++
    cx -= dx
    cy -= dy
  }
  if (isBoardCoordInBounds(cx, cy, boardSize) && !board.has(serializeCoord(cx, cy))) openEnds++
  return { count, openEnds }
}

function isBrokenFourThreat(
  board: Map<string, Cell>,
  move: BotMove,
  player: PlayerColor,
  dx: number,
  dy: number,
  boardSize: number,
): boolean {
  const playerId = String(player)
  for (let start = -4; start <= 0; start++) {
    let own = 0
    let empty = 0
    let blocked = false
    for (let i = 0; i < 5; i++) {
      const x = move.x + dx * (start + i)
      const y = move.y + dy * (start + i)
      if (!isBoardCoordInBounds(x, y, boardSize)) {
        blocked = true
        break
      }
      const isMove = x === move.x && y === move.y
      const cell = isMove ? { playerId } as Cell : board.get(serializeCoord(x, y))
      if (!cell) empty++
      else if (cell.playerId === playerId) own++
      else blocked = true
    }
    if (!blocked && own === 4 && empty === 1) return true
  }
  return false
}

function scoreMove(board: Map<string, Cell>, move: BotMove, player: PlayerColor, context: DatasetBotContext): number {
  const next = withMove(board, move, player)
  if (checkWin(next, move.x, move.y, String(player), context.rules.noOverlines).won) return 1_000_000
  if (checkWin(withMove(board, move, otherPlayer(player)), move.x, move.y, String(otherPlayer(player)), context.rules.noOverlines).won) return 900_000
  return evaluateBoard(next, player, context)
}

function evaluateBoard(board: Map<string, Cell>, player: PlayerColor, context: DatasetBotContext): number {
  const opponent = otherPlayer(player)
  return linePotential(board, player, context) - linePotential(board, opponent, context) * 1.05
}

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

  for (const [key, cell] of board) {
    if (cell.playerId !== playerId) continue
    const { x, y } = deserializeCoord(key)
    score += Math.max(0, context.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, context.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, context.boardSize) && !board.has(serializeCoord(cx, cy))) openEnds++
      if (context.rules.noOverlines && count > 5) continue
      score += Math.pow(10, Math.min(count, 5)) * (openEnds + 1)
    }
  }

  return score
}

function phase1Scores(board: Map<string, Cell>, context: DatasetBotContext): Record<Swap2Phase1Decision, number> {
  const { player1, player2 } = rawSwap2ColorUtilities(board, context)
  return positiveSpecialScores({
    player1,
    player2,
    offer: Math.min(player1, player2) + Math.abs(player1 - player2) * 0.35,
  })
}

function phase2Scores(board: Map<string, Cell>, context: DatasetBotContext): Record<Swap2ColorDecision, number> {
  const { player1, player2 } = rawSwap2ColorUtilities(board, context)
  return positiveSpecialScores({
    player1,
    player2,
  })
}

function choosePhase1MctsDecision(
  board: Map<string, Cell>,
  context: DatasetBotContext,
  rootStats: PuctRootActionStats[],
): Swap2Phase1Decision {
  const utilities = rawSwap2ColorUtilities(board, context)
  const delta = utilities.player1 - utilities.player2
  const margin = Math.abs(normalizeSwap2Utility(delta))
  if (margin <= SWAP2_OFFER_VALUE_MARGIN) return 'offer'

  const searchDecision = choosePhase1FromRootStats(rootStats)
  if (searchDecision) return searchDecision

  if (margin <= SWAP2_OFFER_UNCERTAIN_MARGIN && shouldOfferSwap2FromUncertainty(rootStats)) return 'offer'
  return delta > 0 ? 'player1' : 'player2'
}

function choosePhase1FromRootStats(rootStats: PuctRootActionStats[]): Swap2Phase1Decision | null {
  const stats = rootStats.filter((stat): stat is PuctRootActionStats & { action: { kind: 'special'; decision: Swap2Phase1Decision } } => (
    stat.action.kind === 'special'
    && (stat.action.decision === 'offer' || stat.action.decision === 'player1' || stat.action.decision === 'player2')
  ))
  const totalVisits = stats.reduce((sum, stat) => sum + stat.visits, 0)
  if (stats.length < 3 || totalVisits <= 0) return null
  const best = stats.reduce((current, stat) => stat.visits > current.visits ? stat : current, stats[0])
  return best.action.decision
}

function choosePhase2MctsDecision(
  board: Map<string, Cell>,
  context: DatasetBotContext,
  rng: SeededRandom,
): Swap2ColorDecision {
  const utilities = rawSwap2ColorUtilities(board, context)
  if (utilities.player1 === utilities.player2) return rng.pick(['player1', 'player2'] as const)
  return utilities.player1 > utilities.player2 ? 'player1' : 'player2'
}

function rawSwap2ColorUtilities(
  board: Map<string, Cell>,
  context: DatasetBotContext,
): Record<Swap2ColorDecision, number> {
  return {
    player1: swap2DecisionUtility(board, context, 0),
    player2: swap2DecisionUtility(board, context, 1),
  }
}

function normalizeSwap2Utility(value: number): number {
  return Math.tanh(value / SWAP2_UTILITY_SCALE)
}

function swap2DecisionUtility(board: Map<string, Cell>, context: DatasetBotContext, actorColor: PlayerColor): number {
  const opponent = otherPlayer(actorColor)
  const nextMoveGain = bestImmediateMoveGain(board, 1, context)
  const tempoScore = actorColor === 1 ? nextMoveGain * 2.5 : -nextMoveGain * 2.5
  const threatScore = maxThreatDanger(board, actorColor, context) - maxThreatDanger(board, opponent, context) * 1.1
  const mobilityScore = mobilityPotential(board, actorColor, context) - mobilityPotential(board, opponent, context) * 1.05
  return (
    linePotential(board, actorColor, context)
    - linePotential(board, opponent, context) * 1.05
    + tempoScore
    + threatScore * SWAP2_THREAT_WEIGHT
    + mobilityScore * SWAP2_MOBILITY_WEIGHT
  )
}

function bestImmediateMoveGain(board: Map<string, Cell>, color: PlayerColor, context: DatasetBotContext): number {
  let best = 0
  for (const move of generateCandidates(board, context.boardSize, 2).slice(0, 24)) {
    const next = withMove(board, move, color)
    if (checkWin(next, move.x, move.y, String(color), context.rules.noOverlines).won) return 1_000_000
    best = Math.max(best, linePotential(next, color, context) - linePotential(board, color, context))
  }
  return best
}

function mobilityPotential(board: Map<string, Cell>, color: PlayerColor, context: DatasetBotContext): number {
  let mobility = 0
  for (const move of generateCandidates(board, context.boardSize, 2).slice(0, 24)) {
    const next = withMove(board, move, color)
    const gain = linePotential(next, color, context) - linePotential(board, color, context)
    if (gain >= THREAT_OPEN_THREE) 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
}

export function choosePhase1Heuristic(
  board: Map<string, Cell>,
  context: DatasetBotContext,
  rng: SeededRandom,
  noise: number,
): Swap2Phase1Decision {
  const scores = phase1Scores(board, context)
  const policy = scoreSpecialPolicy(scores, 0.35 + noise)
  const weights = [
    policy.specialActions?.offerTwo ?? 0,
    policy.specialActions?.choosePlayer1 ?? 0,
    policy.specialActions?.choosePlayer2 ?? 0,
  ]
  return (['offer', 'player1', 'player2'] as const)[weightedIndex(weights, rng)]
}

export function choosePhase2Heuristic(
  board: Map<string, Cell>,
  context: DatasetBotContext,
  rng: SeededRandom,
  noise: number,
): Swap2ColorDecision {
  const scores = phase2Scores(board, context)
  const policy = scoreSpecialPolicy(scores, 0.2 + noise)
  const weights = [policy.specialActions?.choosePlayer1 ?? 0, policy.specialActions?.choosePlayer2 ?? 0]
  return (['player1', 'player2'] as const)[weightedIndex(weights, rng)]
}

function chooseTopSpecial<T extends Swap2Phase1Decision | Swap2ColorDecision>(
  scores: Partial<Record<T, number>>,
  rng: SeededRandom,
): T {
  let bestScore = -Infinity
  let best: T[] = []
  for (const [decision, score] of Object.entries(scores) as [T, number][]) {
    if (score > bestScore) {
      bestScore = score
      best = [decision]
    } else if (score === bestScore) {
      best.push(decision)
    }
  }
  return rng.pick(best)
}

function specialDecisionValue<T extends Swap2Phase1Decision | Swap2ColorDecision>(
  scores: Partial<Record<T, number>>,
  decision: T,
): number {
  const values = Object.values(scores).filter((value): value is number => typeof value === 'number')
  const min = Math.min(...values)
  const max = Math.max(...values)
  if (max === min) return 0
  return ((scores[decision] ?? min) - min) / (max - min) * 2 - 1
}

function scoreSpecialPolicy(
  scores: Partial<Record<Swap2Phase1Decision, number>>,
  temperature: number,
): PolicyTarget {
  const entries: [Swap2Phase1Decision, number][] = []
  if (scores.player1 !== undefined) entries.push(['player1', scores.player1])
  if (scores.player2 !== undefined) entries.push(['player2', scores.player2])
  if (scores.offer !== undefined) entries.push(['offer', scores.offer])

  const max = Math.max(...entries.map(([, score]) => score))
  const temp = Math.max(temperature, 0.01)
  const weights = entries.map(([, score]) => Math.exp((score - max) / Math.max(1, Math.abs(max), 1000) / temp))
  const total = weights.reduce((sum, weight) => sum + weight, 0)
  const specialActions: NonNullable<PolicyTarget['specialActions']> = {}
  for (let i = 0; i < entries.length; i++) {
    const probability = weights[i] / total
    if (entries[i][0] === 'player1') specialActions.choosePlayer1 = probability
    if (entries[i][0] === 'player2') specialActions.choosePlayer2 = probability
    if (entries[i][0] === 'offer') specialActions.offerTwo = probability
  }
  return { placements: [], specialActions }
}

export function boardSizeFromRules(rules: GameRules): number {
  return rules.boardSize ?? BOARD_SIZE
}