import nimpy

import sequtils, strutils
let chess = pyImport("chess")
const
  pieceValue = {
    "PAWN": 100,
    "ROOK": 500,
    "KNIGHT": 320,
    "BISHOP": 330,
    "QUEEN": 900,
    "KING": 20000
  }

var
    pawnEvalWhite = @[
      0, 0, 0, 0, 0, 0, 0, 0,
      5, 10, 10, -20, -20, 10, 10, 5,
      5, -5, -10, 0, 0, -10, -5, 5,
      0, 0, 0, 20, 20, 0, 0, 0,
      5, 5, 10, 25, 25, 10, 5, 5,
      10, 10, 20, 30, 30, 20, 10, 10,
      50, 50, 50, 50, 50, 50, 50, 50,
      0, 0, 0, 0, 0, 0, 0, 0]

    pawnEvalBlack: array[64, int] = pawnEvalWhite.reversed()

    knightEval: array[64, int] = [
      -50, -40, -30, -30, -30, -30, -40, -50,
      -40, -20, 0, 0, 0, 0, -20, -40,
      -30, 0, 10, 15, 15, 10, 0, -30,
      -30, 5, 15, 20, 20, 15, 5, -30,
      -30, 0, 15, 20, 20, 15, 0, -30,
      -30, 5, 10, 15, 15, 10, 5, -30,
      -40, -20, 0, 5, 5, 0, -20, -40,
      -50, -40, -30, -30, -30, -30, -40, -50]

    bishopEvalWhite: array[64, int] = [
      -20, -10, -10, -10, -10, -10, -10, -20,
      -10, 5, 0, 0, 0, 0, 5, -10,
      -10, 10, 10, 10, 10, 10, 10, -10,
      -10, 0, 10, 10, 10, 10, 0, -10,
      -10, 5, 5, 10, 10, 5, 5, -10,
      -10, 0, 5, 10, 10, 5, 0, -10,
      -10, 0, 0, 0, 0, 0, 0, -10,
      -20, -10, -10, -10, -10, -10, -10, -20]

    bishopEvalBlack: array[64, int] = bishopEvalWhite.reverse()

    rookEvalWhite: array[64, int] = [
      0, 0, 0, 5, 5, 0, 0, 0,
      -5, 0, 0, 0, 0, 0, 0, -5,
      -5, 0, 0, 0, 0, 0, 0, -5,
      -5, 0, 0, 0, 0, 0, 0, -5,
      -5, 0, 0, 0, 0, 0, 0, -5,
      -5, 0, 0, 0, 0, 0, 0, -5,
      5, 10, 10, 10, 10, 10, 10, 5,
      0, 0, 0, 0, 0, 0, 0, 0]

    rookEvalBlack: array[64, int] = rookEvalWhite.reverse()

    queenEval: array[64, int] = [
      -20, -10, -10, -5, -5, -10, -10, -20,
      -10, 0, 0, 0, 0, 0, 0, -10,
      -10, 0, 5, 5, 5, 5, 0, -10,
      -5, 0, 5, 5, 5, 5, 0, -5,
      0, 0, 5, 5, 5, 5, 0, -5,
      -10, 5, 5, 5, 5, 5, 0, -10,
      -10, 0, 5, 0, 0, 0, 0, -10,
      -20, -10, -10, -5, -5, -10, -10, -20]

    kingEvalWhite: array[64, int] = [
      20, 30, 10, 0, 0, 10, 30, 20,
      20, 20, 0, 0, 0, 0, 20, 20,
      -10, -20, -20, -20, -20, -20, -20, -10,
      20, -30, -30, -40, -40, -30, -30, -20,
      -30, -40, -40, -50, -50, -40, -40, -30,
      -30, -40, -40, -50, -50, -40, -40, -30,
      -30, -40, -40, -50, -50, -40, -40, -30,
      -30, -40, -40, -50, -50, -40, -40, -30]

    kingEvalBlack: array[64, int] = kingEvalWhite.reverse()

    kingEvalEndGameWhite: array[64, int] = [
      50, -30, -30, -30, -30, -30, -30, -50,
      -30, -30, 0, 0, 0, 0, -30, -30,
      -30, -10, 20, 30, 30, 20, -10, -30,
      -30, -10, 30, 40, 40, 30, -10, -30,
      -30, -10, 30, 40, 40, 30, -10, -30,
      -30, -10, 20, 30, 30, 20, -10, -30,
      -30, -20, -10, 0, 0, -10, -20, -30,
      -50, -40, -30, -20, -20, -30, -40, -50]

    kingEvalEndGameBlack: array[64, int] = kingEvalEndGameWhite.reverse()

