#!/usr/bin/env python3

###############################################################################
#
# Copyright (c) 2019-2022 The Board of Trustees of the University of Illinois.
# All rights reserved.
#
# Developed by:  David J. Hardy
#                Theoretical and Computational Biophysics Group
#                Beckman Institute for Advanced Science and Technology
#                University of Illinois at Urbana-Champaign
#                https://www.ks.uiuc.edu
#
# Distributed under the University of Illinois/NCSA Open Source License
# https://opensource.org/licenses/NCSA
#
###############################################################################

###############################################################################
#
# DESCRIPTION:
#
# diff_energy.py
#
# Show "diff" of energy values from two NAMD log files.
# Typical use case is running an identical simulation to compare with a
# preexisting reference simulation, testing accuracy after code changes.
# Defaults to comparing TOTAL energy.
#
###############################################################################

import sys
import string

nargs = len(sys.argv)
atitle = []
if nargs < 3:
  print("ERROR:  Must specify two NAMD log files")
  print("SYNTAX: %s [ENERGIES] TEST.log EXACT.log" % sys.argv[0])
  sys.exit(1)

# two log file names are at end command line arguments
# names in middle are titles to be matched
for i in range(1, nargs-2):
  atitle.append(sys.argv[i])
fname1 = sys.argv[nargs-2]
fname2 = sys.argv[nargs-1]

if not atitle:
  atitle.append("TOTAL")

# attempt to open both log files
try:
  f1 = open(fname1, 'r')
except IOError:
  print("ERROR: Unable to open file \"%s\"" % fname1)
  sys.exit(2)

try:
  f2 = open(fname2, 'r')
except IOError:
  print("ERROR: Unable to open file \"%s\"" % fname2)
  sys.exit(2)

# get full list of titles from the first ETITLE: line of first log file
while True:
  line1 = f1.readline()
  if not line1:
    print("ERROR: Did not file \"ETITLE:\" in log file \"%s\"" % fname1)
    sys.exit(3)
  else:
    s1 = line1.split()
    if s1 and s1[0] == "ETITLE:":
      titles = s1[1:]
      break

# make sure that atitle has title entries that match titles
for t in atitle:
  try:
    titles.index(t)
  except ValueError:
    print("ERROR: \"%s\" is not a valid title" % t)
    sys.exit(4)

# find ENERGY: lines from each file
# report relative errors for corresponding atitle values
print("%7s %10s  %12s  %12s" % ("step", "energy", "abserr", "relerr"))
done = 0
while not done:

  # find next ENERGY: line from log file 1
  while True:
    line1 = f1.readline()
    if not line1:
      done = 1  # signal to leave outer while loop
      break     # break out of inner while loop
    s1 = line1.split()
    if s1 and s1[0] == "ENERGY:":
      value1 = s1[1:]  # store values aligned with titles
      break

  # find next ENERGY: line from log file 2
  while True:
    line2 = f2.readline()
    if not line2:
      done = 1  # signal to leave outer while loop
      break     # break out of inner while loop
    s2 = line2.split()
    if s2 and s2[0] == "ENERGY:":
      value2 = s2[1:]  # store values aligned with titles
      break

  if done: break  # at least one file has been read entirely

  if int(value1[0]) != int(value2[0]):
    print("ERROR: Timesteps are not aligned")
    sys.exit(5)

  # consider value1 as approx and value2 as exact
  ts = int(value1[0])
  for t in atitle:
    i = titles.index(t)
    e1 = float(value1[i])
    e2 = float(value2[i])
    aerr = e2 - e1  # absolute error
    if e2 != 0:  rerr = aerr / e2  # relative error
    else:        rerr = 1
    print("%7d %10s  %12g  %12g" % (ts, t, aerr, rerr))

f1.close()
f2.close()