Optimizers

def init_training(self):
    """Sets up training and misceallenous parameters so the model is ready for training
    """
    self._model.init_training()

def load_state(self, optim_state: Dict, load_from_state: bool = False) -> Union['BaseOptimizer', None]:
    """Reconfigures optimizer from a given state.

    This is the default method for optimizers that don't support state. Does nothing.

    Args:
        optim_state: not used
        load_from_state (optional): not used

    Returns:
        None
    """
    logger.warning("load_state method of optimizer not implemented, cannot load optimizer status")
    return None

def save_state(self) -> Union[Dict, None]:
    """Gets optimizer state.

    This is the default method for optimizers that don't support state. Does nothing.

    Returns:
        None
    """
    logger.warning("save_state method of optimizer not implemented, cannot save optimizer status")
    return None

@abstractmethod
def step(self):
    """Performs an optimisation step and updates model weights.
    """

def get_aux(self) -> Optional[Dict[str, Any]]:
    aux = self.optimizer.get_aux()
    return aux

def load_state(self, optim_state: Dict[str, Any], load_from_state: bool = False) -> 'DeclearnOptimizer':
    """Reconfigures optimizer from a given state (contained in `optim_state` argument).
    Usage:
    ```python
    >>> import torch.nn as nn
    >>> from fedbiomed.common.optimizers import Optimizer
    >>> from fedbiomed.common.models import TorchModel
    >>> model = TorchModel(nn.Linear(4, 2))
    >>> optimizer = Optimizer(lr=.1)
    >>> optim = DeclearnOptimizer(model, optimizer)

    >>> optim.load_state(state)  # provided state contains the state one wants to load the optimizer with
    ```
    If `load_from_state` argument is True, it completes the current optimizer state with `optim_state` argument

    ```python
    >>> import torch.nn as nn
    >>> from fedbiomed.common.optimizers import Optimizer
    >>> from fedbiomed.common.optimizers.declearn import MomentumModule, AdamModule
    >>> from fedbiomed.common.models import TorchModel
    >>> model = TorchModel(nn.Linear(4, 2))
    >>> optimizer = Optimizer(lr=.1, modules=[MomentumModule(), AdamModule()])
    >>> optim_1 = DeclearnOptimizer(model, optimizer)

    >>> optimizer = Optimizer(lr=.1, modules=[AdamModule(), MomentumModule()])
    >>> optim_2 = DeclearnOptimizer(model, optimizer)
    >>> optim_2.load_state(optim_1.save_state())
    >>> optim_2.save_state()['states']
    {'modules': [('momentum', {'velocity': 0.0}),
            ('adam',
            {'steps': 0,
                'vmax': None,
                'momentum': {'state': 0.0},
                'velocity': {'state': 0.0}})]}
    ```
    Modules of DeclearnOptimizer will be reloaded provided that Module is the same and occupying the same index.
    Eg if the state contains following modules: 
    ```modules=[AdamModule(), AdagradModule(), MomemtumModule()]```
     And the Optimizer contained in the TrainingPlan has the following modules:
    ```modules=[AdamModule(), MomemtumModule()]```
    Then only `AdamModule` module will be reloaded, `MomentumModule` will be set with default argument (they don't
    share the same index in the modules list).

    Args:
        optim_state: state of the Optimizer to be loaded. It will change the current state of the optimizer
            with the one loaded
        load_from_state (optional): strategy for loading states: whether to load from saved states (True) or
            from breakpoint (False).
            If set to True, loading is done partially in the sense that if some of the OptimModules is different in
            the optim_state and the original state of the optimizer, it loads only the OptiModule(s) from the
            latest state that both state has in common. Defaults to False.

    Raises:
        FedbiomedOptimizerError: raised if state is not of dict type.

    Returns:
        Optimizer wrapper reloaded from `optim_state` argument.
    """
    # state: breakpoint content for optimizer
    if not isinstance(optim_state, Dict):
        raise FedbiomedOptimizerError(f"{ErrorNumbers.FB626.value}, incorrect type of argument `optim_state`: "
                                      f"expecting a dict, but got {type(optim_state)}")

    if load_from_state:
        # first get Optimizer detailed in the TrainingPlan.

        init_optim_state = self.optimizer.get_state()  # we have to get states since it is the only way we can 
        # gather modules (other methods of `Optimizer are private`)

        optim_state_copy = copy.deepcopy(optim_state)
        optim_state.update(init_optim_state)  # optim_state will be updated with current optimizer state
        # check if opimizer state has changed from last optimizer to the current one
        # if it has changed, find common modules and update common states
        for component in ( 'modules', 'regularizers',):
            components_to_keep: List[Tuple[str, int]] = []  # we store here common Module between current Optimizer
            # and the ones in the `optim_state` tuple (common Module name, index in List)

            if not init_optim_state['states'].get(component) or not optim_state_copy['states'].get(component):
                continue
            self._collect_common_optimodules(
                init_optim_state,
                optim_state_copy,
                component,
                components_to_keep
            )

            for mod in components_to_keep:
                for mod_state in optim_state_copy['states'][component]:
                    if mod[0] == mod_state[0]:
                        # if we do find same module in the current optimizer than the previous one,
                        # we load the previous optimizer module state into the current one
                        optim_state['states'][component][mod[1]] = mod_state

        logger.info("Loading optimizer state from saved state")

    reloaded_optim = FedOptimizer.load_state(optim_state)
    self.optimizer = reloaded_optim
    return self

