import {
  PLAY_ONLY_FEATURE_CHANNELS,
  encodeFeaturePlanes,
  encodePlayOnlyFeaturePlanes,
  type FeaturePhase,
} from './FeatureEncoder.ts'
import type { MaybePromise, PolicyValueEvaluator, SearchAction, SearchState } from './PuctSearch.ts'
import { valueSignForRootActor } from './PuctSearch.ts'

export interface ModelEvaluation {
  policyLogits: readonly number[]
  value: number
}

export interface ModelInferenceBackend {
  featureChannels?: number
  evaluate(features: Float32Array, state: SearchState): MaybePromise<ModelEvaluation>
}

export class EncodedPolicyValueEvaluator implements PolicyValueEvaluator {
  private readonly backend: ModelInferenceBackend

  constructor(backend: ModelInferenceBackend) {
    this.backend = backend
  }

  async evaluate(state: SearchState, legalActions: SearchAction[]): Promise<{ value: number; priors: number[] }> {
    const encoderInput = {
      board: state.board,
      boardSize: state.boardSize,
      rules: state.rules,
      currentPlayer: state.playerToMove ?? state.actorSeat,
      phase: state.phase as FeaturePhase,
    }
    const features = this.backend.featureChannels === PLAY_ONLY_FEATURE_CHANNELS
      ? encodePlayOnlyFeaturePlanes(encoderInput)
      : encodeFeaturePlanes(encoderInput)
    const output = await this.backend.evaluate(features, state)
    const valueActor = state.playerToMove ?? state.actorSeat
    return {
      value: clamp(output.value, -1, 1) * valueSignForRootActor(state, valueActor),
      priors: softmax(legalActions.map(action => output.policyLogits[actionIndex(action, state.boardSize)] ?? -1_000_000)),
    }
  }
}

function actionIndex(action: SearchAction, boardSize: number): number {
  const boardArea = boardSize * boardSize
  if (action.kind === 'place') return action.y * boardSize + action.x
  if (action.decision === 'player1') return boardArea
  if (action.decision === 'player2') return boardArea + 1
  return boardArea + 2
}

function softmax(logits: number[]): number[] {
  if (logits.length === 0) return []
  const max = Math.max(...logits)
  const exps = logits.map(value => Math.exp(value - max))
  const total = exps.reduce((sum, value) => sum + value, 0)
  if (!Number.isFinite(total) || total <= 0) return logits.map(() => 1 / logits.length)
  return exps.map(value => value / total)
}

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