#!/usr/bin/env python3
import ucf
import argparse
import os
import sys
import traceback
import numpy as np

parser = argparse.ArgumentParser(description='Verifies that content of single UCF file is equivalent to multiple UCF files')
parser.add_argument('file1', metavar='file_single',help='input file #1 (single file)')
parser.add_argument('file2', metavar='filebase_multi',help='input file #2 (multiple files)')
parser.add_argument("--debug", help="enable debug output? [default: False]", action="store_true")
parser.add_argument("--stop", help="stop as soon as an error has been detected? [default: False]", action="store_true")
parser.add_argument("--tolerance", metavar='tol', help="floating point tolerance for comparison [default: 1e-12]", type=float)
parser.add_argument("-v","--verbose", help="enable verbose output? [default: False]", action="store_true")
args = parser.parse_args()

file_single  = args.file1
fiba_multi   = args.file2
flag_debug   = args.debug
flag_verbose = args.verbose
fptol        = args.tolerance
flag_stop    = args.stop

if fptol is None:
  fptol = 1e-12

if os.path.isfile(file_single)==False:
  print('File not found:',os.path.isfile(file_single))
  sys.exit(1)

objs = ucf.UCF(file=file_single,verbosity=flag_verbose,debug=flag_debug)
#if obj1.NumTimestep!=obj2.NumTimestep:

errfound=False
for istep in range(0,objs.NumTimestep):
  print('Comparing rank {:5d}'.format(istep))
  file_multi = '{}{:05d}'.format(fiba_multi,istep)
  if os.path.isfile(file_multi)==False:
    print('File not found:',os.path.isfile(file_multi))
    if flag_stop:
      sys.exit(1)
    else:
      errfound=True
  objm = ucf.UCF(file=file_multi,verbosity=flag_verbose,debug=flag_debug)
  if objs._UCF__numSetPerStep[istep]!=objm._UCF__numSetPerStep[0]:
    print("Number of datasets differs: {:d}, {:d}".format(objs._UCF__numSetPerStep[istep],objm._UCF__numSetPerStep[0]))
    if flag_stop:
      sys.exit(2)
    else:
      errfound=True
  for iset in range(0,objs._UCF__numSetPerStep[istep]):
    (data1,param1) = objs.readSet(step=istep+1,dset=iset+1)
    (data2,param2) = objm.readSet(step=1,dset=iset+1)
    if not param1==param2:
      print('Step #{:d}, set #{:d}: parameter mismatch. {}, {}'.format(istep,iset,param1,param2))
      if flag_stop:
        sys.exit(3)
      else:
        errfound=True
    if not np.allclose(data1,data2,rtol=fptol):
      data_diff = np.abs(data1-data2)
      idx = np.argmax(data_diff)
      vald = data_diff.flatten()[idx]
      val1 = data1.flatten()[idx]
      val2 = data2.flatten()[idx]
      print('Step #{:d}, set #{:d}: difference in FP data above tolerance. idx={}, diff={}, val1={}, val2={}'.format(istep,iset,idx,vald,val1,val2))
      if flag_stop:
        sys.exit(4)
      else:
        errfound=True
if errfound:
  sys.exit(5)