import numpy as np
import sys

def A3PI_ParseHistory(file_in, file_out='A3PI_History'):
    """This postprocessing script for A3PI will read a history array generated
    from libEnsemble and output two text files containing the multi-objective
    optimization inputs and corresponding outputs.
    
    Note: the generation data is not parsed but can be reconstructed if a
    uniform population size is used. Each row in the text files correspond to
    an individual evaluation.
    
    The total number of rows N_p corresponds to the population size and number
    of generations. The number of columns in each of the output files depend on
    the dimension of the input space N_i and number of objectives N_o.
    
    N_p = Total evaluations = (population size) * (generations + 1)
    N_i = Dimension of input parameter space for multi-objective optimization
    N_o = Number of objectives in multi-objective optimization
    
    Inputs:
        file_in = filename of libEnsemble history array in .npy format
        
    Outputs:
        file_out_inputs  = text file with N_p rows and N_i columns of data
        file_out_outputs = text file with N_p rows and N_o columns of data
        
    Note: filenames are created with '_inputs' or '_outputs' suffixes for the
    corresponding files given with the file_out prefix name.
    
    Written by David Bizzozero
    March 29, 2021
    """
    
    data = np.load(file_in) #Input history array in libEnsemble format
    N_p = len(data)         #Number of total evaluations in history array
    
    #Write list of inputs to file named with suffix "_inputs"
    with open(file_out+'_inputs','w') as file:
        for row in range(N_p):
            datarow = np.array2string(data[row][4]).replace('\n', '').replace('[', '').replace(']', '')
            file.write(datarow)
            file.write('\n')
    
    #Write list of outputs to file named with suffix "_outputs"
    with open(file_out+'_outputs','w') as file:
        for row in range(N_p):
            datarow = np.array2string(data[row][5]).replace('\n', '').replace('[', '').replace(']', '')
            file.write(datarow)
            file.write('\n')

#This section allows A3PI_ParseHistory to be called like an executable with an
#input file (and optiionally output filename)
if __name__=='__main__':
            
    if len(sys.argv)==2:
        file_in = sys.argv[1]
        A3PI_ParseHistory(file_in)
        
    elif len(sys.argv)==3:
        file_in = sys.argv[1]
        file_out = sys.argv[2]
        A3PI_ParseHistory(file_in, file_out)

    else:
        print('No A3PI history file specified to parse.')