# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import os
import omegaconf
from omegaconf import OmegaConf


def load_config(args=None, config_file=None, overwrite_fairseq=False):
    """TODO (huxu): move fairseq overwrite to another function."""
    if args is not None:
        config_file = args.taskconfig
    config = recursive_config(config_file)

    if config.dataset.subsampling is not None:
        batch_size = config.fairseq.dataset.batch_size // config.dataset.subsampling
        print(
            "adjusting batch_size to {} due to subsampling {}.".format(
                batch_size, config.dataset.subsampling
            )
        )
        config.fairseq.dataset.batch_size = batch_size

    is_test = config.dataset.split is not None and config.dataset.split == "test"
    if not is_test:
        if (
            config.fairseq.checkpoint is None
            or config.fairseq.checkpoint.save_dir is None
        ):
            raise ValueError("fairseq save_dir or save_path must be specified.")

        save_dir = config.fairseq.checkpoint.save_dir
        os.makedirs(save_dir, exist_ok=True)
        if config.fairseq.common.tensorboard_logdir is not None:
            tb_run_dir = suffix_rundir(
                save_dir, config.fairseq.common.tensorboard_logdir
            )
            config.fairseq.common.tensorboard_logdir = tb_run_dir
            print(
                "update tensorboard_logdir as", config.fairseq.common.tensorboard_logdir
            )
        os.makedirs(save_dir, exist_ok=True)
        OmegaConf.save(config=config, f=os.path.join(save_dir, "config.yaml"))

    if overwrite_fairseq and config.fairseq is not None and args is not None:
        # flatten fields.
        for group in config.fairseq:
            for field in config.fairseq[group]:
                print("overwrite args." + field, "as", config.fairseq[group][field])
                setattr(args, field, config.fairseq[group][field])
    return config


def recursive_config(config_path):
    """allows for stacking of configs in any depth."""
    config = OmegaConf.load(config_path)
    if config.includes is not None:
        includes = config.includes
        config.pop("includes")
        base_config = recursive_config(includes)
        config = OmegaConf.merge(base_config, config)
    return config


def suffix_rundir(save_dir, run_dir):
    max_id = -1
    for search_dir in os.listdir(save_dir):
        if search_dir.startswith(run_dir):
            splits = search_dir.split("_")
            cur_id = int(splits[1]) if len(splits) > 1 else 0
            max_id = max(max_id, cur_id)
    return os.path.join(save_dir, run_dir + "_" + str(max_id + 1))


def overwrite_dir(config, replace, basedir):
    for key in config:
        if isinstance(config[key], str) and config[key].startswith(basedir):
            config[key] = config[key].replace(basedir, replace)
        if isinstance(config[key], omegaconf.dictconfig.DictConfig):
            overwrite_dir(config[key], replace, basedir)