fedbiomed.common.models

def apply_updates(self, updates: Union[Dict[str, np.ndarray], NumpyVector]) -> None:
    """Apply incoming updates to the wrapped model's parameters.

    Args:
        updates: Model parameters' updates to add/apply existing model parameters.
    """
    for key, val in self._get_iterator_model_params(updates):
        wgt = getattr(self.model, key)
        setattr(self.model, key, wgt + val)

@abstractmethod
def disable_internal_optimizer(self) -> None:
    """Abstract method to apply;

    Disables scikit learn internal optimizer by setting arbitrary learning rate parameters to the
    scikit learn model, in order to then compute its gradients.

    ''' warning "Call it only if using `declearn` optimizers"
            Method implementation will depend on the attribute used to set up
            these arbitrary arguments.
    """

def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

    !!! warning "Warning":
        This method uses `joblib.dump`, which relies on pickle and
        is therefore hard to trust by third-party loading methods.
    """
    with open(filename, "wb") as file:
        joblib.dump(self.model, file)

def flatten(self) -> List[float]:
    """Gets weights as flatten vector

    Returns:
        to_list: Convert np.ndarray to a list if it is True.
    """

    weights = self.get_weights()
    flatten = []
    for _, w in weights.items():
        w_: List[float] = list(w.flatten().astype(float))
        flatten.extend(w_)

    return flatten

def get_gradients(
    self,
    as_vector: bool = False,
) -> Union[Dict[str, np.ndarray], NumpyVector]:
    """Gets computed gradients

    Args:
        as_vector: Whether to wrap returned gradients into a declearn Vector.

    Raises:
        FedbiomedModelError: raised if gradients have not been computed yet (ie model has not been trained)

    Returns:
        Gradients, as a dictionary mapping parameters' names to their gradient's
            numpy array, or as a declearn NumpyVector wrapping such a dict.
    """
    if self._gradients is None:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Cannot get gradients if model has not been trained beforehand!"
        )
    gradients = self._gradients
    if as_vector:
        return NumpyVector(gradients)
    return gradients

@abstractmethod
def get_learning_rate(self) -> List[float]:
    """Retrieves learning rate of the model. Method implementation will
    depend on the attribute used to set up these arbitrary arguments

    Returns:
        Initial learning rate value(s); a single value if only on learning rate has been used, and
            a list of several learning rates, one for each layer of the model.
    """

def get_params(self, value: Any = None) -> Dict[str, Any]:
    """Gets scikit learn model hyperparameters.

    Please refer to [`baseEstimator documentation`]
    [https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html] `get_params` method
    for further details

    Args:
        value: if specified, returns a specific hyperparameter, otherwise, returns a dictionary
            with all the hyperparameters. Defaults to None.

    Returns:
        Dictionary mapping model hyperparameter names to their values
    """
    if value is not None:
        return self.model.get_params().get(value)
    return self.model.get_params()

def get_weights(
    self,
    as_vector: bool = False,
) -> Union[Dict[str, np.ndarray], NumpyVector]:
    """Returns model's parameters, optionally as a declearn NumpyVector.

    Args:
        as_vector: Whether to wrap returned weights into a declearn Vector.

    Raises:
        FedbiomedModelError: If the list of parameters are not defined.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their
            numpy array, or as a declearn NumpyVector wrapping such a dict.
    """
    if not self.param_list:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Attribute `param_list` is empty. You should "
            f"have initialized the model beforehand (try calling `set_init_params`)"
        )
    # Gather copies of the model weights.
    weights = {}  # type: Dict[str, np.ndarray]
    try:
        for key in self.param_list:
            val = getattr(self.model, key)
            if not isinstance(val, np.ndarray):
                raise FedbiomedModelError(
                    f"{ErrorNumbers.FB622.value}: SklearnModel parameter is not a numpy array."
                )
            weights[key] = val.copy()
    except AttributeError as err:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Unable to access weights of BaseEstimator "
            f"model {self.model} (details {err}"
        ) from err
    # Optionally encapsulate into a NumpyVector, else return as a dict.
    if as_vector:
        return NumpyVector(weights)
    return weights

def init_training(self):
    """Initialises the training by setting up attributes.

    Raises:
        FedbiomedModelError: raised if `param_list` has not been defined
    """
    if not self.param_list:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Attribute `param_list` is empty. You should "
            f"have initialized the model beforehand (try calling `set_init_params`)"
        )
    if self._is_declearn_optim:
        self.disable_internal_optimizer()

def predict(self, inputs: np.ndarray) -> np.ndarray:
    """Computes prediction given input data.

    Args:
        inputs: input data

    Returns:
        Model predictions
    """
    return self.model.predict(inputs)

@abstractmethod
def set_init_params(self, model_args: Dict) -> None:
    """Zeroes scikit learn model parameters.

    Should be used before any training, as it sets the scikit learn model parameters
    and makes them accessible through the use of attributes. Model parameter attribute names
    will depend on the scikit learn model wrapped.

    Args:
        model_args: dictionary that contains specifications for setting initial model
    """

def set_params(self, **params: Any) -> Dict[str, Any]:
    """Sets scikit learn model hyperparameters.

    Please refer to [BaseEstimator][sklearn.base.BaseEstimator]
    [https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html] `set_params` method
    for further details

    Args:
        params: new hyperparameters to set up the model.

    Returns:
        Dict[str, Any]: dictionary containing new hyperparameters values
    """
    self.model.set_params(**params)
    return params

def set_weights(
    self,
    weights: Union[Dict[str, np.ndarray], NumpyVector],
) -> None:
    """Assign new values to the model's trainable weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names to their
            numpy array, or as a declearn NumpyVector wrapping such a dict.
    """
    for key, val in self._get_iterator_model_params(weights):
        setattr(self.model, key, val.copy())

def train(
    self,
    inputs: np.ndarray,
    targets: np.ndarray,
    stdout: Optional[List[List[str]]] = None,
    **kwargs,
) -> None:
    """Trains scikit learn model and internally computes gradients

    Args:
        inputs: inputs data.
        targets: targets, to be fit with inputs data
        stdout: list of console outputs that have been collected
            during training, that contains losses values. Used to plot model losses. Defaults to None.

    Raises:
        FedbiomedModelError: raised if training has not been initialized
    """
    batch_size = inputs.shape[0]
    w_init = self.get_weights(as_vector=False)  # type: Dict[str, np.ndarray]
    w_updt = {key: np.zeros_like(val) for key, val in w_init.items()}
    # Iterate over the batch; accumulate sample-wise gradients (and loss).
    for idx in range(batch_size):
        # Compute updated weights based on the sample. Capture loss prints.
        with capture_stdout() as console:
            self.model.partial_fit(inputs[idx : idx + 1], targets[idx])
        if stdout is not None:
            stdout.append(console)
        # Accumulate updated weights (weights + sum of gradients).
        # Reset the model's weights and iteration counter.
        for key in self.param_list:
            w_updt[key] += getattr(self.model, key)
            setattr(self.model, key, w_init[key])
        self.model.n_iter_ -= 1
    # Compute the batch-averaged gradients (scaled by eta_t).
    # Note: w_init: {w_t}, w_updt: {w_t - eta_t sum_{s=1}^B(grad_s)}
    #       hence eta_t * avg(grad_s) = w_init - (w_updt / B)
    self._gradients = {
        key: w_init[key] - (w_updt[key] / batch_size) for key in self.param_list
    }
    # When using a declearn Optimizer, negate the learning rate.
    if self._is_declearn_optim:
        lrate = self.get_learning_rate()[0]
        for key, val in self._gradients.items():
            val /= lrate
    # Finally, increment the model's iteration counter.
    self.model.n_iter_ += 1

def unflatten(
        self,
        weights_vector: List[float]
) -> Dict[str, np.ndarray]:
    """Unflatten vectorized model weights

    Args:
        weights_vector: Vectorized model weights to convert dict

    Returns:
        Model dictionary
    """

    super().unflatten(weights_vector)

    weights_vector = np.array(weights_vector)
    weights = self.get_weights()
    pointer = 0

    params = {}
    for key, w in weights.items():
        num_param = w.size
        params[key] = weights_vector[pointer: pointer + num_param].reshape(w.shape)

        pointer += num_param

    return params

def disable_internal_optimizer(self):
    self.model.learning_rate_init = self.default_lr_init
    self.model.learning_rate = self.default_lr
    self._is_declearn_optim = True

def get_learning_rate(self) -> List[float]:
    return [self.model.learning_rate_init]

@abstractmethod
def apply_updates(self, updates: Any):
    """Applies updates to the model.

    Args:
        updates (Any): model updates.
    """

@abstractmethod
def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.
    """

