package de.nowchess.bot.bots.nnue

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

class NNUE:

  private val (l1Weights, l1Bias, l2Weights, l2Bias, l3Weights, l3Bias) = loadWeights()

  private def loadWeights(): (Array[Float], Array[Float], Array[Float], Array[Float], Array[Float], Array[Float]) =
    val stream = getClass.getResourceAsStream("/nnue_weights.bin")
    if stream == null then
      throw RuntimeException("NNUE weights file not found in resources")

    try
      val bytes = stream.readAllBytes()
      val buffer = ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN)

      // Read and verify magic number
      val magic = buffer.getInt()
      if magic != 0x4555_4e4e then // "NNUE" in little-endian
        throw RuntimeException(s"Invalid magic number: 0x${magic.toHexString}")

      // Read version
      val version = buffer.getInt()
      if version != 1 then
        throw RuntimeException(s"Unsupported weight version: $version")

      // Read all weight tensors in order
      val l1w = readTensor(buffer)
      val l1b = readTensor(buffer)
      val l2w = readTensor(buffer)
      val l2b = readTensor(buffer)
      val l3w = readTensor(buffer)
      val l3b = readTensor(buffer)

      (l1w, l1b, l2w, l2b, l3w, l3b)
    finally stream.close()

  private def readTensor(buffer: ByteBuffer): Array[Float] =
    // Read shape
    val shapeLen = buffer.getInt()
    val shape = Array.ofDim[Int](shapeLen)
    for i <- 0 until shapeLen do
      shape(i) = buffer.getInt()

    // Calculate total elements
    val totalElements = shape.product

    // Read float data
    val floats = Array.ofDim[Float](totalElements)
    for i <- 0 until totalElements do
      floats(i) = buffer.getFloat()
    floats

  // 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. */
  private 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))
