import chess
import chess.pgn
import numpy as np

def fen_to_numeric(char):
    
    # Switches fen characters to numeric piece values
    # White pieces are positive numbers and black pieces are negative numbers
    #
    # argument : A non-numeric character in a FEN string
    # return   : Numeric value of input piece
    
    switcher = {
        'r' : -5,
        'n' : -3,
        'b' : -4,
        'q' : -9,
        'k' : -10,
        'p' : -1,
        'R' : 5,
        'N' : 3,
        'B' : 4,
        'Q' : 9,
        'K' : 10,
        'P' : 1,
    }
    return(switcher.get(char))

def get_board_state(board): 
    
    # Takes a chess board and converts the PGN to FEN notation. Then subs the 
    # FEN characters for the value of each piece and splits the result into an array
    #
    # argument : chess board object
    # return   : input array for NN

    shredder_fen = list((board.shredder_fen().split(" ")[0]).replace("/",""))
    numeric_fen = []
    for char in shredder_fen:
        if(char.isdigit()): # empty squares between pieces: add zeros
            for i in range(int(char)):
                numeric_fen.append(0)
        else: # convert piece notation to numeric value
            numeric_fen.append(fen_to_numeric(char))
    
    # Whose turn is it?
    if(board.turn == True):
        numeric_fen + [1] # White to move
    else:
        numeric_fen + [0] # Black to move
    
    # Is the king in check?
    if(board.is_check() == True):
        numeric_fen += [1] # Yes
    else:
        numeric_fen += [0] # No
    
    # Is queenside castling available for player?
    if(board.has_queenside_castling_rights(board.turn)):
        numeric_fen += [1] # Yes
    else:
        numeric_fen += [0] # No
        
    # Is kingside castling available for player?
    if(board.has_kingside_castling_rights(board.turn)):
        numeric_fen += [1] # Yes
    else:
        numeric_fen += [0] # No
        
    # Is queenside castling available for opponent?
    if(board.has_queenside_castling_rights(not(board.turn))):
        numeric_fen += [1] # Yes
    else:
        numeric_fen += [0] # No
        
    # Is kingside castling available for player?
    if(board.has_kingside_castling_rights(not(board.turn))):
        numeric_fen += [1] # Yes
    else:
        numeric_fen += [0] # No

    return(numeric_fen)

def get_legal_moves(board, origin):
    
    # Takes a board and an origin (from_square) and finds all legal moves originating on that square
    #
    # argument : chess board object and from_square
    # return   : len 64 list representing board, with 0s for illegal moves and 1s for valid moves
    
    to_squares = []
    for move in board.legal_moves:
        if move.from_square == origin:
            to_squares.append(move.to_square)
    valid_dest_list = []
    for i in range(64):
        if i in to_squares:
            valid_dest_list.append(1)
        else:
            valid_dest_list.append(0)
    
    return(valid_dest_list)

def gen_mappings(board):
    
    # Takes a board and generates two 64 len lists to show the squares that white can capture and the squares black can capture
    #
    # argument : chess board object and from_square
    # return   : len 64 list representing board, with 0s for non-captures and ints for captures
    
    white_check_map = [0] * 64 
    black_check_map = [0] * 64 
    white_capture_map = [0] * 64
    black_capture_map = [0] * 64
    
    for move in board.legal_moves:
        if board.turn == True:
            if board.gives_check(move):
                white_check_map[move.to_square] += 1
            if board.is_capture(move):
                white_capture_map[move.to_square] += 1
        else:
            if board.gives_check(move):
                white_check_map[move.to_square] -= 1
            if board.is_capture(move):
                black_capture_map[move.to_square] -= 1
                
    # pass the turn to the other player
    board.push(chess.Move.null())
    
    for move in board.legal_moves:
        if board.turn == True:
            if board.gives_check(move):
                white_check_map[move.to_square] += 1
            if board.is_capture(move):
                white_capture_map[move.to_square] += 1
        else:
            if board.gives_check(move):
                white_check_map[move.to_square] -= 1
            if board.is_capture(move):
                black_capture_map[move.to_square] -= 1
    
    # reset the board to the original state
    board.pop()
    
    return(white_check_map, black_check_map, white_capture_map, black_capture_map)

