package de.nowchess.bot.bots.nnue

import de.nowchess.api.board.{Board, Color, File, PieceType, Rank, Square}
import de.nowchess.api.game.GameContext

class NNUE:

  private val l1Weights = NNUEWeights.l1_weights
  private val l1Bias = NNUEWeights.l1_bias
  private val l2Weights = NNUEWeights.l2_weights
  private val l2Bias = NNUEWeights.l2_bias
  private val l3Weights = NNUEWeights.l3_weights
  private val l3Bias = NNUEWeights.l3_bias

  // Pre-allocated buffers for inference
  private val features = new Array[Float](768)
  private val l1Output = new Array[Float](256)
  private val l2Output = new Array[Float](32)

  /** Convert a position to 768-dimensional binary feature vector.
   *  12 piece types (white pawn to black king) × 64 squares from white's perspective. */
  def positionToFeatures(board: Board, sideToMove: Color): Array[Float] =
    // Zero out features array
    java.util.Arrays.fill(features, 0f)

    // Piece type to feature index offset: wp=0, wn=64, wb=128, wr=192, wq=256, wk=320, bp=384, bn=448, bb=512, br=576, bq=640, bk=704
    val pieceToFeatureOffset = Array(
      0,    // White Pawn (0)
      64,   // White Knight (1)
      128,  // White Bishop (2)
      192,  // White Rook (3)
      256,  // White Queen (4)
      320,  // White King (5)
      384,  // Black Pawn (6)
      448,  // Black Knight (7)
      512,  // Black Bishop (8)
      576,  // Black Rook (9)
      640,  // Black Queen (10)
      704   // Black King (11)
    )

    // Build features: always from white's perspective
    for
      fileIdx <- 0 until 8
      rankIdx <- 0 until 8
    do
      val file = File.values(fileIdx)
      val rank = Rank.values(rankIdx)
      val square = Square(file, rank)
      val squareNum = rankIdx * 8 + fileIdx

      board.pieceAt(square).foreach { piece =>
        val featureIdx = if sideToMove == Color.Black then
          // Mirror square for black side-to-move
          val mirroredSq = squareNum ^ 56
          val offset = pieceToFeatureOffset(piece.color.ordinal * 6 + piece.pieceType.ordinal)
          offset + mirroredSq
        else
          val offset = pieceToFeatureOffset(piece.color.ordinal * 6 + piece.pieceType.ordinal)
          offset + squareNum

        if featureIdx >= 0 && featureIdx < 768 then
          features(featureIdx) = 1f
      }

    features

  /** Run NNUE inference on the given position.
   *  Returns centipawn score from the perspective of the side-to-move.
   *  No allocations in the hot path (uses pre-allocated buffers). */
  def evaluate(context: GameContext): Int =
    val features = positionToFeatures(context.board, context.turn)

    // Layer 1: Dense(768 -> 256) + ReLU
    for i <- 0 until 256 do
      var sum = l1Bias(i)
      for j <- 0 until 768 do
        sum += features(j) * l1Weights(i * 768 + j)
      l1Output(i) = if sum > 0f then sum else 0f

    // Layer 2: Dense(256 -> 32) + ReLU
    for i <- 0 until 32 do
      var sum = l2Bias(i)
      for j <- 0 until 256 do
        sum += l1Output(j) * l2Weights(i * 256 + j)
      l2Output(i) = if sum > 0f then sum else 0f

    // Layer 3: Dense(32 -> 1), no activation
    var output = l3Bias(0)
    for j <- 0 until 32 do
      output += l2Output(j) * l3Weights(j)

    // Convert from sigmoid(output) back to centipawns (output is trained as sigmoid(eval/400))
    // Inverse sigmoid: eval/400 = ln(output / (1 - output))
    // But for simplicity, just scale directly: output ≈ sigmoid(eval/400), so eval ≈ 400 * (output - 0.5) * 2
    val cp = (output * 400f).toInt
    math.max(-20000, math.min(20000, cp))
