#!/usr/bin/env python

'''Script to compare libEnsemble history arrays in files.

E.g., ./compare_npy.py out1.npy out2.npy

If two *.npy files are provided they are compared with each other.
If one *.npy file is provided if is compared with a hard-coded expected file
(by default located at ../expected.npy)

Default tolerances used for comparison are (rtol=1e-05, atol=1e-08). These
can be overwritten with -r (--rtol) and -a (--atol) flags.

E.g., ./compare_npy.py out1.npy out2.npy -r 1e-03

Nans compare as equal. Variable fields (such as those containing a time)
are ignored. In some cases you may have to ignore further user-defined fields

'''
import sys
import numpy as np
import argparse

desc = "Script to compare libEnsemble history arrays in files"
example = '''examples:

 ./compare_npy.py out1.npy out2.npy
 ./compare_npy.py out1.npy out2.npy --rtol 1e-03 --atol 1e-06
 '''

exclude_fields = ['gen_worker', 'sim_worker', 'gen_time', 'given_time']  # list of fields to ignore
locate_mismatch = True

parser = argparse.ArgumentParser(description=desc, epilog=example,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('-r', '--rtol', dest='rtol', type=float, default=1e-05, help='rel. tolerance')
parser.add_argument('-a', '--atol', dest='atol', type=float, default=1e-08, help='abs. tolerance')
parser.add_argument('args', nargs='*', help='*.npy files to compare')
args = parser.parse_args()

rtol = args.rtol
atol = args.atol
files = args.args

if len(files) >= 1:
    results = np.load(files[0])
    exp_results = np.load(files[1]) if len(files) >= 2 else np.load("../expected.npy")
else:
    parser.print_help()
    sys.exit()

compare_fields = tuple(filter(lambda x: x not in exclude_fields, exp_results.dtype.names))
match = all([np.allclose(exp_results[name], results[name],
                         rtol=rtol, atol=atol, equal_nan=True)
            for name in compare_fields])

print('Compare results: {}\n'.format(match))

if not locate_mismatch:
    assert match, 'Error: Results do NOT match'

if not match:
    for name in compare_fields:
        for i in range(len(results)):
            assert np.allclose(exp_results[name][i],
                   results[name][i], rtol=rtol, atol=atol, equal_nan=True), \
                'Mismatch in row ' + str(i) + ' field: ' + name + '. ' \
                + str(exp_results[name][i]) + ' ' + str(results[name][i])
