#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
from collections import namedtuple
from typing import Iterable, Union

import megengine as mge
import megengine.distributed as dist
import megengine.module as M
import megengine.optimizer as optim
from basecore.config import ConfigDict
from megengine import Parameter
from megengine.amp import GradScaler
from megengine.autodiff import GradManager
from pkg_resources import packaging

from basecls.utils import registers

from .optimizer import LAMB, LARS, SGD
from .weight_decay import get_param_groups

__all__ = ["Solver", "BaseSolver", "DefaultSolver"]

Solver = namedtuple("Solver", ["optimizer", "grad_manager", "grad_scaler"])


class BaseSolver:
    """Base class for solver factory.

    A solver factory should return a :py:class:`~Solver` object, which combines
    an :py:class:`~megengine.optimizer.Optimizer` and
    a :py:class:`~megengine.autodiff.GradManager`.
    """

    @classmethod
    def build(cls, cfg: ConfigDict, model: M.Module) -> Solver:
        """Abstract build function

        Args:
            cfg: config for training.
            model: model for training.

        Returns:
            A solver.

        """
        raise NotImplementedError


@registers.solvers.register()
class DefaultSolver(BaseSolver):
    """The default solver factory.

    According to ``cfg.reduce_mode``, learning rate and weight decay will be scaled automatically
    following the linear scaling rule, see
    `"Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour"
    <https://arxiv.org/abs/1706.02677>`_ for more details.

    It supports ``"sgd"``, ``"adam"`` and ``"adamw"``.

    Note:
        This linear scaling rule can only work well with SGD. We are still looking for
        the applicable scaling rule for Adam and AdamW. Thus we recommend keeping default
        training settings (like learning rate and world size) when using Adam and AdamW.
    """

    @classmethod
    def build(cls, cfg: ConfigDict, model: M.Module) -> Solver:
        """Build function with the linear scaling strategy.

        Args:
            cfg: config for training.
            model: model for training.

        Returns:
            A solver.
        """
        amp_cfg = cfg.amp
        cfg = cfg.solver

        world_size = dist.get_world_size()

        # build optimizer
        lr = cfg.basic_lr * world_size  # linear scaling rule

        optim_params = get_param_groups(model, cfg.weight_decay)

        optimizer = cls.build_optimizer(cfg, optim_params, lr, 0)

        # build grad_manager
        gm = 
GradManager()