shared/bot/FeatureEncoder.ts

import { deserializeCoord, isBoardCoordInBounds, serializeCoord } from '../board.ts'
import type { Cell, GameRules } from '../types.ts'
import { ZobristHasher } from './Zobrist.ts'

export const FEATURE_CHANNELS = 13
export const PLAY_ONLY_FEATURE_CHANNELS = 7
export const CH_CURRENT_STONES = 0
export const CH_OPPONENT_STONES = 1
export const CH_LAST_MOVE = 2
export const CH_PHASE0 = 3
export const CH_PHASE1 = 4
export const CH_PHASE2 = 5
export const CH_NORMAL = 6
export const CH_CURRENT_PLAYER = 7
export const CH_OPENING_STONES = 8
export const CH_OFFERED_STONES = 9
export const CH_CURRENT_THREATS = 10
export const CH_OPPONENT_THREATS = 11
export const CH_OVERLINE_TRAPS = 12

const PLAY_ONLY_CH_CURRENT_STONES = 0
const PLAY_ONLY_CH_OPPONENT_STONES = 1
const PLAY_ONLY_CH_LAST_MOVE = 2
const PLAY_ONLY_CH_CURRENT_PLAYER = 3
const PLAY_ONLY_CH_CURRENT_THREATS = 4
const PLAY_ONLY_CH_OPPONENT_THREATS = 5
const PLAY_ONLY_CH_OVERLINE_TRAPS = 6

const DIRECTIONS: readonly [number, number][] = [[1, 0], [0, 1], [1, 1], [1, -1]]
export const OPEN_THREE_THREAT_LEVEL = 0.65

export type FeaturePhase = 'swap2_phase0' | 'swap2_phase1' | 'swap2_phase2' | 'play' | 'offer_place_1' | 'offer_place_2'

export interface FeatureCoord {
  x: number
  y: number
}

export interface FeatureEncoderInput {
  board: Map<string, Cell> | Record<string, Cell>
  boardSize: number
  rules: GameRules
  currentPlayer: 0 | 1
  phase: FeaturePhase
  lastMove?: FeatureCoord
  openingStones?: FeatureCoord[]
  offeredStones?: FeatureCoord[]
}

export interface ThreatMaps {
  current: Float32Array
  opponent: Float32Array
  overline: Float32Array
}

export class ThreatMapState {
  readonly boardSize: number
  readonly hash: bigint
  readonly maps: ThreatMaps

  private constructor(boardSize: number, hash: bigint, maps: ThreatMaps) {
    this.boardSize = boardSize
    this.hash = hash
    this.maps = maps
  }

  static fromBoard(input: FeatureEncoderInput, hasher = new ZobristHasher(input.boardSize)): ThreatMapState {
    const board = toMap(input.board)
    return new ThreatMapState(
      input.boardSize,
      hasher.hash(board, input.currentPlayer),
      computeThreatMaps(board, input.boardSize, input.rules, input.currentPlayer),
    )
  }

  applyMove(input: FeatureEncoderInput, move: FeatureCoord, player: 0 | 1, hasher = new ZobristHasher(input.boardSize)): ThreatMapState {
    const board = toMap(input.board)
    const boardArray = boardToArray(board, input.boardSize)
    const nextMaps = {
      current: new Float32Array(this.maps.current),
      opponent: new Float32Array(this.maps.opponent),
      overline: new Float32Array(this.maps.overline),
    }
    const affected = affectedCells(move, input.boardSize)
    for (const cell of affected) {
      updateThreatCell(nextMaps, boardArray, cell.x, cell.y, input.boardSize, input.rules, input.currentPlayer)
    }
    const hash = hasher.toggleStone(this.hash, move.x, move.y, player)
    return new ThreatMapState(input.boardSize, hash, nextMaps)
  }
}