@abstractmethod
def flatten(self) -> List[float]:
    """Flattens model weights

    Returns:
        List of model weights as float.
    """

@abstractmethod
def get_gradients(self, as_vector: bool = False) -> Union[Dict[str, Any], _VT]:
    """Return computed gradients attached to the model.

    Args:
        as_vector: Whether to wrap returned gradients into a declearn Vector.

    Returns:
        Gradients, as a dictionary mapping parameters' names to their gradient's
            value, or as a declearn Vector structure wrapping such a dict.
    """

@abstractmethod
def get_weights(self, as_vector: bool = False) -> Union[Dict[str, _DT], _VT]:
    """Return a copy of the model's trainable weights.

    Args:
        as_vector: Whether to wrap returned weights into a declearn Vector.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their
            value, or as a declearn Vector structure wrapping such a dict.
    """

@abstractmethod
def init_training(self):
    """Initializes parameters before model training"""

@abstractmethod
def predict(self, inputs: Any) -> Any:
    """Returns model predictions given input values

    Args:
        inputs (Any): input values.

    Returns:
        Any: predictions.
    """

def reload(self, filename: str) -> None:
    """Import and replace the wrapped model from a dump file.

    Args:
        filename: path to the file where the model has been exported.

    !!! info "Notes":
        This method is designed to load the model from a local dump
        file, that might not be in a trustworthy format. It should
        therefore only be used to re-load data exported locally and
        not received from someone else, including other FL peers.

    Raises:
        FedbiomedModelError: if the reloaded instance is of unproper type.
    """
    model = self._reload(filename)
    if not isinstance(model, self._model_type):
        err_msg = (
            f"{ErrorNumbers.FB622.value}: unproper type for imported model"
            f": expected '{self._model_type}', but 'got {type(model)}'."
        )
        logger.critical(err_msg)
        raise FedbiomedModelError(err_msg)
    self.model = model

