import os.path as osp
from abc import ABCMeta, abstractmethod
import megengine as mge
import megengine.distributed as dist
from megengine.optimizer.optimizer import Optimizer
from megengine.module import Module
from edit.utils import mkdir_or_exist, build_from_cfg, get_root_logger
from ..hook import Hook, HOOKS, get_priority

module_ckpt_suffix = "_module.mge"
optim_ckpt_suffix = "_optim.mge"

class BaseRunner(metaclass=ABCMeta):
    """The base class of Runner, a training helper for Mge.

    All subclasses should implement the following APIs:

    - ``run()``
    - ``train()``
    - ``test()``
    - ``save_checkpoint()``
    - ``resume()``

    Args:
        model (:obj:`megengine.module.Module`): The model to be run.
        optimizers_cfg (dict): optimizer configs
        work_dir (str, optional): The working directory to save checkpoints and logs. Defaults to None.
    """

    def __init__(self, model, optimizers_cfg=None, work_dir=None):
        assert hasattr(model, 'train_step')
        assert hasattr(model, 'test_step')
        assert hasattr(model, 'create_gradmanager_and_optimizers')
        assert hasattr(model, 'cal_for_eval')

        self.model = model
        self.optimizers_cfg = optimizers_cfg
        self.logger = get_root_logger()
        self.work_dir = work_dir
        assert self.work_dir is not None

        # get model name from the model class
        self._model_name = self.model.__class__.__name__
        self.mode = None
        self._hooks = []
        self._epoch = 0
        self._iter = 0
        self._inner_iter = 0
        self._max_epochs = 0
        self._max_iters = 0

    @property
    def model_name(self):
        """str: Name of the model, usually the module class name."""
        return self._model_name

    @property
    def hooks(self):
        """list[:obj:`Hook`]: A list of registered hooks."""
        return self._hooks

    @property
    def epoch(self):
        """int: Current epoch."""
        return self._epoch

    @property
    def iter(self):
        """int: Current iteration."""
        return self._iter

    @property
    def inner_iter(self):
        """int: Iteration in an epoch."""
        return self._inner_iter

    @property
    def max_epochs(self):
        """int: Maximum training epochs."""
        return self._max_epochs

    @property
    def max_iters(self):
        """int: Maximum training iterations."""
        return self._max_iters

    @abstractmethod
    def train(self, data_loader):
        pass

    @abstractmethod
    def test(self, data_loader):
        pass

    @abstractmethod
    def run(self, data_loaders, workflow, max_iters):
        pass

    @abstractmethod
    def save_checkpoint(self, out_dir, create_symlink=True):
        pass

    @abstractmethod
    def resume(self, path2checkpoint):
        pass

    @abstractmethod
    def register_training_hooks(self, lr_config, checkpoint_config, log_config):
        """Register default hooks for training.

            Default hooks include:

            - LrUpdaterHook
            - CheckpointSaverHook
            - log_config
        """
        pass

    def create_gradmanager_and_optimizers(self):
        self.model.create_gradmanager_and_optimizers(self.optimizers_cfg)

    def sync_model_params(self):
        if 
dist.is_distributed()