export function encodeFeaturePlanes(input: FeatureEncoderInput, threatMaps?: ThreatMaps): Float32Array {
  const board = toMap(input.board)
  const size = input.boardSize
  const planes = new Float32Array(FEATURE_CHANNELS * size * size)
  const opponent = otherPlayer(input.currentPlayer)
  const maps = threatMaps ?? computeThreatMaps(board, size, input.rules, input.currentPlayer)

  for (const [key, cell] of board) {
    const { x, y } = deserializeCoord(key)
    if (!isBoardCoordInBounds(x, y, size)) continue
    const player = Number(cell.playerId)
    if (player === input.currentPlayer) setPlane(planes, size, CH_CURRENT_STONES, x, y, 1)
    if (player === opponent) setPlane(planes, size, CH_OPPONENT_STONES, x, y, 1)
  }

  if (input.lastMove) setPlane(planes, size, CH_LAST_MOVE, input.lastMove.x, input.lastMove.y, 1)
  fillPhasePlane(planes, size, input.phase)
  fillPlane(planes, size, CH_CURRENT_PLAYER, input.currentPlayer)
  for (const stone of input.openingStones ?? []) setPlane(planes, size, CH_OPENING_STONES, stone.x, stone.y, 1)
  for (const stone of input.offeredStones ?? []) setPlane(planes, size, CH_OFFERED_STONES, stone.x, stone.y, 1)

  copyMap(planes, size, CH_CURRENT_THREATS, maps.current)
  copyMap(planes, size, CH_OPPONENT_THREATS, maps.opponent)
  copyMap(planes, size, CH_OVERLINE_TRAPS, maps.overline)
  return planes
}

export function encodePlayOnlyFeaturePlanes(input: FeatureEncoderInput, threatMaps?: ThreatMaps): Float32Array {
  const board = toMap(input.board)
  const size = input.boardSize
  const planes = new Float32Array(PLAY_ONLY_FEATURE_CHANNELS * size * size)
  const opponent = otherPlayer(input.currentPlayer)
  const maps = threatMaps ?? computeThreatMaps(board, size, input.rules, input.currentPlayer)

  for (const [key, cell] of board) {
    const { x, y } = deserializeCoord(key)
    if (!isBoardCoordInBounds(x, y, size)) continue
    const player = Number(cell.playerId)
    if (player === input.currentPlayer) setPlane(planes, size, PLAY_ONLY_CH_CURRENT_STONES, x, y, 1)
    if (player === opponent) setPlane(planes, size, PLAY_ONLY_CH_OPPONENT_STONES, x, y, 1)
  }

  if (input.lastMove) setPlane(planes, size, PLAY_ONLY_CH_LAST_MOVE, input.lastMove.x, input.lastMove.y, 1)
  fillPlane(planes, size, PLAY_ONLY_CH_CURRENT_PLAYER, input.currentPlayer)
  copyMap(planes, size, PLAY_ONLY_CH_CURRENT_THREATS, maps.current)
  copyMap(planes, size, PLAY_ONLY_CH_OPPONENT_THREATS, maps.opponent)
  copyMap(planes, size, PLAY_ONLY_CH_OVERLINE_TRAPS, maps.overline)
  return planes
}

export function computeThreatMaps(
  board: Map<string, Cell>,
  boardSize: number,
  rules: GameRules,
  currentPlayer: 0 | 1,
): ThreatMaps {
  const boardArray = boardToArray(board, boardSize)
  const current = new Float32Array(boardSize * boardSize)
  const opponent = new Float32Array(boardSize * boardSize)
  const overline = new Float32Array(boardSize * boardSize)
  for (let y = 0; y < boardSize; y++) {
    for (let x = 0; x < boardSize; x++) {
      updateThreatCell({ current, opponent, overline }, boardArray, x, y, boardSize, rules, currentPlayer)
    }
  }
  return { current, opponent, overline }
}

function updateThreatCell(
  maps: ThreatMaps,
  board: Uint8Array,
  x: number,
  y: number,
  boardSize: number,
  rules: GameRules,
  currentPlayer: 0 | 1,
): void {
  const index = y * boardSize + x
  if (!isBoardCoordInBounds(x, y, boardSize) || board[index] !== 0) {
    maps.current[index] = 0
    maps.opponent[index] = 0
    maps.overline[index] = 0
    return
  }
  const currentOverline = rules.noOverlines && createsOverline(board, x, y, currentPlayer, boardSize)
  maps.current[index] = currentOverline ? 0 : threatScore(board, x, y, currentPlayer, boardSize, rules)
  maps.opponent[index] = threatScore(board, x, y, otherPlayer(currentPlayer), boardSize, rules)
  maps.overline[index] = currentOverline ? 1 : 0
}

function threatScore(board: Uint8Array, x: number, y: number, player: 0 | 1, boardSize: number, rules: GameRules): number {
  let best = 0
  for (const [dx, dy] of DIRECTIONS) {
    const line = scanLine(board, x, y, player, dx, dy, boardSize)
    if (line.count >= 5 && !(rules.noOverlines && line.count > 5)) best = Math.max(best, 1)
    else if (line.count === 4 && line.openEnds >= 1) best = Math.max(best, 0.85)
    else if (brokenFour(board, x, y, player, dx, dy, boardSize)) best = Math.max(best, 0.75)
    else if (line.count === 3 && line.openEnds === 2) best = Math.max(best, OPEN_THREE_THREAT_LEVEL)
    else if (line.count === 3 && line.openEnds === 1) best = Math.max(best, 0.25)
  }
  return best
}

