# Path: util/tools.py
# def batch_flatten(ndarry):
#     return ndarry.reshape(ndarry.shape[0], -1)
# 
# Path: util/gentensor.py
# def gentensor(selfpos, emypos):
#     '''
#     selfpos, emypos: ndarray [9,10,16]
#     return:
#         selfmove, emymove, selfprot, selfprot: ndarray [9,10,16]
#     '''
#     fen = tensor2fen(selfpos, emypos)
#     move = genmove.gen(fen)
#     fen = fen + '\t' + move.replace('\n','\t')
#     selfpos, emypos, selfmove, emymove, selfprot, emyprot = fen2tensor(fen)
#     return [selfmove, emymove, selfprot, emyprot]


import tensorflow as tf
from util.tools import batch_flatten
from util.gentensor import gentensor
from config import config
        next_self_prot = np.zeros(self_prot.shape)
        next_self_pos[mask]  = next_mask_self_pos
        next_enemy_pos[mask] = next_mask_enemy_pos
        next_self_move[mask] = next_mask_self_move
        next_emy_move[mask]  = next_mask_emy_move
        next_self_prot[mask] = next_mask_self_prot
        next_emy_prot[mask]  = next_mask_enemy_prot

        return reward, [next_self_pos, next_enemy_pos], if_terminated, next_self_move, [next_emy_move, next_self_prot, next_emy_prot]

def valid_move(states):
    '''
    states = [self_pos, enemy_pos]
    self_pos.shape = [None, 9, 10, 16]
    enemy_pos.shape = [None, 9, 10, 16]

    return valid_moves.shape = [None, 9, 10, 16]
    '''

    return gentensor(*states)


def change_state(states, move):
    '''
    states = [self_pos, enemy_pos]
    self_pos.shape = [None, 9, 10, 16]
    enemy_pos.shape = [None, 9, 10, 16]

    return valid_moves.shape = [None, 9, 10, 16]
    '''

# Path: util/tools.py
# def batch_flatten(ndarry):
#     return ndarry.reshape(ndarry.shape[0], -1)
# 
# Path: util/gentensor.py
# def gentensor(selfpos, emypos):
#     '''
#     selfpos, emypos: ndarray [9,10,16]
#     return:
#         selfmove, emymove, selfprot, selfprot: ndarray [9,10,16]
#     '''
#     fen = tensor2fen(selfpos, emypos)
#     move = genmove.gen(fen)
#     fen = fen + '\t' + move.replace('\n','\t')
#     selfpos, emypos, selfmove, emymove, selfprot, emyprot = fen2tensor(fen)
#     return [selfmove, emymove, selfprot, emyprot]


# Path: rl_train/value_dataset.py
import tensorflow as tf
from util.tools import batch_flatten
from util.gentensor import gentensor
from config import config

        next_self_prot = np.zeros(self_prot.shape)
        next_self_pos[mask]  = next_mask_self_pos
        next_enemy_pos[mask] = next_mask_enemy_pos
        next_self_move[mask] = next_mask_self_move
        next_emy_move[mask]  = next_mask_emy_move
        next_self_prot[mask] = next_mask_self_prot
        next_emy_prot[mask]  = next_mask_enemy_prot

        return reward, [next_self_pos, next_enemy_pos], if_terminated, next_self_move, [next_emy_move, next_self_prot, next_emy_prot]

def valid_move(states):
    '''
    states = [self_pos, enemy_pos]
    self_pos.shape = [None, 9, 10, 16]
    enemy_pos.shape = [None, 9, 10, 16]

    return valid_moves.shape = [None, 9, 10, 16]
    '''

    return gentensor(*states)


def change_state(states, move):
    '''
    states = [self_pos, enemy_pos]
    self_pos.shape = [None, 9, 10, 16]
    enemy_pos.shape = [None, 9, 10, 16]

    return valid_moves.shape = [None, 9, 10, 16]
    '''

    move = np.argmax(batch_flatten(move), axis=1)