@abstractmethod
def set_weights(self, weights: Union[Dict[str, _DT], _VT]) -> None:
    """Assign new values to the model's trainable weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names to their
            value, or as a declearn Vector structure wrapping such a dict.
    """

@abstractmethod
def train(self, inputs: Any, targets: Any, **kwargs) -> None:
    """Trains model given inputs and targets data

    !!! warning "Warning"
        Please run `init_training` method before running `train` method,
        so to initialize parameters needed for model training"

    !!! warning "Warning"
        This function may not update weights. You may need to call `apply_updates`
        to apply updates to the model

    Args:
        inputs (Any): input (training) data.
        targets (Any): target values.
    """

@abstractmethod
def unflatten(
        self,
        weights_vector: List[float]
) -> None:
    """Revert flatten model weights back model-dict form.

    Args:
        weights_vector: Vectorized model weights to convert dict

    Returns:
        Model dictionary
    """

    if not isinstance(weights_vector, list) or not all([isinstance(w, float) for w in weights_vector]):
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622} `weights_vector should be 1D list of float containing flatten model parameters`"
        )

def set_init_params(self, model_args: Dict[str, Any]) -> None:
    """Initialize the model's trainable parameters."""
    # Set up zero-valued start weights, for binary of multiclass classif.
    n_classes = model_args["n_classes"]
    if n_classes == 2:
        init_params = {
            "intercept_": np.zeros((1,)),
            "coef_": np.zeros((1, model_args["n_features"])),
        }
    else:
        init_params = {
            "intercept_": np.zeros((n_classes,)),
            "coef_": np.zeros((n_classes, model_args["n_features"])),
        }
    # Assign these initialization parameters and retain their names.
    self.param_list = list(init_params)
    for key, val in init_params.items():
        setattr(self.model, key, val)
    # Also initialize the "classes_" slot with unique predictable labels.
    # FIXME: this assumes target values are integers in range(n_classes).
    setattr(self.model, "classes_", np.arange(n_classes))