def optimizer_processing(self) -> SklearnOptimizerProcessing:
    """Provides a context manager able to do some actions before and after setting up an Optimizer, mainly
    disabling scikit-learn internal optimizer.

    Also, checks if `model_args` dictionary contains training parameters that
    won't be used or have any effect on the training, because of disabling the scikit-learn optimizer (
    such as initial learning rate, learnig rate scheduler, ...). If disabling the internal optimizer leads
    to such changes, displays a warning.

    Returns:
        SklearnOptimizerProcessing: context manager providing extra logic

    Usage:
    ```python
        >>> dlo = DeclearnSklearnOptimizer(model, optimizer)
        >>> with dlo.optimizer_processing():
                model.train(inputs,targets)
    ```
    """
    if isinstance(self._model, SkLearnModel):
        return SklearnOptimizerProcessing(self._model, disable_internal_optimizer=True)
    else:
        raise FedbiomedOptimizerError(f"{ErrorNumbers.FB626.value}: Method optimizer_processing should be used "
                                      f"only with SkLearnModel, but model is {self._model}")

def save_state(self) -> Dict:
    """Gets optimizer state.

    Returns:
        optimizer state
    """
    optim_state = self.optimizer.get_state()
    return optim_state

def set_aux(self, aux: Dict[str, Any]):
    # FIXME: for imported tensors in PyTorch sent as auxiliary variables,
    # we should push it on the appropriate device (ie cpu/gpu)
    self.optimizer.set_aux(aux)

def step(self):
    """Performs one optimization step"""
    # NOTA: for sklearn, gradients retrieved are unscaled because we are using learning rate equal to 1.
    # Therefore, it is necessary to disable the sklearn internal optimizer beforehand
    # otherwise, computation will be incorrect
    grad = declearn.model.api.Vector.build(self._model.get_gradients())
    weights = declearn.model.api.Vector.build(self._model.get_weights(
        only_trainable=False,
        exclude_buffers=True
    ))
    updates = self.optimizer.step(grad, weights)
    self._model.apply_updates(updates.coefs)

def zero_grad(self):
    """Zeroes gradients of the Pytorch model. Basically calls the `zero_grad`
    method of the model.

    Raises:
        FedbiomedOptimizerError: triggered if model has no method called `zero_grad`
    """
    # warning: specific for pytorch
    if not isinstance(self._model, TorchModel):
        raise FedbiomedOptimizerError(f"{ErrorNumbers.FB626.value}. This method can only be used for TorchModel, "
                                      f"but got {self._model}")
    self._model.model.zero_grad()

