import numpy as np
import chess
import tensorflow as tf
from keras.models import load_model

class Model_handler :
  def __init__(self) :
    self.model = load_model("model.h5")
    print("Model loaded and is ready to use !!!")

  def predict(self, state, color) :
    prediction = self.model(self.convert_state_to_input(state, color))
    return self.convert_output_to_probs(state, color, prediction[1][0])

  def convert_state_to_input(self, state, color) :
    if type(state) == str :
      temp = state.split("_")
      temp_arr = np.zeros((8,8,12), dtype = np.float16)
      arr2 = np.zeros((8,8,10), dtype = np.float16)
      temp_arr, arr2 = self.convert_board_to_input(temp[-1], color)
      for i in range(1,5) :
        temp_arr = np.concatenate([self.convert_board_to_input(temp[-1-i], color, False), temp_arr], axis = 2)
      return (np.expand_dims(temp_arr, axis = 0), np.expand_dims(arr2, axis = 0))

  def convert_board_to_input(self, state, color, current = True) :
    if current :
      board = chess.Board(state)
      board.turn = color
      arr = np.zeros((8,8,12), dtype = np.float16)
      arr2 = np.zeros((8,8,10), dtype = np.float16)
      piece_to_value = self.get_piece_to_value(color)
      piece_to_value2 = self.get_piece_to_value(color, False)
      for i in range(64) :
        if(board.piece_at(i) is not None) :
          arr[i//8,i%8,piece_to_value[board.piece_at(i).symbol()]] = 1
      for move in board.legal_moves :
        square = move.to_square
        arr[square//8, square%8, piece_to_value[board.piece_at(move.from_square).symbol()]] = 0.5
        symbol = board.piece_at(move.from_square).symbol()
        if move.promotion is not None :
          arr2[move.promotion-2, move.from_square%8, 9] = 1
        else :
          arr2[square//8, square%8, piece_to_value2[symbol]] = 1
          if(piece_to_value2[symbol] == 1 or piece_to_value2[symbol] == 3 or piece_to_value2[symbol] == 5) :
            piece_to_value2[symbol] += 1
      return (arr, arr2)
    else :
      arr = np.zeros((8,8,12), dtype = np.float16)
      if len(state) == 0 :
        return arr
      board = chess.Board(state)
      board.turn = color
      piece_to_value = self.get_piece_to_value(color)
      for i in range(64) :
        if(board.piece_at(i) is not None) :
          arr[i//8,i%8,piece_to_value[board.piece_at(i).symbol()]] = 1
      return arr
    
  def convert_output_to_probs(self, state, color, policy_output) :
    policy = tf.reshape(policy_output, [8,8,10])
    board = chess.Board(state.split("_")[-1])
    board.turn = color
    piece_to_value = self.get_piece_to_value(color,False)
    move_dict = {}
    for move in list(board.legal_moves) :
      to_square = move.to_square
      from_square = move.from_square
      piece_type = piece_to_value[board.piece_at(from_square).symbol()]
      if move.promotion is not None :
        move_dict[move.uci()] = policy[move.promotion-2, from_square%8, 9]
      else :
        move_dict[move.uci()] = policy[to_square//8, to_square%8, piece_type]
        if(piece_type == 1 or piece_type == 3 or piece_type == 5) :
          piece_to_value[board.piece_at(from_square).symbol()] += 1
   # print(list(board.legal_moves))
    move = [item[0] for item in sorted(move_dict.items(), key = lambda x: x[1], reverse = True)][0]
    print(move_dict[move])
    return move
  
  def get_piece_to_value(self, color, inp = True) :
    if inp :
      if(color == 1) :
        return {
        'P': 0, 'N': 1, 'B': 2, 'R': 3, 'Q': 4, 'K': 5,
        'p': 6, 'n':7, 'b': 8, 'r': 9, 'q': 10, 'k': 11
        }
      return {
        'p': 0, 'n': 1, 'b': 2, 'r': 3, 'q': 4, 'k': 5,
        'P': 6, 'N':7, 'B': 8, 'R': 9, 'Q': 10, 'K': 11
      }
    else :
      if(color == 1) :
        return {
        'P': 0, 'N': 1, 'B': 3, 'R': 5, 'Q': 7, 'K': 8
        }
      return {
        'p': 0, 'n': 1, 'b': 3, 'r': 5, 'q': 7, 'k': 8
      }