def set_init_params(self, model_args: Dict[str, Any]):
    """Initialize the model's trainable parameters."""
    init_params = {
        "intercept_": np.array([0.0]),
        "coef_": np.array([0.0] * model_args["n_features"]),
    }
    self.param_list = list(init_params)
    for key, val in init_params.items():
        setattr(self.model, key, val)

def add_corrections_to_gradients(
    self,
    corrections: Union[TorchVector, Dict[str, torch.Tensor]],
) -> None:
    """Adds values to attached gradients in the model

    Args:
        corrections: corrections to be added to model's gradients
    """
    iterator = self._get_iterator_model_params(corrections)
    for name, update in iterator:
        param = self.model.get_parameter(name)
        if param.grad is not None:
            param.grad.add_(update.to(param.grad.device))

def apply_updates(
    self,
    updates: Union[TorchVector, Dict[str, torch.Tensor]],
) -> None:
    """Apply incoming updates to the wrapped model's parameters.

    Args:
        updates: model updates to be added to the model.
    """
    iterator = self._get_iterator_model_params(updates)
    with torch.no_grad():
        for name, update in iterator:
            param = self.model.get_parameter(name)
            param.add_(update.to(param.device))

def export(self, filename: str) -> None:
    """Export the wrapped model to a dump file.

    Args:
        filename: path to the file where the model will be saved.

    !!! info "Notes":
        This method is designed to save the model to a local dump
        file for easy re-use by the same user, possibly outside of
        Fed-BioMed. It is not designed to produce trustworthy data
        dumps and is not used to exchange models and their weights
        as part of the federated learning process.

    !!! warning "Warning":
        This method uses `torch.save`, which relies on pickle and
        is therefore hard to trust by third-party loading methods.
    """
    torch.save(self.model, filename)

def flatten(self) -> List[float]:
    """Gets weights as flatten vector

    Returns:
        to_list: Convert np.ndarray to a list if it is True.
    """

    params: List[float] = torch.nn.utils.parameters_to_vector(
        self.model.parameters()
    ).tolist()

    return params

def get_gradients(
    self,
    as_vector: bool = False,
) -> Union[Dict[str, torch.Tensor], TorchVector]:
    """Return the gradients attached to the model, opt. as a declearn TorchVector.

    Args:
        as_vector: Whether to wrap returned gradients into a declearn Vector.

    Returns:
        Gradients, as a dictionary mapping parameters' names to their gradient's
            torch tensor, or as a declearn TorchVector wrapping such a dict.
    """
    gradients = {
        name: param.grad.detach().clone()
        for name, param in self.model.named_parameters()
        if (param.requires_grad and param.grad is not None)
    }
    if len(gradients) < len(list(self.model.named_parameters())):
        # FIXME: this will be triggered when having some frozen weights even if training was properly conducted
        logger.warning(
            "Warning: can not retrieve all gradients from the model. Are you sure you have "
            "trained the model beforehand?"
        )
    if as_vector:
        return TorchVector(gradients)
    return gradients