def optimizer_processing(self) -> SklearnOptimizerProcessing:
    return SklearnOptimizerProcessing(self._model, disable_internal_optimizer=False)

def step(self):
    """Performs an optimization step and updates model weights."""
    gradients = self._model.get_gradients()
    updates = {k: -v for k, v in gradients.items()}
    self._model.apply_updates(updates)

def get_learning_rate(self) -> Dict[str, float]:
    """Gets learning rates from param groups in Pytorch optimizer.

    For each optimizer param group, it iterates over all parameters in that parameter group and searches for the "
    corresponding parameter of the model by iterating over all model parameters. If it finds a correspondence,
    it saves the learning rate value. This function assumes that the parameters in the optimizer and the model
    have the same reference.


    !!! warning
        This function gathers the base learning rate applied to the model weights,
        including alterations due to any LR scheduler. However, it does not catch
        any adaptive component, e.g. due to RMSProp, Adam or such.

    Returns:
        List[float]: list of single learning rate or multiple learning rates
            (as many as the number of the layers contained in the model)
    """
    logger.warning(
        "`get_learning_rate` is deprecated and will be removed in future Fed-BioMed releases",
        broadcast=True)

    mapping_lr_layer_name: Dict[str, float] = {}

    for param_group in self.optimizer.param_groups:
        for layer_params in param_group['params']:
            for layer_name, tensor in self._model.model.named_parameters():
                if layer_params is tensor:
                    mapping_lr_layer_name[layer_name] = param_group['lr']
    return mapping_lr_layer_name

def step(self):
    """Performs an optimization step and updates model weights
    """
    self.optimizer.step()

def zero_grad(self):
    """Zeroes gradients of the Pytorch model. Basically calls the `zero_grad`
    method of the optimizer.
    """
    self.optimizer.zero_grad()

@classmethod
def from_declearn_optimizer(
    cls,
    declearn_optimizer: DeclearnOptimizer,
) -> Self:
    """Wrap a declearn Optimizer into a fed-biomed one.

    Args:
        declearn_optimizer: [declearn.optimizer.Optimizer][] instance that
            needs to be wrapped.

    Returns:
        Fed-BioMed `Optimizer` instance wrapping a copy of the input
            declearn optimizer.
    """
    config = declearn_optimizer.get_config()
    optim = cls(
        lr=config["lrate"],
        decay=config["w_decay"],
        modules=config["modules"],
        regularizers=config["regularizers"],
    )
    optim._optimizer.set_state(declearn_optimizer.get_state())
    return optim

def get_aux(self) -> Dict[str, Union[Dict[str, Any], Any]]:
    """Return auxiliary variables that need to be shared across network.

    Returns:
        Aux-var dict that associates `module.collect_aux_var()` values to
            `module.name` keys for each and every module plugged in this
            Optimizer that has some auxiliary variables to share.

    !!! info "Note"
        "Auxiliary variables" are information that needs to be shared
        between the nodes and the researcher between training rounds, to
        synchronize some optimizer plug-ins that work by pair. Their
        production via this method can have internal side effects;
        `get_aux` should therefore be called sparingly.
    """
    try:
        return self._optimizer.collect_aux_var()
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: error in 'get_aux': {exc}"
        ) from exc

def get_state(self) -> Dict[str, Any]:
    """Return the configuration and current states of this Optimizer.

    This method is to be used for creating breakpoints.

    Returns:
        State-and-config dict that may be saved as part of a breakpoint
            file, and used to re-create this Optimizer using the
            `Optimizer.load_state` classmethod constructor.
    """
    try:
        config = self._optimizer.get_config()
        states = self._optimizer.get_state()
        return {"config": config, "states": states}
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: error in 'get_state': {exc}"
        ) from exc

