text: # Copyright (c) Facebook, Inc. and its affiliates.
  import collections
  import json
  import logging
  import os
  import warnings
  from ast import literal_eval
  
  import demjson
  import torch
  from mmf.common.registry import registry
  from mmf.utils.env import import_user_module
  from mmf.utils.file_io import PathManager
  from mmf.utils.general import get_absolute_path, get_mmf_root
  from omegaconf import OmegaConf
  
  
  logger = logging.getLogger(__name__)
  
  
  def load_yaml(f):
      # Convert to absolute path for loading includes
      abs_f = get_absolute_path(f)
  
      try:
          mapping = OmegaConf.load(abs_f)
          f = abs_f
      except FileNotFoundError as e:
          # Check if this file might be relative to root?
          # TODO: Later test if this can be removed
          relative = os.path.abspath(os.path.join(get_mmf_root(), f))
          if not PathManager.isfile(relative):
              raise e
          else:
              f = relative
              mapping = OmegaConf.load(f)
  
      if mapping is None:
          mapping = OmegaConf.create()
  
      includes = mapping.get("includes", [])
  
      if not isinstance(includes, collections.abc.Sequence):
          raise AttributeError(
              "Includes must be a list, {} provided".format(type(includes))
          )
  
      include_mapping = OmegaConf.create()
  
      mmf_root_dir = get_mmf_root()
  
      for include in includes:
          original_include_path = include
          include = os.path.join(mmf_root_dir, include)
  
          # If path doesn't exist relative to MMF root, try relative to current file
          if not PathManager.exists(include):
              include = os.path.join(os.path.dirname(f), original_include_path)
  
          current_include_mapping = load_yaml(include)
          include_mapping = OmegaConf.merge(include_mapping, current_include_mapping)
  
      mapping.pop("includes", None)
  
      mapping = OmegaConf.merge(include_mapping, mapping)
  
      return mapping
  
  
  def get_default_config_path():
      directory = os.path.dirname(os.path.abspath(__file__))
      configs_dir = os.path.join(directory, "..", "configs")
  
      # Check for fb defaults
      fb_defaults = os.path.join(configs_dir, "fb_defaults.yaml")
      if PathManager.exists(fb_defaults):
          return fb_defaults
      else:
          return os.path.join(configs_dir, "defaults.yaml")
  
  
  def load_yaml_with_defaults(f):
      default_config = get_default_config_path()
      return OmegaConf.merge(load_yaml(default_config), load_yaml(f))
  
  
  def get_zoo_config(
      key, variation="defaults", zoo_config_path=None, zoo_type="datasets"
  ):
      version = None
      resources = None
      if zoo_config_path is None:
          zoo_config_path = os.path.join("configs", "zoo", f"{zoo_type}.yaml")
      zoo = load_yaml(zoo_config_path)
  
      # Set struct on zoo so that unidentified access is not allowed
      OmegaConf.set_struct(zoo, True)
  
      try:
          item = OmegaConf.select(zoo, key)
      except Exception:
          # Key wasn't present or something else happened, return None, None
          return version, resources
  
      if not item:
          return version, resources
  
      if variation not in item:
          # If variation is not present, then key value should
          # be directly returned if "defaults" was selected as the variation
          assert (
              variation == "defaults"
          ), f"'{variation}' variation not present in zoo config"
          return _get_version_and_resources(item)
      elif "resources" in item:
          # Case where full key is directly passed
          return _get_version_and_resources(item)
      else:
          return _get_version_and_resources(item[variation])
  
  
  def _get_version_and_resources(item):
      assert "version" in item, "'version' key should be present in zoo config {}".format(
          item._get_full_key("")
      )
      assert (
          "resources" in item
      ), "'resources' key should be present in zoo config {}".format(
          item._get_full_key("")
      )
  
      return item.version, item.resources
  
  
  def get_global_config(key=None):
      config = registry.get("config")
      if config is None:
          configuration = Configuration()
          config = configuration.get_config()
          registry.register("config", config)
  
      if key:
          config = OmegaConf.select(config, key)
  
      return config
  
  
  def get_mmf_cache_dir():
      config = get_global_config()
      cache_dir = config.env.cache_dir
      # If cache_dir path exists do not join to mmf root
      if not os.path.exists(cache_dir):
          cache_dir = os.path.join(get_mmf_root(), cache_dir)
      return cache_dir
  
  
  def get_mmf_env(key=None):
      config = get_global_config()
      if key:
          return OmegaConf.select(config.env, key)
      else:
          return config.env
  
  
  def resolve_cache_dir(env_variable="MMF_CACHE_DIR", default="mmf"):
      # Some of this follow what "transformers" does for there cache resolving
      try:
          from torch.hub import _get_torch_home
  
          torch_cache_home = _get_torch_home()
      except ImportError:
          torch_cache_home = os.path.expanduser(
              os.getenv(
                  "TORCH_HOME",
                  os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"),
              )
          )
      default_cache_path = os.path.join(torch_cache_home, default)
  
      cache_path = os.getenv(env_variable, default_cache_path)
  
      if not PathManager.exists(cache_path):
          try:
              PathManager.mkdirs(cache_path)
          except PermissionError:
              cache_path = os.path.join(get_mmf_root(), ".mmf_cache")
              PathManager.mkdirs(cache_path)
  
      return cache_path
  
  
  def resolve_dir(env_variable, default="data"):
      default_dir = os.path.join(resolve_cache_dir(), default)
      dir_path = os.getenv(env_variable, default_dir)
  
      if not PathManager.exists(dir_path):
          PathManager.mkdirs(dir_path)
  
      return dir_path
  
  
  class Configuration:
      def __init__(self, args=None, default_only=False):
          self.config = {}
  
          if not args:
              import argparse
  
              args = argparse.Namespace(opts=[])
              default_only = True
  
          self.args = args
          self._register_resolvers()
  
          self._default_config = self._build_default_config()
  
          if default_only:
              other_configs = {}
          else:
              other_configs = self._build_other_configs()
  
          self.config = OmegaConf.merge(self._default_config, other_configs)
  
          self.config = self._merge_with_dotlist(self.config, args.opts)
          self._update_specific(self.config)
          self.upgrade(self.config)
          # Resolve the config here itself after full creation so that spawned workers
          # don't face any issues
          self.config = OmegaConf.create(
              OmegaConf.to_container(self.config, resolve=True)
          )
          registry.register("config", self.config)
  
      def _build_default_config(self):
          self.default_config_path = get_default_config_path()
          default_config = load_yaml(self.default_config_path)
          return default_config
  
      def _build_other_configs(self):
          opts_config = self._build_opt_list(self.args.opts)
          user_config = self._build_user_config(opts_config)
  
          self._opts_config = opts_config
          self._user_config = user_config
  
          self.import_user_dir()
  
          model_config = self._build_model_config(opts_config)
          dataset_config = self._build_dataset_config(opts_config)
          args_overrides = self._build_demjson_config(self.args.config_override)
          prune_args = self._build_prune_config(opts_config)
          other_configs = OmegaConf.merge(
              model_config, dataset_config, user_config, args_overrides, prune_args
          )
  
          return other_configs
  
      def _build_opt_list(self, opts):
          opts_dot_list = self._convert_to_dot_list(opts)
          return OmegaConf.from_dotlist(opts_dot_list)
  
      def _build_user_config(self, opts):
          user_config = {}
  
          # Update user_config with opts if passed
          self.config_path = opts.config
          if self.config_path is not None:
              user_config = load_yaml(self.config_path)
  
          return user_config
  
      def import_user_dir(self):
          # Try user_dir options in order of MMF configuration hierarchy
          # First try the default one, which can be set via environment as well
          user_dir = self._default_config.env.user_dir
  
          # Now, check user's config
          user_config_user_dir = self._user_config.get("env", {}).get("user_dir", None)
  
          if user_config_user_dir:
              user_dir = user_config_user_dir
  
          # Finally, check opts
          opts_user_dir = self._opts_config.get("env", {}).get("user_dir", None)
          if opts_user_dir:
              user_dir = opts_user_dir
  
          if user_dir:
              import_user_module(user_dir)
  
      def _build_model_config(self, config):
          model = config.model
          if model is None:
              raise KeyError("Required argument 'model' not passed")
          model_cls = registry.get_model_class(model)
  
          if model_cls is None:
              warning = f"No model named '{model}' has been registered"
              warnings.warn(warning)
              return OmegaConf.create()
  
          default_model_config_path = model_cls.config_path()
  
          if default_model_config_path is None:
              warning = "Model {}'s class has no default configuration provided".format(
                  model
              )
              warnings.warn(warning)
              return OmegaConf.create()
  
          return load_yaml(default_model_config_path)
  
      def _build_prune_config(self, config):
          prune = config.prune
          num_of_examples = prune.num_of_examples
          if num_of_examples is None:
              raise KeyError("Required arguments 'prune.num_of_examples' not passed")
          positive_prune = prune.positive_prune
          if positive_prune is None:
              raise KeyError("Required arguments 'prune.positive_prune' not passed")
          is_grad = prune.is_grad
          if is_grad is None:
              raise KeyError("Required arguments 'prune.is_grad' not passed")
          prune_config = OmegaConf.create({'prune.num_of_examples':num_of_examples,'prune.positive_prune': positive_prune, 'prune.is_grad':is_grad })
          prune_config = OmegaConf.create({'prune':prune})
          return prune_config
  
      def _build_dataset_config(self, config):
          dataset = config.dataset
          datasets = config.datasets
  
          if dataset is None and datasets is None:
              raise KeyError("Required argument 'dataset|datasets' not passed")
  
          if datasets is None:
              config.datasets = dataset
              datasets = dataset.split(",")
          else:
              datasets = datasets.split(",")
  
          dataset_config = OmegaConf.create()
  
          for dataset in datasets:
              builder_cls = registry.get_builder_class(dataset)
  
              if builder_cls is None:
                  warning = f"No dataset named '{dataset}' has been registered"
                  warnings.warn(warning)
                  continue
              default_dataset_config_path = builder_cls.config_path()
  
              if default_dataset_config_path is None:
                  warning = (
                      "Dataset {}'s builder class has no default configuration "
                      + f"provided"
                  )
                  warnings.warn(warning)
                  continue
  
              dataset_config = OmegaConf.merge(
                  dataset_config, load_yaml(default_dataset_config_path)
              )
  
          return dataset_config
  
      def get_config(self):
          self._register_resolvers()
          return self.config
  
      def _build_demjson_config(self, demjson_string):
          if demjson_string is None:
              return OmegaConf.create()
  
          demjson_dict = demjson.decode(demjson_string)
          return OmegaConf.create(demjson_dict)
  
      def _get_args_config(self, args):
          args_dict = vars(args)
          return OmegaConf.create(args_dict)
  
      def _register_resolvers(self):
          OmegaConf.clear_resolvers()
          # Device count resolver
          device_count = max(1, torch.cuda.device_count())
          OmegaConf.register_resolver("device_count", lambda: device_count)
          OmegaConf.register_resolver("resolve_cache_dir", resolve_cache_dir)
          OmegaConf.register_resolver("resolve_dir", resolve_dir)
  
      def _merge_with_dotlist(self, config, opts):
          # TODO: To remove technical debt, a possible solution is to use
          # struct mode to update with dotlist OmegaConf node. Look into this
          # in next iteration
          if opts is None:
              opts = []
  
          if len(opts) == 0:
              return config
  
          # Support equal e.g. model=visual_bert for better future hydra support
          has_equal = opts[0].find("=") != -1
  
          if has_equal:
              opt_values = [opt.split("=") for opt in opts]
          else:
              assert len(opts) % 2 == 0, "Number of opts should be multiple of 2"
              opt_values = zip(opts[0::2], opts[1::2])
  
          for opt, value in opt_values:
              if opt == "dataset":
                  opt = "datasets"
  
              splits = opt.split(".")
              current = config
              for idx, field in enumerate(splits):
                  array_index = -1
                  if field.find("[") != -1 and field.find("]") != -1:
                      stripped_field = field[: field.find("[")]
                      array_index = int(field[field.find("[") + 1 : field.find("]")])
                  else:
                      stripped_field = field
                  if stripped_field not in current:
                      raise AttributeError(
                          "While updating configuration"
                          " option {} is missing from"
                          " configuration at field {}".format(opt, stripped_field)
                      )
                  if isinstance(current[stripped_field], collections.abc.Mapping):
                      current = current[stripped_field]
                  elif (
                      isinstance(current[stripped_field], collections.abc.Sequence)
                      and array_index != -1
                  ):
                      current_value = current[stripped_field][array_index]
  
                      # Case where array element to be updated is last element
                      if (
                          not isinstance(
                              current_value,
                              (collections.abc.Mapping, collections.abc.Sequence),
                          )
                          or idx == len(splits) - 1
                      ):
                          logger.info(f"Overriding option {opt} to {value}")
                          current[stripped_field][array_index] = self._decode_value(value)
                      else:
                          # Otherwise move on down the chain
                          current = current_value
                  else:
                      if idx == len(splits) - 1:
                          logger.info(f"Overriding option {opt} to {value}")
                          current[stripped_field] = self._decode_value(value)
                      else:
                          raise AttributeError(
                              "While updating configuration",
                              "option {} is not present "
                              "after field {}".format(opt, stripped_field),
                          )
  
          return config
  
      def _decode_value(self, value):
          # https://github.com/rbgirshick/yacs/blob/master/yacs/config.py#L400
          if not isinstance(value, str):
              return value
  
          if value == "None":
              value = None
  
          try:
              value = literal_eval(value)
          except ValueError:
              pass
          except SyntaxError:
              pass
          return value
  
      def freeze(self):
          OmegaConf.set_struct(self.config, True)
  
      def defrost(self):
          OmegaConf.set_struct(self.config, False)
  
      def _convert_to_dot_list(self, opts):
          if opts is None:
              opts = []
  
          if len(opts) == 0:
              return opts
  
          # Support equal e.g. model=visual_bert for better future hydra support
          has_equal = opts[0].find("=") != -1
  
          if has_equal:
              return opts
  
          return [(opt + "=" + value) for opt, value in zip(opts[0::2], opts[1::2])]
  
      def pretty_print(self):
          if not self.config.training.log_detailed_config:
              return
  
          logger.info("=====  Training Parameters    =====")
          logger.info(self._convert_node_to_json(self.config.training))
  
          logger.info("======  Dataset Attributes  ======")
          datasets = self.config.datasets.split(",")
  
          for dataset in datasets:
              if dataset in self.config.dataset_config:
                  logger.info(f"======== {dataset} =======")
                  dataset_config = self.config.dataset_config[dataset]
                  logger.info(self._convert_node_to_json(dataset_config))
              else:
                  logger.warning(f"No dataset named '{dataset}' in config. Skipping")
  
          logger.info("======  Optimizer Attributes  ======")
          logger.info(self._convert_node_to_json(self.config.optimizer))
  
          if self.config.model not in self.config.model_config:
              raise ValueError(f"{self.config.model} not present in model attributes")
  
          logger.info(f"======  Model ({self.config.model}) Attributes  ======")
          logger.info(
              self._convert_node_to_json(self.config.model_config[self.config.model])
          )
  
      def _convert_node_to_json(self, node):
          container = OmegaConf.to_container(node, resolve=True)
          return json.dumps(container, indent=4, sort_keys=True)
  
      def _update_specific(self, config):
          # tp = self.config.training
  
          # if args["seed"] is not None or tp['seed'] is not None:
          #     print(
          #         "You have chosen to seed the training. This will turn on CUDNN "
          #         "deterministic setting which can slow down your training "
          #         "considerably! You may see unexpected behavior when restarting "
          #         "from checkpoints."
          #     )
  
          # if args["seed"] == -1:
          #     self.config["training"]["seed"] = random.randint(1, 1000000)
  
          if config.learning_rate:
              if "optimizer" in config and "params" in config.optimizer:
                  lr = config.learning_rate
                  config.optimizer.params.lr = lr
  
          if not torch.cuda.is_available() and "cuda" in config.training.device:
              warnings.warn(
                  "Device specified is 'cuda' but cuda is not present. "
                  + "Switching to CPU version."
              )
              config.training.device = "cpu"
  
          return config
  
      def upgrade(self, config):
          mapping = {
              "training.resume_file": "checkpoint.resume_file",
              "training.resume": "checkpoint.resume",
              "training.resume_best": "checkpoint.resume_best",
              "training.load_pretrained": "checkpoint.resume_pretrained",
              "training.pretrained_state_mapping": "checkpoint.pretrained_state_mapping",
              "training.run_type": "run_type",
          }
  
          for old, new in mapping.items():
              value = OmegaConf.select(config, old)
              if value:
                  OmegaConf.update(config, new, value)
  
  
  # This is still here due to legacy reasons around
  # older checkpoint loading from v0.3
  class ConfigNode(collections.OrderedDict):
      pass
  