"""Functional implementation of FedAdam 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_adam(  # noqa: PLR0913, PLR0917
    server_state: ServerState,
    fedavg_result: dict[tuple[str, ModelStateNames], np.ndarray],
    beta_1: float = 0.9,
    beta_2: float = 0.999,
    eta: float = 1e-1,
    tau: float = 1e-9,
) -> None:
    """Perform federated averaging with Adam optimizer.

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


    Parameters
    ----------
    server_state : ServerState
        The state of the server, including model parameters and momentum vectors.
    fedavg_result : dict[tuple[str, ModelStateNames], np.ndarray]
        The result of the federated averaging step, containing the updated model
        parameters.
    beta_1 : float, optional
        The exponential decay rate for the first moment estimate, by default 0.9
    beta_2 : float, optional
        The exponential decay rate for the second moment estimate, by default 0.999
    eta : float, optional
        The learning rate, by default 1e-1
    tau : float, optional
        The smoothing term, by default 1e-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"
    # Get the server second momentum from the server state
    second_momentum_vector = server_state.second_momentum_vector
    assert (
        second_momentum_vector is not None
    ), "Second 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_second_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

        # Compute first momentum of layer i
        momentum_vector[i] = (
            beta_1 * momentum_vector[i] + (1 - beta_1) * layer_pseudo_gradient
        )
        # Compute second momentum of layer i
        second_momentum_vector[i] = beta_2 * second_momentum_vector[i] + (
            1 - beta_2
        ) * np.multiply(
            layer_pseudo_gradient,
            layer_pseudo_gradient,
        )
        # Compute the new weights of layer i
        layer_fedadam_result = x + eta * np.divide(
            momentum_vector[i] * (1 / (1 - beta_1**server_state.current_round)),
            (
                np.sqrt(
                    second_momentum_vector[i]
                    * (1 / (1 - beta_2**server_state.current_round)),
                )
                + tau
            ),
        )
        # Assign new values to the parameters variable
        parameters[i] = layer_fedadam_result

        # Metrics collection
        layerwise_l2_norms_pseudo_gradient.append(l2_norm([layer_pseudo_gradient]))
        layerwise_l2_norms_momentum_vector.append(
            l2_norm([momentum_vector[i]]),
        )
        layerwise_l2_norms_second_momentum_vector.append(
            l2_norm([second_momentum_vector[i]]),
        )
        layerwise_l2_norms_fedavg_result.append(
            l2_norm([x - layer_pseudo_gradient]),
        )
        layerwise_l2_norms_model.append(l2_norm([layer_fedadam_result]))
