# Copyright (c) Alibaba, Inc. and its affiliates.

from typing import Dict, List, Optional, Tuple

from modelscope.utils.config import Config

DEFAULT_CONFIG = Config({
    'framework': 'pytorch',
    'train': {
        'work_dir': '/tmp',
        'max_epochs': 10,
        'dataloader': {
            'batch_size_per_gpu': 16,
            'workers_per_gpu': 0
        },
        'optimizer': {
            'type': 'SGD',
            'lr': 1e-3
        },
        'lr_scheduler': {
            'type': 'StepLR',
            'step_size': 2
        },
        'checkpoint': {
            'period': {
                'interval': 1
            }
        }
    },
    'evaluation': {
        'dataloader': {
            'batch_size_per_gpu': 16,
            'workers_per_gpu': 0,
            'shuffle': False
        },
    }
})

DEFAULT_HOOKS_CONFIG = {
    'train': {
        'hooks': [{
            'type': 'CheckpointHook',
            'interval': 1
        }, {
            'type': 'TextLoggerHook',
            'interval': 10
        }, {
            'type': 'IterTimerHook'
        }]
    }
}

_HOOK_KEY_CHAIN_MAP = {
    'TextLoggerHook': 'train.logging',
    'CheckpointHook': 'train.checkpoint.period',
    'BestCkptSaverHook': 'train.checkpoint.best',
    'EvaluationHook': 'evaluation.period',
}


def merge_cfg(cfg: Config):
    """Merge the default config into the input cfg.

    This function will pop the default CheckpointHook when the BestCkptSaverHook exists in the input cfg.

    Aegs:
        cfg: The input cfg to be merged into.
    """
    cfg.merge_from_dict(DEFAULT_HOOKS_CONFIG, force=False)


def merge_hooks(cfg: Config) -> List[Dict]:
    hooks = cfg.train.hooks.copy()
    for hook_type, key_chain in _HOOK_KEY_CHAIN_MAP.items():
        hook = _key_chain_to_hook(cfg, key_chain, hook_type)
        if hook is not None:
            hooks.append(hook)
    return hooks


def update_cfg(cfg: Config) -> Config:
    if 'hooks' not in cfg.train:
        return cfg
    key_chain_map = {}
    for hook in cfg.train.hooks:
        if not hook:
            continue
        key, value = _hook_split(hook)
        if key not in _HOOK_KEY_CHAIN_MAP:
            continue
        key_chain_map[_HOOK_KEY_CHAIN_MAP[key]] = value
        hook.clear()
    cfg.train.hooks = list(filter(bool, cfg.train.hooks))
    cfg.merge_from_dict(key_chain_map)
    return cfg


def _key_chain_to_hook(cfg: Config, key_chain: str,
                       hook_type: str) -> Optional[Dict]:
    if not _check_basic_hook(cfg, key_chain, hook_type):
        return None
    hook_params: Dict = cfg.safe_get(key_chain)
    hook = {'type': hook_type}
    hook.update(hook_params)
    return hook


def _check_basic_hook(cfg: Config, key_chain: str, hook_type: str) -> bool:
    if cfg.safe_get(key_chain) is None:
        return False
    hooks = list(
        filter(lambda hook: hook['type'] == hook_type, cfg.train.hooks))
    assert len(hooks) == 0, f'The key_chain {key_chain} and the traditional hook ' \
                            f'cannot exist at the same time, ' \
                            f'please delete {hook_type} in the configuration file.'
    return True


def _hook_split(hook: Dict) -> Tuple[str, Dict]:
    hook = hook.copy()
    return hook.pop('type'), hook
