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

from concurrent.futures import ThreadPoolExecutor
from logging import WARNING

import numpy as np
from flwr.common import log as flwr_log

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_nesterov_momentum(  # noqa: PLR0913, PLR0915, PLR0914
    server_state: ServerState,
    fedavg_result: dict[tuple[str, ModelStateNames], np.ndarray],
    aggregation_num_workers_process_ndarrays: int | None = None,
    server_learning_rate: float = 1.0,
    *,
    server_momentum: float = 0.9,
    ndarrays_parallelization: bool = True,
) -> None:
    """Perform federated averaging with Nesterov momentum for specified layers.

    This function applies a standard federated averaging update to most layers
    and a Nesterov momentum update to layers specified as EXP_AVG parameters.
    It computes and stores norms of the pseudo-gradient, momentum vector,
    update result, and final model state for each layer.

    Parameters
    ----------
    server_state : ServerState
        The current state of the server, including model parameters and layer structure.
    fedavg_result : dict[tuple[str, ModelStateNames], np.ndarray]
        The aggregated results from clients, where keys are tuples of layer names and
        model state names.
    aggregation_num_workers_process_ndarrays : int | None, optional
        The number of worker threads to use for processing NDArrays. If None, no
        parallelization is applied.
    server_learning_rate : float, optional
        The learning rate for the server during the aggregation. Defaults to 1.0.
    server_momentum : float, optional
        The momentum factor for the Nesterov update. Defaults to 0.9.
    ndarrays_parallelization : bool, optional
        Whether to use parallel processing for NDArrays. Defaults to True.

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

    # --- Thread Pool Executor ---
    ndarrays_agg_executor: ThreadPoolExecutor | None = None
    if ndarrays_parallelization:
        ndarrays_agg_executor = ThreadPoolExecutor(
            max_workers=aggregation_num_workers_process_ndarrays,
        )

    # --- Index Maps ---
    # Global map for all model states
    global_key_to_index = {
        (key, typ): i for i, (key, typ) in enumerate(server_state.layer_names_and_types)
    }
    # Corrected map for momentum vector, which only contains parameters
    momentum_key_to_index = {
        (key, typ): i
        for i, (key, typ) in enumerate(
            (
                k_t
                for k_t in server_state.layer_names_and_types
                if k_t[1] == ModelStateNames.PARAMETERS
            ),
        )
    }

    # 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] = []

    # --- Update Functions ---
    def update_layer(
        item: tuple[tuple[str, ModelStateNames], np.ndarray],
    ) -> tuple[float, float, float, float]:
        (layer_name, model_state_name), layer_res = item
        i = global_key_to_index[layer_name, model_state_name]
        x = parameters[i]

        layer_pseudo_gradient = x - layer_res
        layer_fedavg_result = x - server_learning_rate * layer_pseudo_gradient
        parameters[i] = layer_fedavg_result

        # For non-momentum layers, the momentum norm is 0
        momentum_norm = 0.0

        return (
            l2_norm([layer_pseudo_gradient]),
            momentum_norm,
            l2_norm([x - layer_pseudo_gradient]),
            l2_norm([layer_fedavg_result]),
        )

    def update_layer_nesterov(
        item: tuple[tuple[str, ModelStateNames], np.ndarray],
    ) -> tuple[float, float, float, float]:
        # --- Setup ---
        (layer_name, model_state_name), layer_res = item
        g_idx = global_key_to_index[layer_name, model_state_name]
        m_idx = momentum_key_to_index[layer_name, model_state_name]

        # Use names similar to PyTorch for clarity (parameter, gradient, buffer)
        param = parameters[g_idx]
        momentum_buffer = momentum_vector[m_idx]

        # --- PyTorch Equivalent Logic (with nesterov=True, dampening=0) ---

        # 1. Calculate the gradient (in this context, the pseudo-gradient)
        # In PyTorch: grad = grads[i]
        grad = param - layer_res

        # 2. Update the momentum buffer
        # In PyTorch: buf.mul_(momentum).add_(grad)
        # v_new <- mu * v_prev + grad
        new_momentum_buffer = server_momentum * momentum_buffer + grad

        # 3. Calculate the final update term using the new momentum buffer
        # In PyTorch: grad = grad.add(buf, alpha=momentum)
        # g_update <- g_original + mu * v_new
        grad_update = grad + server_momentum * new_momentum_buffer

        # 4. Apply the final update to the parameter
        # In PyTorch: param.add_(grad, alpha=-lr)
        # p_new <- p_old - lr * g_update
        updated_param = param - server_learning_rate * grad_update

        # --- Update Server State ---
        parameters[g_idx] = updated_param
        momentum_vector[m_idx] = new_momentum_buffer

        # --- Return Metrics ---
        return (
            l2_norm([grad]),  # Norm of the original pseudo-gradient
            l2_norm([new_momentum_buffer]),  # Norm of the updated momentum buffer
            l2_norm([param - grad]),  # Norm of original param - original grad
            l2_norm([updated_param]),  # Norm of the final resulting parameter
        )

    # --- Partition and Dispatch Tasks ---
    nesterov_items = [
        ((layer_name, model_state_name), res)
        for (layer_name, model_state_name), res in fedavg_result.items()
        if model_state_name == ModelStateNames.PARAMETERS
    ]
    fedavg_items = [
        ((layer_name, model_state_name), res)
        for (layer_name, model_state_name), res in fedavg_result.items()
        if model_state_name != ModelStateNames.PARAMETERS
    ]

    if ndarrays_agg_executor is not None:
        fedavg_futures = [
            ndarrays_agg_executor.submit(update_layer, item) for item in fedavg_items
        ]
        nesterov_futures = [
            ndarrays_agg_executor.submit(update_layer_nesterov, item)
            for item in nesterov_items
        ]

        for future in fedavg_futures + nesterov_futures:
            norm_pseudo, norm_momentum, norm_fedavg, norm_model = future.result()
            layerwise_l2_norms_pseudo_gradient.append(norm_pseudo)
            layerwise_l2_norms_momentum_vector.append(norm_momentum)
            layerwise_l2_norms_fedavg_result.append(norm_fedavg)
            layerwise_l2_norms_model.append(norm_model)

        ndarrays_agg_executor.shutdown()
    else:
        # Fallback to serial updates
        for item in fedavg_items:
            norm_pseudo, norm_momentum, norm_fedavg, norm_model = update_layer(item)
            layerwise_l2_norms_pseudo_gradient.append(norm_pseudo)
            layerwise_l2_norms_momentum_vector.append(norm_momentum)
            layerwise_l2_norms_fedavg_result.append(norm_fedavg)
            layerwise_l2_norms_model.append(norm_model)

        for nesterov_item in nesterov_items:
            norm_pseudo, norm_momentum, norm_fedavg, norm_model = update_layer_nesterov(
                nesterov_item,
            )
            layerwise_l2_norms_pseudo_gradient.append(norm_pseudo)
            layerwise_l2_norms_momentum_vector.append(norm_momentum)
            layerwise_l2_norms_fedavg_result.append(norm_fedavg)
            layerwise_l2_norms_model.append(norm_model)

    # --- Log Metrics ---
    flwr_log(
        WARNING,
        f"Layerwise L2 norms - Pseudo Gradient: {layerwise_l2_norms_pseudo_gradient}, "
        f"Momentum Vector: {layerwise_l2_norms_momentum_vector}, "
        f"FedAvg Result: {layerwise_l2_norms_fedavg_result}, "
        f"Model: {layerwise_l2_norms_model}",
    )