def init_round(self) -> None:
    """Trigger start-of-training-round behavior of wrapped regularizers."""
    try:
        self._optimizer.start_round()
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: error in 'init_round': {exc}"
        ) from exc

@classmethod
def load_state(cls, state: Dict[str, Any]) -> Self:
    """Instantiate an Optimizer from its breakpoint state dict.

    Args:
        state: state-and-config dict created using the `get_state` method.

    Returns:
        Optimizer instance re-created from the `state` dict.

    Raises:
        FedbiomedOptimizerError: If the input `state` dict has improper keys
            or fails to set up a declearn Optimizer and set back its state.
    """
    try:
        optim = DeclearnOptimizer.from_config(state["config"])
        optim.set_state(state["states"])
    except KeyError as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: Missing field in the breakpoints state: {exc}"
        ) from exc
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: `Optimizer.load_state`: {exc}"
        ) from exc
    return cls(
        lr=optim.lrate,
        decay=optim.w_decay,
        modules=optim.modules,
        regularizers=optim.regularizers,
    )

def set_aux(self, aux: Dict[str, Dict[str, Any]]) -> None:
    """Update plug-in modules based on received shared auxiliary variables.

    Args:
        aux: Auxiliary variables received from the counterpart optimizer
            (on the other side of the node-researcher frontier), that are
            to be a `{module.name: module.collect_aux_var()}` *or* a
            `{module.name: {node: module.collect_aux_var()}}` dict
            (depending on which side this optimizer is on).

    Raises:
        FedbiomedOptimizerError: If a key from `aux_var` does not match the
            name of any module plugged in this optimizer (i.e. if received
            variables cannot be mapped to a destinatory module).

    !!! info "Note"
        "Auxiliary variables" are information that is shared between the
        nodes and researcher between training rounds, to synchronize some
        optimizer plug-ins that work by pair. The inputs to this method are
        not simply stored by the Optimizer, but are processed into internal
        side effects; this method should therefore be called sparingly.
    """
    try:
        self._optimizer.process_aux_var(aux)
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: `Optimizer.set_aux`: {exc}"
        ) from exc

def step(self, grads: Vector, weights: Vector) -> Vector:
    """Run an optimization step to compute and return model weight updates.

    Use the pre-assigned `weights` and `grads` (set using the `set_weights`
    and `set_grads` methods) to compute weight updates, using the pipeline
    defined by this instance.

    Args:
        grads: Raw gradients based on which to compute weights updates,
            wrapped into a declearn Vector structure.
        weights: Current values of the weights with respect to which the
            gradients were computed, wrapped into a declearn Vector with
            the same concrete type as `grads`.

    Returns:
        Updates to be applied to the model weights, computed by:
            - running wrapped gradients and weights through the regularizer
              plug-ins (that add loss-regularization terms' derivatives);
            - running resulting gradients through the optimodule plug-ins
              (that perform any defined gradient-alteration operation);
            - adding a decoupled weight-decay term, if one is to be used;
            - scaling the updates by the base learning rate.
            The results are wrapped into a declearn Vector structure, the
            concrete type of which is same as input `grads` and `weights`.
    """
    # This code mostly replicates that of
    # `declearn.optimizer.Optimizer.compute_updates_from_gradients`.
    try:
        # Add loss-regularization terms' derivatives to the raw gradients.
        for reg in self._optimizer.regularizers:
            grads = reg.run(grads, weights)
        # Iteratively refine updates by running them through the optimodules.
        for mod in self._optimizer.modules:
            grads = mod.run(grads)
        # Apply the base learning rate.
        updates = - self._optimizer.lrate * grads
        # Optionally add the decoupled weight decay term.
        if self._optimizer.w_decay:
            updates -= self._optimizer.w_decay * weights
        # Return the model updates.
        return updates
    except Exception as exc:
        raise FedbiomedOptimizerError(
            f"{ErrorNumbers.FB621.value}: error in 'step': {exc}"
        ) from exc