#!/usr/bin/env python
  
  from __future__ import print_function
  import re
  import os
  import argparse
  import sys
  import warnings
  import copy
  import glob
  
  
  if __name__ == "__main__":
      # we add compulsory arguments as named arguments for readability
      parser = argparse.ArgumentParser(description="Create a list of models suitable for averaging "
                                                   "based on their train objf values.",
                                       epilog="See steps/nnet3/lstm/train.sh for example.")
  
      parser.add_argument("--difference-threshold", type=float,
                          help="The threshold for discarding models, "
                          "when objf of the model differs more than this value from the best model "
                          "it is discarded.",
                          default=1.0)
  
      parser.add_argument("num_models", type=int,
                          help="Number of models.")
  
      parser.add_argument("logfile_pattern", type=str,
                          help="Pattern for identifying the log-file names. "
                          "It specifies the entire log file name, except for the job number, "
                          "which is replaced with '%'. e.g. exp/nneet3/tdnn_sp/log/train.4.%.log")
  
  
      args = parser.parse_args()
  
      assert(args.num_models > 0)
  
      parse_regex = re.compile("LOG .* Overall average objective function for 'output' is ([0-9e.\-+]+) over ([0-9e.\-+]+) frames")
      loss = []
      for i in range(args.num_models):
          model_num = i + 1
          logfile = re.sub('%', str(model_num), args.logfile_pattern)
          lines = open(logfile, 'r').readlines()
          this_loss = -100000
          for line_num in range(1, len(lines) + 1):
              # we search from the end as this would result in
              # lesser number of regex searches. Python regex is slow !
              mat_obj = parse_regex.search(lines[-1*line_num])
              if mat_obj is not None:
                  this_loss = float(mat_obj.groups()[0])
                  break;
          loss.append(this_loss);
      max_index = loss.index(max(loss))
      accepted_models = []
      for i in range(args.num_models):
          if (loss[max_index] - loss[i]) <= args.difference_threshold:
              accepted_models.append(i+1)
  
      model_list = " ".join([str(x) for x in accepted_models])
      print(model_list)
  
      if len(accepted_models) != args.num_models:
          print("WARNING: Only {0}/{1} of the models have been accepted for averaging, based on log files {2}.".format(len(accepted_models), args.num_models, args.logfile_pattern), file=sys.stderr)
          print("         Using models {0}".format(model_list), file=sys.stderr)