from pathlib import Path
from importlib import import_module

import torch

import ignite
import ignite.distributed as idist
from ignite.utils import manual_seed
from util.misc import setup_logger

import fire
from omegaconf import OmegaConf


def log_basic_info(logger, config):
    logger.info(f"Train {config.name}")
    logger.info(f"- PyTorch version: {torch.__version__}")
    logger.info(f"- Ignite version: {ignite.__version__}")
    if idist.get_world_size() > 1:
        logger.info("Distributed setting:\n")
        idist.show_config()


def running(local_rank, config, task, backup_config=False, setup_output_dir=False, setup_random_seed=False):
    if setup_random_seed:
        manual_seed(config.misc.random_seed + idist.get_rank())
    output_dir = Path(config.result_dir) / config.name if config.output_dir is None else config.output_dir
    config.output_dir = str(output_dir)

    if setup_output_dir and config.resume_from is None:
        if output_dir.exists():
            # assert not any(output_dir.iterdir()), "output_dir must be empty"
            contains = list(output_dir.iterdir())
            assert (len(contains) == 0) or (len(contains) == 1 and contains[0].name == "config.yml"), \
                f"output_dir must by empty or only contains config.yml, but now got {len(contains)} files"
        else:
            if idist.get_rank() == 0:
                output_dir.mkdir(parents=True)
                print(f"mkdir -p {output_dir}")

        if backup_config and idist.get_rank() == 0:
            with open(output_dir / "config.yml", "w+") as f:
                print(config.pretty(), file=f)
    logger = setup_logger(name=config.name, distributed_rank=local_rank, filepath=output_dir / "train.log")
    logger.info(f"output path: {config.output_dir}")
    log_basic_info(logger, config)

    OmegaConf.set_readonly(config, True)

    engine = import_module(f"engine.{config.engine}")
    engine.run(task, config, logger)


def run(task, config: str, *omega_options, **kwargs):
    omega_options = [str(o) for o in omega_options]
    cli_conf = OmegaConf.from_cli(omega_options)
    if len(cli_conf) > 0:
        print(cli_conf.pretty())
    conf = OmegaConf.merge(OmegaConf.load(config), cli_conf)
    backend = kwargs.get("backend", "nccl")
    backup_config = kwargs.get("backup_config", False)
    setup_output_dir = kwargs.get("setup_output_dir", False)
    setup_random_seed = kwargs.get("setup_random_seed", False)
    with idist.Parallel(backend=backend) as parallel:
        parallel.run(running, conf, task, backup_config=backup_config, setup_output_dir=setup_output_dir,
                     setup_random_seed=setup_random_seed)


if __name__ == '__main__':
    fire.Fire(run)