"""Functional implementation of FedMom with efficient parameter handling."""

import numpy as np

from mt_dao.masks_utils import ModelStateNames
from mt_dao.server.server_state import ServerState
from mt_dao.utils import l2_norm


def federated_averaging_with_momentum(
    server_state: ServerState,
    fedavg_result: dict[tuple[str, ModelStateNames], np.ndarray],
    server_learning_rate: float = 1.0,
    server_momentum: float = 0.9,
) -> None:
    """Perform federated averaging with momentum.

    This function applies the FedMom algorithm to the given federated
    averaging result, using the server's momentum vector to update the
    model parameters.

    Parameters
    ----------
    server_state : ServerState
        The current state of the server, including model parameters and
        momentum vector.
    fedavg_result : dict[tuple[str, ModelStateNames], np.ndarray]
        The result of the federated averaging step.
    server_learning_rate : float, optional
        The learning rate for the server, by default 1.0
    server_momentum : float, optional
        The momentum for the server, by default 0.9

    """
    # Get the parameters from the server state
    parameters = server_state.model_states
    # Get the server first momentum from the server state
    momentum_vector = server_state.momentum_vector
    assert (
        momentum_vector is not None
    ), "Momentum vector must be initialized in the server state"

    keys_to_index = {
        (key, typ): i for i, (key, typ) in enumerate(server_state.layer_names_and_types)
    }

    # Initialize the metrics
    layerwise_l2_norms_pseudo_gradient: list[float] = []
    layerwise_l2_norms_momentum_vector: list[float] = []
    layerwise_l2_norms_fedavg_result: list[float] = []
    layerwise_l2_norms_model: list[float] = []
    # Loop over layer, apply the server optimizer and compute metrics
    for (layer_name, model_state_name), layer_res in fedavg_result.items():
        i = keys_to_index[layer_name, model_state_name]
        x = parameters[i]
        # Layer i pseudo-gradient
        layer_pseudo_gradient = x - layer_res

        # Using FedMom implementation
        # Layer i new momentum vector
        layer_new_momentum_vector = x - server_learning_rate * layer_pseudo_gradient
        # Layer i new values
        layer_fedavgm_result = (
            1 + server_momentum
        ) * layer_new_momentum_vector - server_momentum * momentum_vector[i]
        # Assign new values to the parameters variable
        parameters[i] = layer_fedavgm_result
        # Assign new values to the momentum vector
        momentum_vector[i] = layer_new_momentum_vector
        # Metrics collection
        layerwise_l2_norms_pseudo_gradient.append(l2_norm([layer_pseudo_gradient]))
        layerwise_l2_norms_momentum_vector.append(
            l2_norm([layer_new_momentum_vector]),
        )
        layerwise_l2_norms_fedavg_result.append(l2_norm([x]))
        layerwise_l2_norms_model.append(l2_norm([layer_fedavgm_result]))
