# Copyright (c) Facebook, Inc. and its affiliates.
import collections
import json
import os
import warnings
from ast import literal_eval

import demjson
import torch
from omegaconf import OmegaConf

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


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_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", "{}.yaml".format(zoo_type))
    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"
        ), "'{}' variation not present in zoo config".format(variation)
        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)

        registry.register("config", self.config)

    def _build_default_config(self):
        self.default_config_path = self._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)
        other_configs = OmegaConf.merge(
            model_config, dataset_config, user_config, args_overrides
        )

        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 = "No model named '{}' has been registered".format(model)
            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_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 = "No dataset named '{}' has been registered".format(dataset)
                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 "
                    + "provided".format(dataset)
                )
                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),
                    ):
                        print("Overriding option {} to {}".format(opt, 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:
                        print("Overriding option {} to {}".format(opt, 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

        self.writer = registry.get("writer")

        self.writer.write("=====  Training Parameters    =====", "info")
        self.writer.write(self._convert_node_to_json(self.config.training), "info")

        self.writer.write("======  Dataset Attributes  ======", "info")
        datasets = self.config.datasets.split(",")

        for dataset in datasets:
            if dataset in self.config.dataset_config:
                self.writer.write("======== {} =======".format(dataset), "info")
                dataset_config = self.config.dataset_config[dataset]
                self.writer.write(self._convert_node_to_json(dataset_config), "info")
            else:
                self.writer.write(
                    "No dataset named '{}' in config. Skipping".format(dataset),
                    "warning",
                )

        self.writer.write("======  Optimizer Attributes  ======", "info")
        self.writer.write(self._convert_node_to_json(self.config.optimizer), "info")

        if self.config.model not in self.config.model_config:
            raise ValueError(
                "{} not present in model attributes".format(self.config.model)
            )

        self.writer.write(
            "======  Model ({}) Attributes  ======".format(self.config.model), "info"
        )
        self.writer.write(
            self._convert_node_to_json(self.config.model_config[self.config.model]),
            "info",
        )

    def _convert_node_to_json(self, node):
        container = OmegaConf.to_container(node, resolve=True)
        return json.dumps(container, indent=4, sort_keys=True)

    def _get_default_config_path(self):
        directory = os.path.dirname(os.path.abspath(__file__))
        return os.path.join(directory, "..", "configs", "defaults.yaml")

    def _update_specific(self, config):
        self.writer = registry.get("writer")
        # 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