def get_weights(
    self,
    as_vector: bool = False,
    only_trainable: bool = False,
) -> Union[Dict[str, torch.Tensor], TorchVector]:
    """Return the model's parameters, optionally as a declearn TorchVector.

    Args:
        only_trainable (bool, optional): whether to gather weights only on trainable layers (ie
            non-frozen layers) or all layers (trainable and frozen). Defaults to False, (trainable and
            frozen ones)
        as_vector: Whether to wrap returned weights into a declearn Vector.

    Returns:
        Model weights, as a dictionary mapping parameters' names to their
            torch tensor, or as a declearn TorchVector wrapping such a dict.
    """
    parameters = {
        name: param.detach().clone()
        for name, param in self.model.named_parameters()
        if param.requires_grad or not only_trainable
    }
    if as_vector:
        return TorchVector(parameters)
    return parameters

def init_training(self) -> None:
    """Initializes and sets attributes before the training.

    Initializes `init_params` as a copy of the initial parameters of the model
    """
    # initial aggregated model parameters
    self.init_params = {
        key: param.data.detach().clone()
        for key, param in self.model.named_parameters()
    }
    self.model.train()  # pytorch switch for training
    self.model.zero_grad()

def predict(
    self,
    inputs: torch.Tensor,
) -> np.ndarray:
    """Computes prediction given input data.

    Args:
        inputs: input data

    Returns:
        Model predictions returned as a numpy array
    """
    self.model.eval()  # pytorch switch for model inference-mode
    with torch.no_grad():
        pred = self.model(inputs)
    return pred.cpu().numpy()

def send_to_device(
    self,
    device: torch.device,
) -> None:
    """Sends model to device

    Args:
        device: device set for using GPU or CPU.
    """
    self.model.to(device)

def set_weights(
    self,
    weights: Union[Dict[str, torch.Tensor], TorchVector],
) -> None:
    """Sets model weights.

    Args:
        weights: Model weights, as a dict mapping parameters' names to their
            torch tensor, or as a declearn TorchVector wrapping such a dict.
    """
    state_dict = dict(self._get_iterator_model_params(weights))
    incompatible = self.model.load_state_dict(state_dict, strict=False)
    if incompatible.missing_keys:
        logger.warning(
            "'TorchModel.set_weights' received inputs that did not cover all"
            "model parameters; missing weights: %s",
            incompatible.missing_keys
        )
    if incompatible.unexpected_keys:
        logger.warning(
            "'TorchModel.set_weights' received inputs with unexpected names: %s",
            incompatible.unexpected_keys
        )

def train(
    self,
    inputs: torch.Tensor,
    targets: torch.Tensor,
    **kwargs,
) -> None:
    # TODO: should we pass loss function here? and do the backward prop?
    if not self.init_params:
        raise FedbiomedModelError(
            f"{ErrorNumbers.FB622.value}. Training has not been initialized, please initialize it beforehand"
        )

def unflatten(
        self,
        weights_vector: List[float]
) -> Dict[str, torch.Tensor]:
    """Unflatten vectorized model weights using [`vector_to_parameters`][torch.nn.utils.vector_to_parameters]

    This method does not manipulate current model weights modify model parameters.

    Args:
        weights_vector: Vectorized model weights to convert dict

    Returns:
        Model dictionary
    """

    super().unflatten(weights_vector)

    # Copy model to make sure global model parameters won't be overwritten
    model = copy.deepcopy(self.model)
    vector = torch.as_tensor(weights_vector).type(torch.DoubleTensor)

    # Following operation updates model parameters of copied model object
    try:
        torch.nn.utils.vector_to_parameters(vector, model.parameters())
    except TypeError as e:
        FedbiomedModelError(
            f"{ErrorNumbers.FB622.value} Can not unflatten model parameters. {e}"
        )

    return TorchModel(model).get_weights()