function brokenFour(board: Uint8Array, x: number, y: number, player: 0 | 1, dx: number, dy: number, boardSize: number): boolean {
  const playerCell = player + 1
  for (let start = -4; start <= 0; start++) {
    let own = 0
    let empty = 0
    let blocked = false
    for (let i = 0; i < 5; i++) {
      const cx = x + dx * (start + i)
      const cy = y + dy * (start + i)
      if (!isBoardCoordInBounds(cx, cy, boardSize)) {
        blocked = true
        break
      }
      const isPlacedCell = cx === x && cy === y
      const cell = isPlacedCell ? playerCell : board[cy * boardSize + cx]
      if (cell === 0) empty++
      else if (cell === playerCell) own++
      else blocked = true
    }
    if (!blocked && own === 4 && empty === 1) return true
  }
  return false
}

function createsOverline(board: Uint8Array, x: number, y: number, player: 0 | 1, boardSize: number): boolean {
  return DIRECTIONS.some(([dx, dy]) => scanLine(board, x, y, player, dx, dy, boardSize).count > 5)
}

function scanLine(
  board: Uint8Array,
  x: number,
  y: number,
  player: 0 | 1,
  dx: number,
  dy: number,
  boardSize: number,
): { count: number; openEnds: number } {
  const playerCell = player + 1
  let count = 1
  let openEnds = 0
  let cx = x + dx
  let cy = y + dy
  while (isBoardCoordInBounds(cx, cy, boardSize) && board[cy * boardSize + cx] === playerCell) {
    count++
    cx += dx
    cy += dy
  }
  if (isBoardCoordInBounds(cx, cy, boardSize) && board[cy * boardSize + cx] === 0) openEnds++
  cx = x - dx
  cy = y - dy
  while (isBoardCoordInBounds(cx, cy, boardSize) && board[cy * boardSize + cx] === playerCell) {
    count++
    cx -= dx
    cy -= dy
  }
  if (isBoardCoordInBounds(cx, cy, boardSize) && board[cy * boardSize + cx] === 0) openEnds++
  return { count, openEnds }
}

function boardToArray(board: Map<string, Cell>, boardSize: number): Uint8Array {
  const array = new Uint8Array(boardSize * boardSize)
  for (const [key, cell] of board) {
    const { x, y } = deserializeCoord(key)
    if (!isBoardCoordInBounds(x, y, boardSize)) continue
    const player = Number(cell.playerId)
    if (player === 0 || player === 1) array[y * boardSize + x] = player + 1
  }
  return array
}

function affectedCells(move: FeatureCoord, boardSize: number): FeatureCoord[] {
  const cells = new Map<string, FeatureCoord>()
  for (const [dx, dy] of DIRECTIONS) {
    for (let offset = -5; offset <= 5; offset++) {
      const x = move.x + dx * offset
      const y = move.y + dy * offset
      if (!isBoardCoordInBounds(x, y, boardSize)) continue
      cells.set(serializeCoord(x, y), { x, y })
    }
  }
  cells.set(serializeCoord(move.x, move.y), move)
  return [...cells.values()]
}

function fillPhasePlane(planes: Float32Array, boardSize: number, phase: FeaturePhase): void {
  if (phase === 'swap2_phase0') fillPlane(planes, boardSize, CH_PHASE0, 1)
  else if (phase === 'swap2_phase1') fillPlane(planes, boardSize, CH_PHASE1, 1)
  else if (phase === 'swap2_phase2') fillPlane(planes, boardSize, CH_PHASE2, 1)
  else fillPlane(planes, boardSize, CH_NORMAL, 1)
}

function fillPlane(planes: Float32Array, boardSize: number, channel: number, value: number): void {
  const offset = channel * boardSize * boardSize
  planes.fill(value, offset, offset + boardSize * boardSize)
}

function copyMap(planes: Float32Array, boardSize: number, channel: number, map: Float32Array): void {
  planes.set(map, channel * boardSize * boardSize)
}

function setPlane(planes: Float32Array, boardSize: number, channel: number, x: number, y: number, value: number): void {
  if (!isBoardCoordInBounds(x, y, boardSize)) return
  planes[channel * boardSize * boardSize + y * boardSize + x] = value
}

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

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