proc moveValue(board: chess.Board, move: chess.Move, endgame: bool): float =
  if move.promotion != NONE:
    return if board.turn == BLACK: float(-inf) else: float(inf)

  let piece = board.piece_at(move.from_square)
  if piece.is_some:
    let fromValue = float(evaluatePiece(piece.get, move.from_square, endgame))
    let toValue = float(evaluatePiece(piece.get, move.to_square, endgame))
    let positionChange = toValue - fromValue
  else:
    raise newException("A piece was expected at ", $move.from_square)

  var captureValue = 0.0
  if board.is_capture(move):
    captureValue = evaluateCapture(board, move)

  var currentMoveValue = captureValue + positionChange
  if board.turn == BLACK:
    currentMoveValue = -currentMoveValue

  return currentMoveValue

proc evaluateCapture(board: chess.Board, move: chess.Move): float =
      if board.is_en_passant(move):
        return float(pieceValue[PAWN])
      let to = board.piece_at(move.to_square)
      let from = board.piece_at(move.from_square)
      if to.isNone() or from.isNone():
        raise newException(
          msg="Pieces were expected at both ",
          `&` & move.to_square & " and " & move.from_square
        )
      return float(pieceValue[to.piece_type] - pieceValue[from.piece_type])

proc evaluatePiece(piece: chess.Piece, square: chess.Square, endGame: bool): int =
      let pieceType = piece.pieceType
      var mapping: seq[int]
      if pieceType == PAWN:
        mapping = if piece.color == WHITE: pawnEvalWhite else pawnEvalBlack
      if pieceType == KNIGHT:
        mapping = knightEval
      if pieceType == BISHOP:
        mapping = if piece.color == WHITE: bishopEvalWhite else bishopEvalBlack
      if pieceType == ROOK:
        mapping = if piece.color == WHITE: rookEvalWhite else rookEvalBlack
      if pieceType == QUEEN:
        mapping = queenEval
      if pieceType == KING:
        # use end game piece-square tables if neither side has a queen
        if endGame:
          mapping = if piece.color == WHITE: kingEvalEndGameWhite else kingEvalEndGameBlack
        else:
          mapping = if piece.color == WHITE: kingEvalWhite else kingEvalBlack
      return mapping[square]

    if move.promotion != NONE:
      result = if board.turn == BLACK: float(-inf) else float(inf)
    elif board.piece_at(move.from_square).isSome:
      let fromValue = float(evaluatePiece(board.piece_at(move.from_square), move.from_square, endgame)
      let toValue = float(evaluatePiece(board.piece_at(move.from_square), move.to_square, endgame)
      let positionChange = toValue - fromValue
      let captureValue = if board.is_capture(move): evaluateCapture(board, move) else 0.0
      result = float(captureValue + positionChange)
    else:
      raise newException("A piece was expected at ", $move.from_square)

    if board.turn == BLACK:
      result = -result
    return result

proc evaluateBoard(board: chess.Board): float =
  var total = 0.0
  var endGame = checkEndGame(board)

  for square in SQUARES:
    if let piece = board.piece_at(square):
      let pieceType = piece.piece_type
      let value = pieceValue[piece_type] + evaluatePiece(piece, square, endGame)
      total += if piece.color == WHITE: value else -value

  return total

proc checkEndGame(board: chess.Board): bool =
      var queens = 0
      var minors = 0
      for square in SQUARES:
        if let piece = board.piece_at(square):
          case piece.piece_type
          of QUEEN:
            queens += 1
          of BISHOP, KNIGHT:
            minors += 1
      if queens == 0 or (queens == 2 and minors <= 1):
        return true
      return false

    var total = 0.0
    let endgame = checkEndGame(board)

for square in SQUARES:
      if let piece = board.piece_at(square):
        let value = pieceValue[piece.piece_type] + evaluatePiece(piece, square, endgame)
        total += if piece.color == WHITE: value else -value

return total