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

from concurrent.futures import ThreadPoolExecutor

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_efficient(
    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,
    *,
    ndarrays_parallelization: bool = True,
) -> None:
    """Perform federated averaging with efficient handling of parameters.

    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.
    ndarrays_parallelization : bool, optional
        Whether to use parallel processing for NDArrays. Defaults to True.

    """
    # Get the parameters from the server state
    parameters = server_state.model_states
    # If the executor is not set, initialize it
    # TODO(anon): We are recreating the executor each time whilst it might be better
    # to reuse it across multiple calls to this function.
    ndarrays_agg_executor: ThreadPoolExecutor | None = None
    if ndarrays_parallelization:
        ndarrays_agg_executor = ThreadPoolExecutor(
            max_workers=aggregation_num_workers_process_ndarrays,
        )

    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_fedavg_result: list[float] = []
    layerwise_l2_norms_model: list[float] = []

    def update_layer(
        item: tuple[tuple[str, ModelStateNames], np.ndarray],
    ) -> tuple[float, float, float]:
        # item is a tuple: ((layer_name, model_state_name), aggregated_value)
        (layer_name, model_state_name), layer_res = item
        i = keys_to_index[layer_name, model_state_name]
        # Layer i pseudo-gradient
        x = parameters[i]
        # Compute pseudo-gradient and new values
        layer_pseudo_gradient = x - layer_res
        layer_fedavg_result = x - server_learning_rate * layer_pseudo_gradient
        # Update parameter tensor
        parameters[i] = layer_fedavg_result
        # Return computed metrics
        return (
            l2_norm([layer_pseudo_gradient]),
            l2_norm([x - layer_pseudo_gradient]),
            l2_norm([layer_fedavg_result]),
        )

    if ndarrays_agg_executor is not None:
        # Schedule each layer update concurrently
        futures = [
            ndarrays_agg_executor.submit(update_layer, item)
            for item in fedavg_result.items()
        ]
        for future in futures:
            norm_pseudo, norm_fedavg, norm_model = future.result()
            layerwise_l2_norms_pseudo_gradient.append(norm_pseudo)
            layerwise_l2_norms_fedavg_result.append(norm_fedavg)
            layerwise_l2_norms_model.append(norm_model)
    else:
        # Fallback to serial updates
        for item in fedavg_result.items():
            norm_pseudo, norm_fedavg, norm_model = update_layer(item)
            layerwise_l2_norms_pseudo_gradient.append(norm_pseudo)
            layerwise_l2_norms_fedavg_result.append(norm_fedavg)
            layerwise_l2_norms_model.append(norm_model)
