# Copyright (c) OpenMMLab. All rights reserved.

import copy as cp

import os

import os.path as osp

import time

import numpy as np

import torch

import torch.distributed as dist

from mmcv.parallel import MMDataParallel, MMDistributedDataParallel

from mmcv.runner import (DistSamplerSeedHook, EpochBasedRunner, OptimizerHook,

                         build_optimizer, get_dist_info)

from mmcv.runner.hooks import Fp16OptimizerHook

from ..core import (DistEvalHook, EvalHook, OmniSourceDistSamplerSeedHook,

                    OmniSourceRunner)

from ..datasets import build_dataloader, build_dataset

from ..utils import PreciseBNHook, get_root_logger

from .test import multi_gpu_test

def init_random_seed(seed=None, device='cuda'):

    """Initialize random seed.

    If the seed is not set, the seed will be automatically randomized,

    and then broadcast to all processes to prevent some potential bugs.

    Args:

        seed (int, Optional): The seed. Default to None.

        device (str): The device where the seed will be put on.

            Default to 'cuda'.

    Returns:

        int: Seed to be used.

    """

    if seed is not None:

        return seed

    # Make sure all ranks share the same random seed to prevent

    # some potential bugs. Please refer to

    # https://github.com/open-mmlab/mmdetection/issues/6339

    rank, world_size = get_dist_info()

    seed = np.random.randint(2**31)

    if world_size == 1:

        return seed

    if rank == 0:

        random_num = torch.tensor(seed, dtype=torch.int32, device=device)

    else:

        random_num = torch.tensor(0, dtype=torch.int32, device=device)

    dist.broadcast(random_num, src=0)

    return random_num.item()

