import os
import yaml
import time
import argparse
from easydict import EasyDict as edict


def parse_arguments():
  parser = argparse.ArgumentParser(
      description="Running Experiments of Deep Prediction")
  parser.add_argument(
      '-c',
      '--config_file',
      type=str,
      default="config/resnet101_cifar.json",
      required=True,
      help="Path of config file")
  parser.add_argument(
      '-l',
      '--log_level',
      type=str,
      default='INFO',
      help="Logging Level, \
        DEBUG, \
        INFO, \
        WARNING, \
        ERROR, \
        CRITICAL")
  parser.add_argument('--logdir', type=str, default=None)
  parser.add_argument('-m', '--comment', help="Experiment comment")
  parser.add_argument('-t', '--test', help="Test model", action='store_true')
  args = parser.parse_args()

  return args


def get_config(config_file, exp_dir=None, is_test=False, args=None):
  """ Construct and snapshot hyper parameters """
  config = edict(yaml.load(open(config_file, 'r'), Loader=yaml.FullLoader))
  
  if args is not None and args.logdir is not None:
    config.save_dir = args.logdir
  else:
    config.save_dir = os.path.join(config.exp_dir, config.exp_name)

  # create hyper parameters
  config.run_id = str(os.getpid())
  config.exp_name = '_'.join([
      time.strftime('%Y-%m-%d-%H-%M-%S'),
      config.exp_name, config.run_id
  ])

  if exp_dir is not None:
    config.exp_dir = exp_dir

  save_name = os.path.join(config.save_dir, 'config.yaml')
  # snapshot hyperparameters
  mkdir(config.exp_dir)
  mkdir(config.save_dir)

  yaml.dump(edict2dict(config), open(save_name, 'w'), default_flow_style=False)

  return config


def edict2dict(edict_obj):
  dict_obj = {}

  for key, vals in edict_obj.items():
    if isinstance(vals, edict):
      dict_obj[key] = edict2dict(vals)
    else:
      dict_obj[key] = vals

  return dict_obj


def mkdir(folder):
  if not os.path.isdir(folder):
    os.makedirs(folder)