def game_to_data(game):
    
    # Takes a chess game and processes it into a inputs to the NN
    # The board state is a numeric representation generated by get_board_state()
    # target_origin move is the space from which a piece is moved: target for first NN
    # target_dest move is the space to which a piece is moved: target for second NN
    #
    # argument : chess game object
    # return   : [flattened board states, target origin moves, target destination moves]
        
    X_SaP = []
    X_MaP = []
    target_origin = []
    target_dest = []
    target_piece = []
    board_color = []
    board_period = []
    
    board = game.board()
    
    for move in game.mainline_moves():
        
        white_check_map_tmp, black_check_map_tmp, white_cap_map_tmp, black_cap_map_tmp = gen_mappings(board)
        
        # add data about this board state to lists
        X = get_board_state(board) + white_check_map_tmp + black_check_map_tmp + white_cap_map_tmp + black_cap_map_tmp
        X_SaP.append(X)
        X_MaP.append(X + get_legal_moves(board, move.from_square))
        target_origin.append(move.from_square)
        target_dest.append(move.to_square)
        target_piece.append(str(board.piece_at(move.from_square)))
        board_color.append(board.turn)
        board_period.append(board.ply())
        board.push(move) # play the next move on the board
    
    return X_SaP, X_MaP, target_origin, target_dest, target_piece, board_color, board_period

def aggregate_game_data(game1, game2):
    
    # Combines game data 
    #
    # argument : chess game object
    # return   : [flattened board states, target origin moves, target destination moves]
    
    X_SaP_g1              = game1[0]
    X_MaP_g1              = game1[1]
    target_origin_g1      = game1[2]
    target_dest_g1        = game1[3]
    target_piece_g1       = game1[4]
    board_color_g1        = game1[5]
    board_period_g1       = game1[6]
    
    X_SaP_g2              = game2[0]
    X_MaP_g2              = game2[1]
    target_origin_g2      = game2[2]
    target_dest_g2        = game2[3]
    target_piece_g2       = game2[4]
    board_color_g2        = game2[5]
    board_period_g2       = game2[6]

    return X_SaP_g1 + X_SaP_g2, X_MaP_g1 + X_MaP_g2, target_origin_g1 + target_origin_g2, target_dest_g1 + target_dest_g2, target_piece_g1 + target_piece_g2, board_color_g1 + board_color_g2, board_period_g1 + board_period_g2

def gen_nn_data(pgn_games):
    
    # Wrapper for the other functions. Goes through the list of games and converts the
    # games to training and target data
    #
    # argument : list of chess game objects
    # return   : [flattened board states, target origin moves, target destination moves]

    nn_data = game_to_data(pgn_games[0])
    pgn_games = pgn_games[1:]
    
    for game in pgn_games:
        nn_data = aggregate_game_data(nn_data, game_to_data(game))
    
    return np.array(nn_data[0]), np.array(nn_data[1]), np.array(nn_data[2]), np.array(nn_data[3]).astype(int), np.array(nn_data[4]), np.array(nn_data[5]), np.array(nn_data[6])

def read_pgn_from_file(files, depth = np.inf):
    
    # Reads in all of the games in PGN form from a .pgn file and saves them as a list
    #
    # argument : .pgn file path
    # return   : list of chess game objects
    
    pgn_games = []
    depth_count = 0
    
    for file in files:
        pgn = open(file)
        while(True):
            game = chess.pgn.read_game(pgn)
            if(game == None):
                break
            pgn_games.append(game)
            depth_count += 1
            if depth_count >= depth:
                return pgn_games 
            
    return(pgn_games)

def mod_MaP_inputs(X, y1, legal_moves):

    # Used for the Move-a-Piece NN. Adds the selected piece (y1 target) and the available legal moves as an input to the second NN
    #
    # argument : SaP input array
    # return   : MaP input Array (72 elements per row)
    
    MaP_input = []
    for i in range(len(X)):
        MaP_input.append(np.append(X[i], y1[i], legal_moves[i]))
    return np.array(MaP_input)