def train_model(model,

                dataset,

                cfg,

                distributed=False,

                validate=False,

                test=dict(test_best=False, test_last=False),

                timestamp=None,

                meta=None):

    """Train model entry function.

    Args:

        model (nn.Module): The model to be trained.

        dataset (:obj:`Dataset`): Train dataset.

        cfg (dict): The config dict for training.

        distributed (bool): Whether to use distributed training.

            Default: False.

        validate (bool): Whether to do evaluation. Default: False.

        test (dict): The testing option, with two keys: test_last & test_best.

            The value is True or False, indicating whether to test the

            corresponding checkpoint.

            Default: dict(test_best=False, test_last=False).

        timestamp (str | None): Local time for runner. Default: None.

        meta (dict | None): Meta dict to record some important information.

            Default: None

    """

    logger = get_root_logger(log_level=cfg.log_level)

    # prepare data loaders

    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]

    dataloader_setting = dict(

        videos_per_gpu=cfg.data.get('videos_per_gpu', 1),

        workers_per_gpu=cfg.data.get('workers_per_gpu', 1),

        persistent_workers=cfg.data.get('persistent_workers', False),

        num_gpus=len(cfg.gpu_ids),

        dist=distributed,

        seed=cfg.seed)

    dataloader_setting = dict(dataloader_setting,

                              **cfg.data.get('train_dataloader', {}))

    if cfg.omnisource:

        # The option can override videos_per_gpu

        train_ratio = cfg.data.get('train_ratio', [1] * len(dataset))

        omni_videos_per_gpu = cfg.data.get('omni_videos_per_gpu', None)

        if omni_videos_per_gpu is None:

            dataloader_settings = [dataloader_setting] * len(dataset)

        else:

            dataloader_settings = []

            for videos_per_gpu in omni_videos_per_gpu:

                this_setting = cp.deepcopy(dataloader_setting)

                this_setting['videos_per_gpu'] = videos_per_gpu

                dataloader_settings.append(this_setting)

        data_loaders = [

            build_dataloader(ds, **setting)

            for ds, setting in zip(dataset, dataloader_settings)

        ]

    else:

        data_loaders = [

            build_dataloader(ds, **dataloader_setting) for ds in dataset

        ]

    # put model on gpus

    if distributed:

        find_unused_parameters = cfg.get('find_unused_parameters', False)

        # Sets the `find_unused_parameters` parameter in

        # torch.nn.parallel.DistributedDataParallel

        model = MMDistributedDataParallel(

            model.cuda(),

            device_ids=[torch.cuda.current_device()],

            broadcast_buffers=False,

            find_unused_parameters=find_unused_parameters)

    else:

        model = MMDataParallel(

            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)

    # build runner

    optimizer = build_optimizer(model, cfg.optimizer)

    Runner = OmniSourceRunner if cfg.omnisource else EpochBasedRunner

    runner = Runner(

        model,

        optimizer=optimizer,

        work_dir=cfg.work_dir,

        logger=logger,

        meta=meta)

    # an ugly workaround to make .log and .log.json filenames the same

    runner.timestamp = timestamp

    # fp16 setting

    fp16_cfg = cfg.get('fp16', None)

    if fp16_cfg is not None:

        optimizer_config = Fp16OptimizerHook(

            **cfg.optimizer_config, **fp16_cfg, distributed=distributed)

    elif distributed and 'type' not in cfg.optimizer_config:

        optimizer_config = OptimizerHook(**cfg.optimizer_config)

    else:

        optimizer_config = cfg.optimizer_config

    # register hooks

    runner.register_training_hooks(cfg.lr_config, optimizer_config,

                                   cfg.checkpoint_config, cfg.log_config,

                                   cfg.get('momentum_config', None))

    if distributed:

        if cfg.omnisource:

            runner.register_hook(OmniSourceDistSamplerSeedHook())

        else:

            runner.register_hook(DistSamplerSeedHook())

    # precise bn setting

    if cfg.get('precise_bn', False):

        precise_bn_dataset = build_dataset(cfg.data.train)

        dataloader_setting = dict(

            videos_per_gpu=cfg.data.get('videos_per_gpu', 1),

            workers_per_gpu=1,  # save memory and time

            persistent_workers=cfg.data.get('persistent_workers', False),

            num_gpus=len(cfg.gpu_ids),

            dist=distributed,

            seed=cfg.seed)

        data_loader_precise_bn = build_dataloader(precise_bn_dataset,

                                                  **dataloader_setting)

        precise_bn_hook = PreciseBNHook(data_loader_precise_bn,

                                        **cfg.get('precise_bn'))

        runner.register_hook(precise_bn_hook)

    if validate:

        eval_cfg = cfg.get('evaluation', {})

        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))

        dataloader_setting = dict(

            videos_per_gpu=cfg.data.get('videos_per_gpu', 1),

            workers_per_gpu=cfg.data.get('workers_per_gpu', 1),

            persistent_workers=cfg.data.get('persistent_workers', False),

            # cfg.gpus will be ignored if distributed

            num_gpus=len(cfg.gpu_ids),

            dist=distributed,

            shuffle=False)

        dataloader_setting = dict(dataloader_setting,

                                  **cfg.data.get('val_dataloader', {}))

        val_dataloader = build_dataloader(val_dataset, **dataloader_setting)

        eval_hook = DistEvalHook(val_dataloader, **eval_cfg) if distributed \

            else EvalHook(val_dataloader, **eval_cfg)

        runner.register_hook(eval_hook)

    if cfg.resume_from:

        runner.resume(cfg.resume_from)

    elif cfg.load_from:

        runner.load_checkpoint(cfg.load_from)

    runner_kwargs = dict()

    if cfg.omnisource:

        runner_kwargs = dict(train_ratio=train_ratio)

    runner.run(data_loaders, cfg.workflow, cfg.total_epochs, **runner_kwargs)

    dist.barrier()

    time.sleep(5)

    if test['test_last'] or test['test_best']:

        best_ckpt_path = None

        if test['test_best']:

            ckpt_paths = [x for x in os.listdir(cfg.work_dir) if 'best' in x]

            ckpt_paths = [x for x in ckpt_paths if x.endswith('.pth')]

            if len(ckpt_paths) == 0:

                runner.logger.info('Warning: test_best set, but no ckpt found')

                test['test_best'] = False

                if not test['test_last']:

                    return

            elif len(ckpt_paths) > 1:

                epoch_ids = [

                    int(x.split('epoch_')[-1][:-4]) for x in ckpt_paths

                ]

                best_ckpt_path = ckpt_paths[np.argmax(epoch_ids)]

            else:

                best_ckpt_path = ckpt_paths[0]

            if best_ckpt_path:

                best_ckpt_path = osp.join(cfg.work_dir, best_ckpt_path)

        test_dataset = build_dataset(cfg.data.test, dict(test_mode=True))

        gpu_collect = cfg.get('evaluation', {}).get('gpu_collect', False)

        tmpdir = cfg.get('evaluation', {}).get('tmpdir',

                                               osp.join(cfg.work_dir, 'tmp'))

        dataloader_setting = dict(

            videos_per_gpu=cfg.data.get('videos_per_gpu', 1),

            workers_per_gpu=cfg.data.get('workers_per_gpu', 1),

            persistent_workers=cfg.data.get('persistent_workers', False),

            num_gpus=len(cfg.gpu_ids),

            dist=distributed,

            shuffle=False)

        dataloader_setting = dict(dataloader_setting,

                                  **cfg.data.get('test_dataloader', {}))

        test_dataloader = build_dataloader(test_dataset, **dataloader_setting)

        names, ckpts = [], []

        if test['test_last']:

            names.append('last')

            ckpts.append(None)

        if test['test_best'] and best_ckpt_path is not None:

            names.append('best')

            ckpts.append(best_ckpt_path)

        for name, ckpt in zip(names, ckpts):

            if ckpt is not None:

                runner.load_checkpoint(ckpt)

            outputs = multi_gpu_test(runner.model, test_dataloader, tmpdir,

                                     gpu_collect)

            rank, _ = get_dist_info()

            if rank == 0:

                out = osp.join(cfg.work_dir, f'{name}_pred.pkl')

                test_dataset.dump_results(outputs, out)

                eval_cfg = cfg.get('evaluation', {})

                for key in [

                        'interval', 'tmpdir', 'start', 'gpu_collect',

                        'save_best', 'rule', 'by_epoch', 'broadcast_bn_buffers'

                ]:

                    eval_cfg.pop(key, None)

                eval_res = test_dataset.evaluate(outputs, **eval_cfg)

                runner.logger.info(f'Testing results of the {name} checkpoint')

                for metric_name, val in eval_res.items():

                    runner.logger.info(f'{metric_name}: {val:.04f}')